Multimodal Datasets with Controllable Mutual Information

AI summary

Overview

Research area: Statistics and machine learning (stat.ML) — specifically information theory, causal modeling, flow-based generative models, mutual information (MI) estimation, and multimodal self-supervised learning (SSL).

Technical level: Advanced. The paper assumes familiarity with mutual information, pointwise mutual information (PMI), Gaussian covariance algebra, normalizing flows, and flow matching.

Scope: The paper introduces a generative framework for building realistic multimodal datasets in which the mutual information between modalities and a target variable is analytically known and controllable, and demonstrates its use for benchmarking MI estimators, regression, and multimodal SSL study design.

What This Paper Is About

Mutual information is central to how multimodal SSL is thought to work, but the exact relationship between the MI shared by modalities and the quality of learned representations is not well understood — partly because MI is hard to estimate from samples, and because existing test datasets with known MI are built from simple analytic distributions that do not resemble real data. This paper builds a dataset generation pipeline that starts from simple uncorrelated Gaussian "proto-latents," mixes them through interpretable linear causal structures to produce correlated latents with closed-form MI, and then pushes those latents through pretrained invertible generative models so the resulting data are high-dimensional and realistic while retaining the exact same MI.

Key Contributions

1. A three-step framework for controllable MI (u → z → x): uncorrelated Gaussian proto-latents are combined by linear structural equations defined on a directed acyclic graph (DAG) to give correlated Gaussian latents with known MI, which are then mapped through per-modality invertible transformations into a realistic feature space.
2. Analytic MI formulas: the framework provides both numerical and closed-form expressions for MI between latents, in terms of the structural-equation hyperparameters, the number of modalities, and the per-modality dimensionality — with the closed forms verified against the numerical calculations.
3. The "template" mechanism: a way to independently control how shared information is distributed across the components of each latent (for example, spread uniformly or concentrated in a single component) while holding total MI fixed, enabling ablations that separate architectural effects from information-content effects.
4. A benchmark suite of realistic datasets with known ground-truth MI, used to evaluate a suite of MI estimators and to test how regression performance on a target variable scales with MI, plus three worked causal "stories" (black hole mass estimation, atmosphere effects, and a scalable many-modality model).

Main Findings

• MI estimators track ground truth: using ten datasets generated from the causal model in Figure 1, with $I(X_1;X_2)$ ranging from 0.0284 to 1.39 and each dataset containing 10,000 paired CIFAR-like images, all eight tested estimators showed regressed MI that increased linearly with ground-truth MI. Reported correlation/RMSE: DV 0.995/0.1094, JS 0.993/0.3049, InfoNCE 0.991/0.1981, MINE 0.993/0.0983, NWJ 0.996/0.0851, SMILE-1 0.998/0.0935, SMILE-5 0.993/0.0974, SMILE-inf 0.993/0.0969.

• Regression improves with MI: on the same ten datasets, where $I(\theta; X_1)$ ranges from 0.134 to 1.73, the best achievable RMSE from a shallow convolutional network trained for 500 epochs (best model selected) decreased monotonically as $I(X_1;\theta)$ increased. The paper reports example error distributions for two representative MI values, where the lower-MI dataset has a wider error spread and the higher-MI dataset a sharper peak, with both distributions centered at zero (no bias). Absolute RMSE values are not reported.

• Framework generates visually uncorrelated but measurably correlated images: Figure 2 shows eight example pairs $(x_1, x_2)$ drawn from a conditional flow-matching model pretrained on CIFAR-10, where image class label is used as a proxy for modality (e.g. "automobile" and "frog"), with $d = 32 \times 32 \times 3 = 3072$. There is no clear visual connection in these pairs, yet the framework specifies their MI exactly.

• Astrophysics example yields interpretable MI structure: for the two-telescope scenario (Event Horizon Telescope and Hubble) with all $\rho$ values set to 1 and $d = 3072$, when $\theta$ is the black hole mass ($\eta_1 = 1, \eta_2 = 0$) the MI values are $I(\theta;X_1) = 2.77$, $I(\theta;X_2) = 0$, and $I(X_1;X_2) = 2.63$; when $\theta$ is the atmospheric effect ($\eta_1 = 0, \eta_2 = 1$) they are $I(\theta;X_1) = 2.77$, $I(\theta;X_2) = 3.33$, and $I(X_1;X_2) = 2.63$. (Note: the running text describes the zero-MI case as the atmospheric scenario, whereas the table as printed lists 0 in the $I(\theta;X_2)$ column for the black hole mass scenario.)

