ALICE: In-context, Zero-shot, Mutual Information Estimation
Giulio Franzese, Simone Rossi, Pietro Michiardi
cs.LG
2026-09-28
One frozen transformer estimates MI for any unseen distribution in-context: 0.092-nat MAE at 1k samples, beating every per-distribution-trained estimator by 2x or more.
Mutual information (MI) quantifies nonlinear statistical dependence between two random variables. InfoNCE in contrastive learning is one of its lower bounds; biology and neuroscience use it to measure how reliably a signaling channel transmits its input. Estimating MI from finite samples is an old and genuinely hard problem: the estimand depends on the full joint density, MI is unbounded and dominated by rare high-density events, and guarantees are fragile in high dimension.
Existing neural estimators share a structural flaw on top of that: each is fit from scratch for every distribution. MINE, InfoNCE, NWJ, and the diffusion-based MINDE, InfoBridge, and FMMI all work this way, with the compute and tuning costs that implies. InfoAtlas amortizes through a hypernetwork and skips per-distribution training, but its MAE sits between 0.25 and 0.28 nats at every budget, a real accuracy penalty. Then there is the small-sample regime: neural estimators need hundreds of thousands of samples and collapse at 1k, while biology and neuroscience datasets often provide a few hundred to a few thousand pairs. ALICE attacks all three constraints at once: no per-distribution training, small-sample viability, and native support for varying dimensionality and data types.
ALICE splits MI estimation into two steps: estimate velocity fields in-context, then apply a fixed identity.
The backbone is rectified flow: interpolate linearly between a clean sample z0 and Gaussian noise ε, and define the velocity field as the conditional mean of the denoising direction. A known identity expresses the KL divergence between two distributions as a time integral of squared velocity differences; the paper derives the MI form: MI equals a weighted integral of the squared difference between the joint velocity field and two block-conditional fields, one block noised, the other held clean as evidence. MI estimation reduces to evaluating three velocity fields.
The design move is representing all three with a single transformer. Input is a context of clean samples from the unseen distribution plus a query carrying position, time t, and a noising mask; the mask decides which coordinates are diffused, three mask patterns produce the three fields, and a fixed Monte Carlo integral returns the estimate at the cost of a few forward passes. Before inference, each marginal is mapped to a standard normal through the empirical CDF, which settles data scaling along the way.
Architecture choices serve any dimension and any context length:
Training data is entirely synthetic: copula mixtures, latent warps, nonparametric regressions, and manifolds, over joint widths 2 to 100 at 70,000 episodes per width. The loss is plain masked velocity regression; no MI labels are needed. Small has 19M parameters, Base 85M.
On the 40-task Beyond Normal benchmark (Czyż et al., 2023; joint widths 2 to 100; closed-form ground truth), every estimator gets the same budget (MAE, nats):
| Method | @1k | @5k | @10k |
| ALICE Base (zero-shot) | 0.092 | 0.063 | 0.060 |
| ALICE Small (zero-shot) | ≈0.10 | ≈0.10 | ≈0.10 |
| InfoAtlas (zero-shot) | 0.25–0.28 | 0.25–0.28 | 0.25–0.28 |
| CCA (classic) | 0.195 | — | — |
| MINDE (per-distribution) | 0.353 | 0.070 | 0.065 |
At the 1k budget ALICE Base beats every competitor by at least 2x. The trained estimators need the jump to 5k/10k to become usable: MINDE crawls from 0.353 to 0.070, InfoNCE from 0.887 to 0.123, and NWJ diverges outright at 1k. On compute, an 8,192-sample budget gets Base to 0.060 nats for 1.6×10¹³ FLOPs per task, while InfoAtlas spends 1.9×10¹³ FLOPs for 0.274.
Three applications on data types absent from the training corpus:
For ML practitioners, MI estimation shifts from training a model per distribution to one frozen model plus a few forward passes, the same move TabPFN made for tabular classification. Anything downstream of MI, from representation analysis to dependence testing, sees a different cost structure.
For AI4Science, the regime where ALICE is strongest is exactly where neural estimators collapse: a few hundred to a few thousand samples, straight to a number. Skipping training also unlocks analyses that need hundreds of estimates, such as sliding windows, per-session splits, and Blahut–Arimoto capacity optimization. At 19M/85M parameters it runs locally on modest hardware.
To be clear about the margins: in the large-data regime the edge narrows (0.060 vs MINDE's 0.065 at 10k). The value concentrates in small samples and the absence of training.
The authors' own list is thin: research code, not thoroughly optimized; model size, training budget, corpus dimensionality, and training context length all have headroom.
Further points from a close read: the training corpus tops out at joint width 100, with higher dimensions untested; contexts beyond the training window are explicitly flagged as extrapolation in the appendix; the relation graph is d×d, so very high dimensions become a bottleneck, and the paper itself scopes the design to long-context, moderate-dimension problems; two of the three applications reproduce known results rather than produce new ones, and the genuinely new claim, the experience-dependent redundancy effect in neural data, rests on internal consistency without independent validation; edge gates start nearly closed and open only on contextual evidence, so exact disconnection under independence is not guaranteed, and at very small contexts the graph itself carries sampling noise.