🤖 machine learning

InfoAtlas: A Foundation Model for Zero-Shot Statistical Dependence Estimate

InfoAtlas is a foundation model that achieves real-time, zero-shot statistical dependence estimation by directly inferring mutual information in a single forward pass, offering state-of-the-art accuracy with a 100-fold speedup over traditional iterative neural estimators.

Original authors: Zhengyang Hu, Yanzhi Chen, Hanxiang Ren, Qunsong Zeng, Youyi Zheng, Adrian Weller, Kaibin Huang, Yanchao Yang

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhengyang Hu, Yanzhi Chen, Hanxiang Ren, Qunsong Zeng, Youyi Zheng, Adrian Weller, Kaibin Huang, Yanchao Yang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to figure out if two people at a party are secretly communicating. You have a list of their actions (what they said, where they looked, what they drank).

The Old Way (Traditional Methods):
In the past, to figure out if these two people are connected, you would have to hire a new detective for every single new party. This detective would spend hours (or days) studying the specific guests, running endless tests, and adjusting their theory until they finally found a pattern. If you wanted to check a second party, you'd have to hire a new detective and start all over again. It's accurate, but it's incredibly slow and expensive.

The New Way (InfoAtlas):
The paper introduces InfoAtlas, which is like hiring a "Super Detective" who has already studied millions of different parties, from tiny gatherings to massive festivals. This Super Detective has seen every possible way people can interact.

Now, when you bring in data from a new party, you don't need to hire a new detective or wait for them to study the guests. You just show the data to InfoAtlas, and it instantly says, "Yes, these two are definitely connected," or "No, they are just strangers." It does this in a single glance (a "forward pass") without needing to do any extra homework.

How It Works (The Magic Tricks)

1. The "Atlas" of Patterns
The authors didn't just train this model on real-world data; they built a massive "atlas" (a map) of synthetic data. Imagine generating millions of fake scenarios where variables (like temperature and ice cream sales) are linked in every weird, complex, and nonlinear way imaginable.

The Analogy: It's like training a chef by having them cook every possible combination of ingredients in a virtual kitchen before they ever touch real food. Because they've seen every possible "flavor" of connection, they can recognize the taste of a new dish immediately.

2. The Dual-Path Brain
InfoAtlas looks at the data in two different ways simultaneously, like a detective with two pairs of eyes:

Eye 1 (The Joint View): Looks at the pairs together (e.g., "What did Person A do while Person B did X?"). This looks for direct links.
Eye 2 (The Separate View): Looks at them separately (e.g., "What does Person A usually do on their own?"). This establishes a baseline of "random chance."
The Comparison: The model compares the "Joint View" against the "Separate View." If the two people act differently when together than they do apart, the model knows they are dependent.

3. The "Noise" Trick for Big Data
Sometimes, the data is too big (like a video with thousands of points). InfoAtlas has a limit on how big a single "bite" of data it can eat at once.

The Analogy: Imagine trying to eat a giant pizza. Instead of trying to swallow the whole thing, InfoAtlas slices the pizza into smaller, manageable pieces. It tastes a few slices, averages the flavor, and tells you how the whole pizza tastes. This allows it to handle massive, high-dimensional data (like 1,000+ dimensions) without getting overwhelmed.

Why This Matters (The Results)

The paper claims three main victories:

Speed: It is 100 times faster than the old methods. While traditional methods might take minutes or hours to analyze one dataset, InfoAtlas does it in a fraction of a second.
Accuracy: Despite being instant, it is just as accurate as the slow, hard-working methods. It correctly identifies complex relationships that simple math (like linear correlation) would miss.
Versatility: It works on data of different sizes and shapes without needing to be retrained. Whether you have 100 data points or 10,000, or variables with 5 dimensions or 20, the same model handles it all.

Real-World Examples from the Paper

The authors tested this "Super Detective" on several real-world scenarios:

