Source-linked AI summary
MIST: Mutual Information Estimation Via Supervised Training
German Gritsai, Megan Richards, Maxime Méloux, Kyunghyun Cho, Maxime Peyrard
TL;DR
Estimating mutual information from samples is difficult in challenging settings that prior evaluations have rarely tested, especially in combination. MIST learns MI directly from synthetic data, substantially outperforming existing methods across the studied conditions; its quantile-based approach also supports uncertainty estimates.
Problem
Method
Results
MIST substantially outperforms existing methods across the studied conditions and requires approximately half as many samples on average as the next-best estimator for reliable estimates.
Source: §1 · p. 2Takeaways & Limitations
MI’s invariance to invertible transformations allows normalizing flows to map meta-datasets to target modalities without altering ground-truth MI.
Source: §6 · p. 15MIST lacks theoretical consistency guarantees; consistency has only been verified empirically, and it is not self-consistent on data-processing and additivity tests.
Source: §5.1 · p. 13
Abstract
from arXiv · showhide
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.
1 Introduction
MI estimation from finite samples remains challenging in settings that are common in practice but underexplored in prior evaluations. The paper introduces MIST, a supervised estimator trained on synthetic distributions, and reports gains across a broad range of conditions.
1 Introduction
- Because the joint and marginal distributions are rarely known, MI is commonly inferred from finite samples using density or density-ratio estimators.MI quantifies nonlinear statistical dependence and is applied across data science tasks.Source: §1 · pp. 1–2
- Prior evaluations rarely test challenging conditions individually or in combination, often using fewer than five distributions and at least 10,000 samples.The challenging settings include high dimensions, limited samples, complex distributions, and high MI.Source: §1 · p. 2
- MIST learns to predict MI directly from samples, training a neural network end-to-end on synthetic joint distributions with known ground-truth MI.The estimator handles arbitrary sample sizes and dimensions and includes built-in uncertainty estimates.Source: §1 · p. 2
- MIST substantially outperforms existing methods across low-sample, higher-dimensional settings and 23 distribution types, especially in low-sample regimes.It requires approximately half as many samples on average as the next-best estimator for reliable estimates.Source: §1 · p. 2
2 Related Work
Related MI estimators are grouped into density and density-ratio approaches, while a broader line of work learns statistical inference procedures from datasets.
2 Related Work
- Density estimators first approximate density functions, while density-ratio estimators estimate the relationship between the joint and product-of-marginal densities.These are the two main families used to group existing MI estimation methods.Source: §2 · p. 3
- KSG extends k-nearest-neighbor entropy estimation to MI and serves as a primary benchmark, while density-estimation methods also use generative models and variational bounds.Other density approaches include diffusion-based MINDE, copula-based VCE, MI-preserving flows, and compression followed by KSG.Source: §2 · p. 3
- Density-ratio methods use variational bounds or discriminative tasks; their limitations include high variance, high sample requirements, and poor estimates at high MI.InfoNet instead emulates a trained MINE critic, whereas MIST learns to predict MI directly from samples.Source: §2 · pp. 3–4
- Related statistical-inference work uses neural processes, transformers, and Deep Sets to predict latent quantities from datasets while enforcing permutation invariance.The broader direction also includes prior-fitted networks, simulation-based inference, amortized Bayesian inference, and learned frequentist estimators.Source: §2 · p. 4
3 Learning Framework
Section 3 frames MI estimation as direct supervised prediction from finite paired samples rather than estimating densities or their ratios. It trains MIST on synthetic distributions with known MI and adds a quantile-regression variant to estimate sampling uncertainty.
3 Learning Framework
- MIST replaces density or density-ratio intermediates with an end-to-end neural predictor that maps paired samples to MI.The estimator is trained on a meta-dataset of synthetic joint distributions with known ground-truth MI.Source: §3 · p. 4
- Training minimizes mean squared error over datasets sampled from distributions in a meta-distribution, whose labels are their known MI values.With sufficient model capacity and perfect optimization, the Bayes-optimal predictor is the conditional expectation of MI given the observed dataset.Source: §3 · p. 5
- Finite datasets can leave the generating distribution unidentified, creating irreducible epistemic uncertainty and MSE.VCE also fails on some low-sample, high-dimensional points when its estimated covariance matrix is not positive definite or invertible.Source: §3 · pp. 5, 7
- MISTQR predicts the τ-quantile of MI's sampling distribution, and querying multiple τ values approximates that distribution with a forward pass.The paper evaluates the calibration and reliability of these intervals against bootstrap alternatives.Source: §3 · pp. 5–6
- The training meta-dataset contains 625,000 datapoints from 16 base distribution families; BMI transformations produce variants across dimensions 2–32, with test families held out from training.The extended test set contains 806,000 datapoints from 13 families, six shared with training.Source: §3 · p. 6
- SetTransformer++ handles variable-size samples, while dimension-axis attention with learned inducing-point pooling accommodates varying dimensions; direct MI prediction works best when labels span a sufficient range.The dimension-axis attention is not required to be permutation-invariant.Source: §3 · p. 6
- MIST estimators substantially outperform existing methods across sample sizes and distributions, including BMI-generated distributions unseen during training.Table 1 reports averaged MSE grouped by sample count.Source: Table 1 · p. 7
4 Experiments
Across the experiments, MIST performs strongly against existing estimators across dimensions and sample sizes, including unseen distributions, and the study evaluates its uncertainty, generalization, efficiency, and feature-selection use.
4.1 MIST Outperforms Existing Estimators, Even on Unseen BMI Distributions
- MIST achieves approximately 10× lower error on training-seen distributions (IMD) and ∼5× lower error on unseen distributions (OoMD) than KSG.KSG is identified as the next strongest method in this benchmark.Source: §4.1 · p. 7
4.2 MIST Exhibits Lower Bias and Variance Across Settings
- Up to 100× lower loss occurs in high-dimensional, low-sample settings; MIST is nearly unbiased across settings and its estimates track high MI more closely than KSG’s.MIST’s confidence intervals achieve approximately 2× better coverage than KSG across settings.Source: §4.2 · p. 8, Figure 3 · p. 8
4.3 MIST Requires Fewer Samples to Scale to Higher Dimensions
- MIST and MISTQR require roughly half as many samples as KSG to meet selected MSE thresholds as dimensionality rises.For some thresholds, KSG requires more than 500 samples and still does not achieve reliable performance.Source: Figure 4 · p. 9, §4.3 · p. 9
4.4 MIST Quantile-Based Intervals are More Reliable Than Bootstrapping
- MISTQR is nearly perfectly calibrated on training-seen distributions and remains well-calibrated on unseen families; both MIST variants outperform KSG in calibration.MISTQR obtains uncertainty estimates by querying the model at different quantile values; confidence intervals require only additional forward passes.Source: Figure 5 · p. 9, §4.4 · p. 10
4.5 MIST Generalizes to Unseen BMI Distributions, Dimensions, and Sample Sizes
- Models trained with dimensions below 16 and sample sizes below 300 generalize to larger dimensions and sample sizes on in-distribution data, but struggle on unseen families at higher dimensions.They also generalize to unseen families when dimensionality is similar to training.Source: §4.5 · p. 10, Figure 6 · p. 10
4.6 MIST Achieves Large Gains in Inference Efficiency
- MIST inference is 1.7× faster than the most efficient baseline and 4–80× faster than KSG.Training MIST and MISTQR each took 3 hours and 45 minutes, with that one-time cost amortized across later inference tasks.Source: §4.6 · p. 11, Figure 7 · p. 11
4.7 Training on Multiple Dimensions Benefits High-Dimensional Estimation
- Scaling both dimensionality and dataset size yields 2–3× lower MSE for dimensions D ≥16, but does not improve performance for D ≤8.At matched training-set size, mixed-dimension training performs worse than a specialized model; the authors hypothesize it requires more samples and longer training.Source: §4.7 · pp. 11–12
4.8 Feature Selection
- In synthetic feature selection, MIST outperforms KSG and MINE for direct selection, and MIST-guided learned selection is more accurate, lower-variance, and more stable across sample sizes and noise levels.The study compares individual feature scoring with gradient ascent on collective mutual information.Source: §4.8 · p. 12, Figure 9 · p. 13
5 Limitations
MIST trades classical theoretical guarantees for empirically learned behavior, with generalization and calibration dependent on the training distribution. The section describes boundaries on out-of-distribution use and practical mitigation strategies.
5.1 Theoretical Guarantees
- MIST lacks classical consistency guarantees; data-processing and additivity tests reveal inconsistency and high-variance prediction ratios, so general-purpose use as a learning objective may require further work.Consistency is verified empirically rather than guaranteed theoretically, and reliability during optimization remains an open challenge.Source: §5.1 · p. 13
5.2 Generalization Out of Meta-Distribution
- Performance depends on meta-training coverage, and generalization to distant out-of-distribution families cannot be guaranteed; degradation may affect predictions and confidence intervals.Unseen distributions in the experiments remain BMI-generated and do not represent arbitrary or adversarial shifts such as discrete or highly structured data.Source: §5.2 · pp. 13–14
- Unseen dimensionalities, especially combined with distribution shift, cause larger degradation and can reduce confidence-interval coverage, although MIST may still outperform baselines at low sample sizes.MIST generalizes smoothly to larger unseen sample sizes, but classical baselines can partially catch up while MIST remains orders of magnitude faster.Source: §5.2 · p. 14
- MIST’s predicted MI range is limited to values seen during training, so applications expecting larger MI values require a mitigation strategy.Normalized-quantity variants may suit larger MI values, though experiments found slightly worse performance than standard MIST.Source: §5.2 · p. 14, §5.3 · p. 14
- MISTQR calibration depends on the training distribution; out-of-distribution bound failures increase, with upper-bound failures rising from 4.3% to 27.5%.Lower-bound failures rise from 9.6% in-distribution to 15.8% out-of-distribution.Source: §5.2 · p. 14
5.3 Mitigation Strategies
- For high-dimensional inputs, the authors recommend projecting below 32 dimensions; embedding-space distance from training datasets can flag prediction and calibration degradation.The distance correlates with MSE degradation (Spearman 0.69) and confidence-interval coverage (−0.33).Source: §5.3 · p. 14
6 Discussion
The discussion presents MIST as a flexible, data-driven alternative whose behavior can be shaped through training data and objectives, while emphasizing speed and future design directions. The appendix sections review equivalent MI definitions and categorize existing estimators by their estimation task.
6 Discussion
- MIST’s behavior is learned end-to-end from a synthetic meta-dataset with known MI, replacing analytic density-based or variational derivations with a flexible empirical design.This is a deliberate trade of theoretical motivations for data-driven flexibility.Source: §6 · p. 15
- Because MI is invariant to invertible transformations, normalizing flows can adapt meta-datasets to target modalities while preserving ground-truth MI, supporting specialized estimators.The authors describe potential targets including images, text, tabular data, and multimodal signals.Source: §6 · p. 15
- After training, MIST inference takes one forward pass, or two for a quantile-model interval, and is three to four orders of magnitude faster than neural MI estimators and one order faster than classical KSG implementations.The authors attribute the speedups to amortizing inference cost during meta-training.Source: §6 · p. 15
- Future designs could incorporate theoretical results as inductive biases, such as penalizing data-processing inequality violations, or use ranking losses to prioritize relative ordering.The authors frame trainable, meta-learned estimators as flexible, efficient, customizable alternatives to classical estimators.Source: §6 · p. 15
A.1 Definitions of Mutual Information
- MI can be expressed as the KL divergence between the joint distribution and the product of marginals, or as entropy reduction; independence makes the KL-based MI zero.The conditional-entropy form measures uncertainty reduction about one variable given the other.Source: §A.1 · p. 21
A.2 MI Estimators
- The appendix groups estimators by task: density estimators approximate densities, while density-ratio estimators model the relation between joint and marginal densities.This categorization is intended to distinguish estimation tasks from modeling choices, which can differ for a method such as FMMI.Source: §A.2 · p. 21
A.3 Empirical Evaluation Scope of MI Estimators
The appendix compares the empirical evaluation scope of prominent MI estimators, including benchmark settings and reported dimensions, sample sizes, and MI ranges. It excludes downstream evaluations when ground-truth MI is unavailable.
A.3 Empirical Evaluation Scope of MI Estimators
- Table 3 summarizes direct MI-estimation evaluations by task scope; it excludes downstream settings without ground-truth MI values and marks unavailable or unspecified entries with dashes.The summary includes evaluations from the papers introducing each estimator.Source: Table 3 · p. 23, §A.3 · p. 23
- Early histogram and KDE studies used simulated time series without ground-truth MI, while KSG evaluations included synthetic distributions up to dimension 8 and MI values up to 1.2.KSG’s minimum sample size was 125 for 1D Gaussian experiments and 10K for other distributions, with 50K used for larger dimensions.Source: §A.3 · p. 23
- Latent-MI evaluations ranged from 0–2 bits on multivariate Gaussians, included sample-requirement tests through dimension 50, and covered MNIST and protein embeddings.The embedding evaluations used dimensions 784 and 1024, respectively.Source: §A.3 · p. 24
- VCE evaluations covered synthetic distributions at dimensions 2–64 and MI up to 16, plus image and language embeddings reduced to 16 dimensions.The reported embedding settings included image MI up to 7 and language-embedding MI up to 2.1.Source: §A.3 · p. 24
- MINE was evaluated on 2D and 20D Gaussians with MI ranges up to 2 and 40, whereas FMMI covered four distribution types with MI up to 80 nats.The FMMI evaluation used 100K training samples and 1K test samples.Source: §A.3 · p. 25
- Other neural evaluations included InfoNet on 1D Gaussians and GMMs, BMI-based MINDE tests, and GM/SMILE experiments that increased MI up to 10.The MIST evaluation used 23 BMI-generated distribution types, maximum dimension 32, MI up to 40, and sample sizes starting at 10.Source: §A.3 · pp. 24–25
B Details of Meta-Dataset Creation
The meta-dataset pairs synthetic joint samples with known MI values, spanning varied distribution families, dimensions, and MI-preserving transformations. Disjoint training and test families support evaluation both within and outside the training distribution.
B Details of Meta-Dataset Creation
- The meta-dataset pairs samples from synthetic joint distributions with ground-truth MI, and partitions distribution families into disjoint training and test groups across dimensions 2–32.It includes smaller Mtest and larger Mtest-extended evaluation corpora for comparison with existing methods and assessment on novel, higher-dimensional distributions.Source: B Details of Meta-Dataset Creation · p. 25
Distribution families
- The distribution families include Gaussian and Student’s t variants with dense or sparse structure, plus bounded-support additive-noise distributions whose MI sums per-dimension contributions.Gaussian and Student’s t variants include two-pair cases; Student’s t variants use heavier tails, while additive-noise MI equals the sum of per-dimension contributions.Source: Distribution families · p. 26
MI-Preserving Transformations
- MI-preserving transformations include identity, asinh, signed-power halfcube, and wigglify maps; wigglify adds high-frequency nonlinear structure while preserving value ordering.The transformations are applied independently to the X and Y marginals.Source: MI-Preserving Transformations · pp. 26–27
Meta-Dataset Composition
- Table 4 labels each family by base distribution, structure, and transformation; reduced 1k-sample IMD and OoMD test subsets preserve their full counterparts’ family composition.The true MI values are computed analytically or numerically depending on tractability.Source: Distribution families · p. 26, Table 4 · p. 27
- Figure 10 shows the distribution of true MI values within Mtrain.Source: Figure 10 · p. 28
C.1 Experiments for Normalized and Multi-Task Predictions
The experiments compare normalized and multi-task prediction strategies for MI estimation. Direct MI prediction and a four-coefficient model with physics-informed loss performed best, and the paper proceeds with direct prediction.
C.1 Experiments for Normalized and Multi-Task Predictions
- Target normalization constrains the MI prediction range to mitigate outlier effects and improve numerical stability across wide MI values.The experiments also add normalized predictions for H(X), H(Y), and H(X,Y), either reconstructing MI from the three entropy coefficients or jointly predicting all four coefficients.Source: §C.1 · p. 28
- For the four-coefficient configuration, the physics-informed term aligns two independent MI estimates; λ_phys = 0.3 balances regularization strength in training on dim = 9 samples.The coefficients include normalized predictions for MI and the three entropies, and the training objective takes an expectation over a batch.Source: §C.1 · p. 28
- Direct MI prediction and the four-coefficient formulation with physics-informed loss obtained the best results on Mtest at dim = 9; subsequent experiments use direct MI prediction.The three-coefficient model performed comparatively worse, which the authors attribute to accumulated errors from independently predicted entropy terms.Source: §C.1 · p. 29
C.2 Median results
This section reports median MSE results alongside the main results’ mean MSE.
C.2 Median results
- Median MSE results are grouped by estimator, sample count, and distribution, with the best and second-best rows marked in each column.Source: §C.2 · p. 29, Table 6 · p. 29
C.3 InfoNet results
The appendix reports additional estimator comparisons, dimensionality analyses, inference-time measurements, and model-size results. The supplied passages describe these evaluations and their visuals, but do not provide enough labeled result values to establish the outcomes of those comparisons.
C.3 InfoNet results
- InfoNet comparisons report average and median MSE on base multivariate Gaussian distributions and average MSE across all multivariate Gaussian distributions in Mtest.The base-Gaussian comparison addresses InfoNet’s training on Gaussian mixtures, which comprise less than half of Mtest.Source: §C.3 · p. 29, Table 7 · p. 29, Table 8 · p. 30
C.4 Dimensionality Ablation
- The dimensionality ablation compares variable- and fixed-dimensionality MIST and MISTQR using average MSE, bias magnitude, and variance.Figure 12 separates MISTQR results for IMD and OoMD subsets; Figures 13–15 summarize proposed models on a test set similar to training.Source: §C.4 · p. 30, Figure 11 · p. 30, Figure 12 · p. 30
C.5 Comparing With Existing Estimators
- Table 9 reports mean MI estimates across distribution families, distinguishing IMD from OoMD and marking best values and those within one order of magnitude of the best.The appendix also reports MSE averaged separately over each distribution family.Source: §C.5 · p. 30, Table 9 · p. 32
C.6 Inference Time
- Table 10 reports average inference time per sample on Mtest-extended.Source: §C.6 · p. 33, Table 10 · p. 33
C.7 Scaling Behavior
- The scaling analysis compares the main 5.5M model with smaller models and presents MIST and MISTQR MSE by sample-size group and dimension.Table 11 summarizes average performance metrics by model size; Figure 16 provides the fine-grained MSE heatmap.Source: §C.7 · p. 33, Table 11 · p. 33
D Estimating OOMD Performance Degradation via MIST Representations
Average distance to nearby training embeddings provides a fast proxy for MIST performance degradation on out-of-distribution datasets. Greater distances are associated with higher MSE and lower confidence-interval coverage.
D Estimating OOMD Performance Degradation via MIST Representations
- The proxy uses dataset-level embeddings from MIST’s last ISAB layer and average cosine distance to the m nearest training embeddings.The reference set contains K training datasets; at inference, the input dataset is embedded and compared with that set.Source: D Estimating OOMD Performance… · pp. 33, 35
- Larger sample sizes improve average MSE across evaluated dimensionalities, while small-sample regimes make MI estimation more difficult.The comparison covers two proposed models on Mtest-extended, IMD.Source: Figure 16 · p. 34
- Greater average m-NN distance predicts increased MSE and decreased CI coverage, providing a test-time signal for MIST prediction reliability.MISTQR remains at higher performance than MIST for larger deviations from the training set.Source: D Estimating OOMD Performance… · p. 35, Figure 17 · p. 35
- Increasing the reference-set size strengthens the correlation, while varying m from 1 to 10 has little effect; average m-NN distance outperformed the tested alternatives.The alternatives were distance to the embedding centroid and Mahalanobis distance.Source: D Estimating OOMD Performance… · p. 35
- 0.68 and −0.32 are the Spearman correlations between average 10-NN distance and MSE, and CI coverage, respectively.The reported correlations use rank statistics because MSE values are skewed.Source: Figure 17 · p. 35, D Estimating OOMD Performance… · p. 36
- About 10% more inference time is needed to compute 10-NN distances over K = 50,000 embeddings, while embedding recording is negligible.The authors report MIST remains significantly faster than competing baselines and orders of magnitude faster than neural estimators.Source: D Estimating OOMD Performance… · p. 36
E.1 Self-Consistency of Quantile Regression
The section tests MISTQR quantile consistency and examines its limits under distribution shift. It also describes three consistency checks for the learned MI quantity.
E.1 Self-Consistency of Quantile Regression
- 1.1% of IMD and 1.4% of OoMD points violate quantile monotonicity; bound failures rise from 9.6% lower and 4.3% upper on IMD to 15.8% lower and 27.5% upper on OoMD.Monotonicity failures occur only when true MI is below 0.05.Source: §E.1 · p. 36, Table 12 · p. 37
- MISTQR quantiles are conditioned on training target-MI values, and larger train–test location shifts correspond to more bound failures.Downward IMD shifts coincide with more lower-bound failures, while upward OoMD shifts coincide with more upper-bound failures.Source: §E.1 · p. 36, Table 13 · p. 37
- Quantile regression captures sampling uncertainty but not all variation in MIST predictions, including model error, and may under-cover under substantial distribution shift.The authors state that careful validation is required.Source: §E.1 · p. 37
E.2 Self-Consistency of the Learned Quantity
- The learned quantity is checked for independence, data processing, and additivity over independent samples.The tests respectively check zero MI for independent variables, non-increasing MI after further processing, and doubled MI for a second independent sample.Source: §E.2 · p. 37
- The MNIST figure captions define ratios for these checks: independence expects a non-decreasing ratio approaching 1, data processing a ratio near 1, and additivity a ratio near 2.The data-processing comparison adds a further masked image version; the additivity comparison uses two independent samples.Source: Figure 18 · p. 38, Figure 19 · p. 38, Figure 20 · p. 39
F Feature Selection
The feature-selection experiments compare MIST-guided optimization with MINE and KSG in synthetic settings. MIST-guided selection is reported as more accurate and lower variance than MINE.
F Feature Selection
- The synthetic task uses 32 normally distributed features, with two randomly selected informative features and noise variances of 0.3, 0.5, or 1.0.The goal is to recover weights assigning 1 to the informative features and 0 to the others.Source: F Feature Selection · p. 39
- Direct selection scores each feature by its estimated MI with the outcome; ROC curves are shown across noise levels and feature estimates for a random seed.The feature-estimate example uses 50 samples and noise variance 1.0.Source: F Feature Selection · p. 39, Figure 22 · p. 41, Figure 24 · p. 42
- In learned selection, gradient ascent optimizes continuous feature weights using MI between the true label and the label predicted from candidate weights.MIST, KSG, and MINE are evaluated; MINE uses both continual training and training to convergence.Source: F Feature Selection · p. 39
- MIST-guided feature selection is more accurate and lower variance than MINE, achieving perfect accuracy at 100 samples.The comparison includes continually trained and convergence-trained MINE; MIST-guided optimization most reliably distinguishes informative from uninformative features.Source: F Feature Selection · pp. 39–40, Figure 21 · p. 40, Figure 23 · p. 42