InfoAtlas: A Foundation Model for Zero-Shot Statistical Dependence Estimate
Abstract
Measuring statistical dependency between high-dimensional random variables is a fundamental task in data science and machine learning. Neural mutual information (MI) estimators offer a promising avenue, but they typically require costly iterative optimization for each new dataset, making them impractical for real-time applications. We present InfoAtlas, a foundation model-like architecture that eliminates this bottleneck by directly inferring MI in a single forward pass. Pretrained on large-scale synthetic data with rich dependence patterns, InfoAtlas learns to identify diverse dependence structures and predict MI directly from the dataset. Comprehensive experiments demonstrate that InfoAtlas matches state-of-the-art neural estimators in accuracy while achieving 100× speedup, can flexibly handle varying dimensions and sample sizes through a single unified model, and generalizes effectively to complex, real-world scenarios. By reformulating MI estimation as an inference task, InfoAtlas establishes a foundation for real-time dependency analysis.
Lay Summary
Measuring statistical dependence between high-dimensional random variables is a fundamental problem in data science and machine learning. We develop a foundation-style model that enables direct prediction of the strength of dependence between high-dimensional variables, eliminating the need for costly per-dataset optimization. This enables real-time assessment of complex relationships between variables across domains such as finance, medicine, science, and AI.