Robotics: They used it to figure out which parts of a robot's movement are linked. If a robot is picking up a cube, the model correctly identified that the hand and the cube move together, helping the robot learn better.
Video Analysis: They analyzed video footage of moving objects. The model could instantly tell which dots in the video belonged to the same object (like a person walking) versus dots on different objects, just by looking at how their paths were statistically linked.
Image & Text: They checked if images and their descriptions (captions) were truly related. The model successfully detected when the connection between an image and its text was weak (by adding noise) and strong.

The Bottom Line

InfoAtlas changes the game by turning a slow, repetitive math problem (optimizing a model for every new dataset) into a fast, one-step recognition task. It's like upgrading from a librarian who has to search the stacks for every single book you ask about, to a librarian who has memorized the entire library and can tell you where a book is instantly.

Note: The paper explicitly states that while it works well for moderate sample sizes (500+), it might struggle with very tiny datasets (under 400 samples), and for extremely high-dimensional data, it relies on the "slicing" technique to function.

Technical Summary: InfoAtlas

1. Problem Statement

Measuring statistical dependence between high-dimensional random variables is a fundamental task in data science, with Mutual Information (MI) serving as a principled metric capable of capturing both linear and nonlinear relationships. However, estimating MI from empirical samples is computationally prohibitive for real-time applications.

Existing neural MI estimators (e.g., MINE, InfoNCE, MINDE) rely on the Donsker-Varadhan (DV) representation, which requires training a neural network (the "critic") from scratch via iterative gradient-based optimization for every new dataset. This process involves thousands of gradient steps, resulting in $O(T)$ computational complexity per dataset. While recent work like InfoNet attempted to bypass this via pretraining, it is fundamentally limited to univariate inputs due to its reliance on lookup tables, which become intractable for dimensions $d > 1$ . Consequently, there is a lack of scalable, real-time methods for estimating MI in multivariate settings with varying dimensions and sample sizes.

2. Methodology: InfoAtlas

The authors propose InfoAtlas, a foundation model-like architecture designed to estimate statistical dependence in a single forward pass without per-dataset optimization. The method reframes MI estimation from an optimization problem into a direct inference task.

2.1. Architecture: Attentive Dual-Path Hypernetwork

InfoAtlas employs a hypernetwork $H$ that maps a dataset $D = \{(x_i, y_i)\}_{i=1}^n$ directly to the parameters $\theta^*$ of the optimal critic network required for the DV representation. The architecture consists of three key components:

Joint Distribution Path: Processes paired samples $\{(x_i, y_i)\}$ as a sequence of tokens. Using cross-attention with a learnable query vector, it aggregates features to capture correlation patterns inherent in the joint distribution $p(x, y)$ . This path utilizes 16 self-attention layers.
Marginal Distribution Path: Processes samples from the product of marginals $p(x)p(y)$ by breaking the pairing (via random permutation). It employs bidirectional cross-attention between projected $x$ and $y$ representations to capture symmetric independence patterns. This path uses 8 self-attention layers.
Integration and Generation Module: Fuses the joint and marginal encodings via cross-attention (where marginal features query joint features). A parameter-generation MLP then maps this fused representation to the flattened parameter vector $\theta$ for the critic network.

Handling Variable Dimensions: To accommodate inputs with varying dimensions $d < D_{max}$ , the model employs a noise padding strategy. Variables are padded with independent Gaussian noise to reach a fixed dimension $D$ . Theoretically, this preserves MI exactly ( $I(x, y) = I([x; n_x], [y; n_y])$ ) when noise is independent of the data.

High-Dimensional Scaling (Slicing): For dimensions exceeding the pretraining limit ( $d > 20$ ), InfoAtlas utilizes Sliced Mutual Information (Sliced MI). It projects high-dimensional data onto multiple low-dimensional subspaces (slices) and aggregates the MI estimates. Unlike traditional methods that require separate optimization for each slice, InfoAtlas processes all slices in a single batched forward pass, reducing complexity from $O(S \cdot T)$ to $O(1)$ .

