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

Zhengyang Hu, Yanzhi Chen, Hanxiang Ren, Qunsong Zeng, Youyi Zheng, Adrian Weller, Kaibin Huang, Yanchao Yang
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:44799-44818, 2026.

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$\times$ 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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-hu26i, title = {{I}nfo{A}tlas: A Foundation Model for Zero-Shot Statistical Dependence Estimate}, author = {Hu, Zhengyang and Chen, Yanzhi and Ren, Hanxiang and Zeng, Qunsong and Zheng, Youyi and Weller, Adrian and Huang, Kaibin and Yang, Yanchao}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {44799--44818}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/hu26i/hu26i.pdf}, url = {https://proceedings.mlr.press/v306/hu26i.html}, 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$\times$ 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.} }
Endnote
%0 Conference Paper %T InfoAtlas: A Foundation Model for Zero-Shot Statistical Dependence Estimate %A Zhengyang Hu %A Yanzhi Chen %A Hanxiang Ren %A Qunsong Zeng %A Youyi Zheng %A Adrian Weller %A Kaibin Huang %A Yanchao Yang %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-hu26i %I PMLR %P 44799--44818 %U https://proceedings.mlr.press/v306/hu26i.html %V 306 %X 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$\times$ 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.
APA
Hu, Z., Chen, Y., Ren, H., Zeng, Q., Zheng, Y., Weller, A., Huang, K. & Yang, Y.. (2026). InfoAtlas: A Foundation Model for Zero-Shot Statistical Dependence Estimate. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:44799-44818 Available from https://proceedings.mlr.press/v306/hu26i.html.

Related Material