MIST: Informazione Mutua Tramite Addestramento Supervisionato

MIST: Mutual Information Via Supervised Training

November 24, 2025

Autori: German Gritsai, Megan Richards, Maxime Méloux, Kyunghyun Cho, Maxime Peyrard

cs.AI

Abstract

Proponiamo un approccio completamente data-driven per la progettazione di stimatori di informazione mutua (MI). Poiché qualsiasi stimatore di MI è una funzione del campione osservato da due variabili casuali, parametrizziamo questa funzione con una rete neurale (MIST) e la addestriamo end-to-end per prevedere i valori di MI. L'addestramento viene eseguito su un ampio meta-dataset di 625.000 distribuzioni congiunte sintetiche con MI noto. Per gestire dimensioni del campione e dimensionalità variabili, impieghiamo uno schema di attenzione bidimensionale che garantisce l'invarianza permutazionale tra i campioni di input. Per quantificare l'incertezza, ottimizziamo una loss di regressione quantilica, permettendo allo stimatore di approssimare la distribuzione campionaria del MI piuttosto che restituire una singola stima puntuale. Questo programma di ricerca si discosta dai lavori precedenti prendendo una rotta completamente empirica, scambiando garanzie teoriche universali per flessibilità ed efficienza. Empiricamente, gli stimatori appresi superano ampiamente i baseline classici su diverse dimensioni del campione e dimensionalità, incluso su distribuzioni congiunte non viste durante l'addestramento. Gli intervalli basati sui quantili risultanti sono ben calibrati e più affidabili degli intervalli di confidenza basati sul bootstrap, mentre l'inferenza è ordini di grandezza più veloce rispetto ai baseline neurali esistenti. Oltre ai vantaggi empirici immediati, questo framework produce stimatori addestrabili e completamente differenziabili che possono essere integrati in pipeline di apprendimento più ampie. Inoltre, sfruttando l'invarianza del MI rispetto alle trasformazioni invertibili, i meta-dataset possono essere adattati a modalità di dati arbitrarie tramite normalizing flows, consentendo un addestramento flessibile per diverse meta-distribuzioni target.

English

We propose a fully data-driven approach to designing mutual information (MI) estimators. Since any MI estimator is a function of the observed sample from two random variables, we parameterize this function with a neural network (MIST) and train it end-to-end to predict MI values. Training is performed on a large meta-dataset of 625,000 synthetic joint distributions with known ground-truth MI. To handle variable sample sizes and dimensions, we employ a two-dimensional attention scheme ensuring permutation invariance across input samples. To quantify uncertainty, we optimize a quantile regression loss, enabling the estimator to approximate the sampling distribution of MI rather than return a single point estimate. This research program departs from prior work by taking a fully empirical route, trading universal theoretical guarantees for flexibility and efficiency. Empirically, the learned estimators largely outperform classical baselines across sample sizes and dimensions, including on joint distributions unseen during training. The resulting quantile-based intervals are well-calibrated and more reliable than bootstrap-based confidence intervals, while inference is orders of magnitude faster than existing neural baselines. Beyond immediate empirical gains, this framework yields trainable, fully differentiable estimators that can be embedded into larger learning pipelines. Moreover, exploiting MI's invariance to invertible transformations, meta-datasets can be adapted to arbitrary data modalities via normalizing flows, enabling flexible training for diverse target meta-distributions.