2.2. Large-Scale Synthetic Pre-training

InfoAtlas is pretrained on a massive "atlas" of synthetic data generated via a meta-distribution $p(D)$ . The generation strategy ensures diversity in two dimensions:

Dependence Diversity: Achieved through random mixtures of copulas (Gaussian and Student's t) with varying correlation matrices and tail dependencies.
Marginal Diversity: Achieved via random flow-based transformations and differentiable copula transformations (softrank) to normalize marginals, ensuring the model learns correlation structures rather than specific marginal distributions.

The pretraining objective minimizes the negative expected empirical MI estimate across the meta-distribution, effectively learning to predict the optimal critic parameters for any dataset drawn from this distribution.

3. Key Contributions

First Zero-Shot Multivariate MI Estimator: InfoAtlas is the first pretrained architecture capable of estimating MI between multivariate variables of varying dimensions and sample sizes without gradient-based optimization.
Dual-Path Hypernetwork Design: The authors propose a novel architecture that directly generates critic parameters from observed samples, eliminating the need for iterative training while maintaining full differentiability for integration into AI pipelines.
Scalability and Efficiency: By leveraging sliced MI and batched inference, the model scales to high-dimensional data (up to 1024D) with a theoretical speedup of $100\times$ compared to state-of-the-art neural estimators.
Comprehensive Evaluation: The method is validated on synthetic benchmarks with known ground-truth MI, independence testing, CLIP embedding analysis, motion trajectory modeling, and robotic manipulation concept discovery.

4. Experimental Results

Accuracy vs. Speed: On synthetic benchmarks with known ground truth, InfoAtlas achieves accuracy comparable to state-of-the-art neural methods (MINE, InfoNCE, MINDE) but executes approximately 300× faster (0.09s vs. 25.9s for MINE on a standard task).
Independence Testing: In high-dimensional independence testing (up to 128 dimensions), InfoAtlas demonstrates strong statistical power, matching gradient-based estimators when sample sizes exceed 500, while significantly outperforming the univariate-limited InfoNet.
Real-World Applications:
- CLIP Analysis: InfoAtlas successfully detects noise-induced reductions in image-text dependence with high sensitivity, outperforming InfoNet in accuracy due to its ability to utilize multi-dimensional slicing.
- Motion Trajectory Modeling: On the PointOdyssey dataset, the model correctly identifies points belonging to the same object based on trajectory correlations, achieving competitive segmentation accuracy in seconds.
- Robotics: In concept discovery tasks for robotic manipulation (ManiSkill 2), policies trained using key states identified by InfoAtlas achieved higher success rates than those using other estimators, even with reduced time budgets.

5. Significance and Claims

The paper claims that InfoAtlas establishes a paradigm shift in statistical dependence estimation. By reformulating MI estimation as a one-step inference task rather than an optimization problem, it enables real-time, scalable dependency analysis.

The authors emphasize that while Sliced MI is an approximation for high dimensions, it provides a favorable accuracy-efficiency trade-off for practical applications where exact MI values are less critical than relative dependence strength. The work positions InfoAtlas as a versatile foundation for real-time applications in healthcare, autonomous driving, biology, and robotics, where rapid assessment of variable relationships is essential.

Limitations Acknowledged:

Small Sample Sizes: Performance degrades when sample sizes are very small ( $n < 400$ ), as the transformer struggles to extract informative signals from limited populations.
High-Dimensional Approximation: For dimensions $d > 20$ , the method relies on slicing, which may miss certain dependence structures or dilute nonlinear dependencies that are clear in the original high-dimensional space.
Pretraining Scope: While effective up to 20 dimensions natively, scaling to extremely high dimensions requires the slicing approximation, which introduces a bias dependent on the projection dimension $k$ and number of slices $S$ .