• Scale-out to many modalities works as designed: using a flow-matching model trained on 10 CIFAR class labels, the paper builds correlated tuples for up to $N_z = 10$ modalities with all $\rho$ set to 1 and $d = 3072$. With decay hyperparameters $\alpha$ and $\beta$, $I(X_1;X_i)$ decays faster as $\alpha$ increases and $I(\theta;X_i)$ decays faster as $\beta$ increases, as expected. Extending beyond 10 modalities is described as straightforward.

• Templates decouple information amount from information placement: because MI does not specify how shared information is spread across the components of a random variable, templates let the authors vary information distribution independently of total MI — a capability the paper frames as a prerequisite for clean ablation studies.

Methodology in Plain English

The pipeline has three stages. First, sample a large vector of independent standard normal variables, called proto-latents. Second, mix them with a small set of linear coefficients arranged according to an interpretable causal diagram: some proto-latents act as "common causes" that feed both a target scalar $\theta$ and the latents for each modality, while others only feed individual modalities. Because the mixing is linear, the resulting latents form a multivariate Gaussian whose covariance matrix is $\mathbf{A}\mathbf{A}^{\top}$, and Gaussian MI is a simple function of block covariance determinants using only logarithms and determinants — so the MI can be computed in closed form, and further simplified formulas are derived for specific causal structures.

Third, split the latent vector into blocks, one per modality, and pass each block through its own invertible transformation. These transformations are continuous-time normalizing flows trained with flow matching on CIFAR-10, with image class used as a stand-in for modality. Since the maps are continuous and bijective, MI is preserved exactly (a consequence of the data-processing inequality), so the realistic images inherit the known MI of the latents. Templates generalize the mixing step by replacing scalar coefficients with $d$-dimensional vectors, controlling whether a common cause's information lands uniformly across a latent or is concentrated in one component — for instance, concentrating a black hole's signal near the galactic center while spreading atmospheric effects everywhere.

Why This Matters

Impact on research: prior datasets with analytically tractable MI are simple synthetic distributions, so MI estimators get validated on cases unlike real data. This framework supplies realistic, high-dimensional, multimodal data with exact ground-truth MI, giving both MI-estimation research and multimodal SSL research a controlled testbed. It also directly targets the gap that contrastive methods like InfoNCE have an information-theoretic interpretation while non-contrastive methods such as multimodal masked modeling and the $N > 2$ multimodal regime do not.

Real-world applications (as described or implied by the paper):

• Multi-detector astrophysics — for example, combining radio data from the Event Horizon Telescope with Hubble observations to estimate a supermassive black hole's mass while accounting for atmospheric variability.
• Simulating realistic observations in scientific domains where the amount of shared information between instruments can be specified exactly.
• Designing and stress-testing mutual information estimators before deploying them on real datasets.
• Designing multimodal SSL architectures and objectives, especially as the number of input modalities grows.

Industry relevance: teams building multimodal foundation models need to know how much shared information between modalities is enough, and how it should be distributed, for a given architecture and training objective to succeed. Controlled benchmarks in this framework let those choices be tested reproducibly rather than inferred after the fact from expensive large-scale training runs.

Future Directions

• Applying the framework to study how the distribution of shared information across modalities affects the quality of learned SSL embeddings, including in the $N > 2$ regime and for non-contrastive objectives.
• Using templates in controlled ablation studies that isolate algorithmic choices (such as data augmentation and training objectives) from architectural inductive biases.
• Instantiating the astrophysics example more concretely, with the flow models trained on realistic instrument data rather than CIFAR-10 classes as modality proxies.
• Scaling the many-modality model beyond the $N_z = 10$ demonstrated, and continuing to evaluate whether MI estimators retain their accuracy when ground-truth MI is known but the data are highly realistic and high-dimensional.

Target Audience

Researchers working on information-theoretic machine learning, mutual information estimation, and multimodal self-supervised learning; practitioners who need controlled benchmarks for multimodal model design; and scientists in data-rich multi-instrument fields such as astrophysics who want synthetic but realistic data with a precisely specified information structure. Readers without a background in information theory or generative modeling will find the mathematical core demanding, though the high-level framework is describable in casual terms.

Authors’ abstract

We introduce a framework for generating highly multimodal datasets with explicitly calculable mutual information (MI) between modalities. This enables the construction of benchmark datasets that provide a novel testbed for systematic studies of mutual information estimators and multimodal self-supervised learning (SSL) techniques. Our framework constructs realistic datasets with known MI using a flow-based generative model and a structured causal framework for generating correlated latent variables. We benchmark a suite of MI estimators on datasets with varying ground truth MI values and verify that regression performance improves as the MI increases between input modalities and the target value. Finally, we describe how our framework can be applied to contexts including multi-detector astrophysics and SSL studies in the highly multimodal regime.