ALICE: In-context, Zero-shot, Mutual Information Estimation

Mutual-information estimators often require a new neural network for every distribution, making repeated analyses expensive when samples are scarce. We introduce ALICE, a transformer pretrained on synthetic distributions that estimates dependence from a clean sample context with frozen weights. It predicts rectified-flow velocity fields and combines their differences into mutual information. On the 40-task Beyond Normal benchmark, ALICE Base achieves the lowest mean absolute error at matched budgets of 1,000, 5,000, and 10,000 samples.

Turning statistical dependence into velocity differences

For variables $X$ and $Y$ with joint density $p_{XY}$ and marginals $p_X,p_Y$, mutual information measures departure from independence:

$$\operatorname{MI}(X;Y)=\operatorname{KL}\!\left[p_{XY}\,\|\,p_X\otimes p_Y\right].$$

Direct density estimation is difficult in high dimensions. We instead represent a distribution through a field along a path from data to independent Gaussian noise. For $Z_0\sim p_{XY}$, independent $\epsilon\sim\mathcal N(0,I)$, and $t\in[0,1]$, we use

$$Z_t=(1-t)Z_0+t\epsilon, \qquad v_t(z)=\mathbb E[Z_0-\epsilon\mid Z_t=z].$$

The sign convention points toward clean data. Let $z_t=(x_t,y_t)$, and let $v_t(x_t\mid y_0)$ denote the field obtained by noising $X$ while holding $Y$ clean. For smooth densities with finite second moments, finite KL divergence, and vanishing boundary terms, mutual information satisfies

$$\begin{aligned} \operatorname{MI}(X;Y) &=\int_0^1\frac{1-t}{t}\,\mathbb E\Big[ \left\|v_t(z_t)|_X-v_t(x_t\mid y_0)\right\|^2\\ &\hspace{5.6em}+\left\|v_t(z_t)|_Y-v_t(y_t\mid x_0)\right\|^2 \Big]\,\mathrm dt. \end{aligned}$$

Here $|_X$ and $|_Y$ select coordinate blocks. A shared noise draw aligns the three evaluations. The squared differences measure how the field changes when the other block becomes clean evidence. Monte Carlo averaging over held-out samples, times, and noise approximates the integral.

Learning fields from a sample context

ALICE conditions its prediction $v_\theta(z_t,t,m;C)$ on clean samples $C$. A binary mask $m$ selects noised coordinates; the others remain observed. For synthetic distributions $p\sim\mathcal T$, we minimize

$$\mathcal L(\theta)=\mathbb E\!\left[ \frac{\left\|m\odot\big(v_\theta(z_t,t,m;C)-(z_0-\epsilon)\big)\right\|^2}{\|m\|_1} \right].$$

The expectation covers distributions, contexts, clean queries, uniform times, independent Gaussian noise, and nonzero masks. The sampled target is available directly, so pretraining needs neither density values nor mutual-information labels. At inference, disjoint context and evaluation sets avoid reusing query samples as conditioning evidence; a coordinate-wise copula map is fitted on the context.

A fixed bank of induced tokens compresses each coordinate's context. Learned nonlinear feature covariances form a relation graph that carries cross-coordinate dependence. Context-dependent attention grows linearly with sample count at a fixed latent budget; the relation graph retains a quadratic dependence on dimension. A boundary parameterization enforces $v_0(z)=z$, controlling positive bias near the integral's large $(1-t)/t$ weight.

We also use frozen context-conditioned predictors to construct conformal regression intervals.

Accuracy, context length, and model capacity

Every benchmark estimator receives the same sample budget. ALICE reserves 64 samples for queries and uses the remainder as context, with 64 time draws per query and eight independent context draws. Base achieves MAEs of 0.092, 0.063, and 0.060 nats at the three budgets; the best competing values are 0.195, 0.070, and 0.065.

Contexts above the final 2,048-sample training window require extrapolation. Base underestimates spiral embeddings and overestimates dense high-dimensional Gaussians, locating specific limits of the pretrained estimator.

Scientific analyses with the same frozen model

The applications test transfer to data absent from pretraining. In single-cell signaling, contexts contain 1,024 cells; estimated channel capacity compares individual measurements with prefixes of NF-$\mathcal K$B response trajectories.

For promoter localization, we score six-base windows using 1,024 context pairs and 1,024 held-out pairs from a balanced Arabidopsis TATA-promoter dataset, a positive control for localization. In mouse visual cortex, matched 128-flash contexts support separate estimates for each session and flash type.

In the neuroscience data, familiar sessions precede novel sessions by one day, so experience and recording order remain coupled. The scientific applications assess localization and consistency; their underlying mutual information is unknown.