InfoAtlas：用于零样本统计依赖估计的基础模型

# 零样本统计依赖性估计的基础模型 来源：https://arxiv.org/html/2606.00241 Yanzhi Chen\*Hanxiang Ren\*Qunsong Zeng\*Youyi Zheng\*Adrian Weller\*Kaibin Huang\*Yanchao Yang\* ###### 摘要 衡量高维随机变量之间的统计依赖性数据科学和机器学习中的一项基础任务。神经互信息（MI）估计器提供了一种有前景的途径，但它们通常需要对每个新数据集进行昂贵的迭代优化，这使得它们不适用于实时应用。我们提出了InfoAtlas，一种类似基础模型的架构，通过直接在一次前向传播中推断MI来消除这一瓶颈。在大规模合成数据上预训练，这些数据包含丰富的依赖模式，InfoAtlas学会识别多样的依赖结构并直接从数据集中预测MI。全面的实验表明，InfoAtlas在准确性上与最先进的神经估计器相当，同时实现了100倍的加速，可以通过单一统一模型灵活处理不同维度和样本大小的情况，并能有效泛化到复杂的现实世界场景。通过将MI估计重新表述为一项推理任务，InfoAtlas为实时依赖性分析奠定了基础。项目页面：InfoAtlas-project (https://datou30.github.io/InfoAtlas-page/) 互信息，统计依赖性，基础模型 ## 1 引言 理解变量之间的统计依赖性数据科学和机器学习的基础。量化变量之间的相互影响揭示了驱动复杂系统的隐藏结构和因果机制。应用涵盖广泛领域：在医疗保健中，识别生活方式因素与疾病风险之间的依赖关系能够实现个性化预防策略(Du et al., 2024 (https://arxiv.org/html/2606.00241#bib.bib190))；在自动驾驶中，建模传感器信号与路况之间的依赖关系能提升安全性和决策能力(Maanpää et al., 2025 (https://arxiv.org/html/2606.00241#bib.bib192))；在生物学中，评估蛋白质序列之间的依赖性能够揭示其功能关系的见解(Gowri et al., 2024 (https://arxiv.org/html/2606.00241#bib.bib180))；在机器人学中，最大化观测状态之间的统计依赖性已被证明对策略发现有用(Zhou and Yang, 2024 (https://arxiv.org/html/2606.00241#bib.bib194))。 互信息（MI）(Shannon, 1948 (https://arxiv.org/html/2606.00241#bib.bib23))长期以来一直是衡量依赖性的有原则的度量，它能够以可解释的比特单位唯一地捕捉多元变量的复杂非线性关系。其通用性使其成为数据分析、生成建模和表示学习中的核心工具(Chen et al., 2016 (https://arxiv.org/html/2606.00241#bib.bib38); Oord et al., 2018 (https://arxiv.org/html/2606.00241#bib.bib107); Chen et al., 2022 (https://arxiv.org/html/2606.00241#bib.bib165))。然而，从经验样本中计算MI是出了名的困难：闭式解仅存在于某些分布中，而神经估计器(Belghazi et al., 2018 (https://arxiv.org/html/2606.00241#bib.bib170); Choi and Lee, 2020 (https://arxiv.org/html/2606.00241#bib.bib222); Franzese et al., 2023 (https://arxiv.org/html/2606.00241#bib.bib185); Tschannen et al., 2019 (https://arxiv.org/html/2606.00241#bib.bib101); Chen et al., 2020 (https://arxiv.org/html/2606.00241#bib.bib181); Tsai et al., 2020 (https://arxiv.org/html/2606.00241#bib.bib28)) 需要对每个数据集进行昂贵的基于梯度的优化，使其不适用于实时或大规模应用。 参见图注 图1：概念对比：先前方法 vs 我们的方法。现有的神经 MI 估计器（左）需要迭代的基于梯度的优化来为每个新数据集训练一个神经网络。相比之下，我们使用一个*预训练*的架构来直接在一次前向传播中生成 MI 估计（右），消除了每个数据集的训练，在保持可比精度的同时实现了加速。 在这项工作中，我们引入了InfoAtlas，一种用于快速准确估计*多元*随机变量之间统计依赖性的基础模型风格架构。InfoAtlas在单次前向传播中预测依赖性强度——这种能力让人联想到基础模型(Hollmann et al., 2025 (https://arxiv.org/html/2606.00241#bib.bib238); Comanici et al., 2025 (https://arxiv.org/html/2606.00241#bib.bib210))。这种能力是通过在覆盖广泛依赖结构和数据模式的大规模合成数据集上进行大规模预训练获得的，使InfoAtlas能够直接推断统计关系，而无需每个数据集的优化。关键的是，InfoAtlas保持了完全可微性，便于无缝集成到更大的 AI 流水线中。大量实验表明，InfoAtlas从合成基准有效泛化到复杂的现实世界数据，准确捕捉广泛的依赖性，成为快速理解变量关系的多功能工具。我们的主要贡献是： - • 我们引入了InfoAtlas，第一个用于*多元*变量之间互信息零样本估计的预训练架构。InfoAtlas达到了与最先进的神经方法相当的精度，无需任何基于梯度的优化，并且能够灵活处理不同维度和样本大小的变量，使用*单个*模型。 - • 我们提出了一种基于注意力双路径超网络的架构，该架构在覆盖多样依赖性结构的大规模合成数据集上预训练。这种设计使InfoAtlas能够在单次推理步骤中预测依赖性强度，并能有效泛化到未见过的现实世界场景，无需特定任务的微调。 - • 我们在合成基准和现实世界任务上全面评估了InfoAtlas，包括独立性检验、CLIP 嵌入分析(Radford et al., 2021 (https://arxiv.org/html/2606.00241#bib.bib160))、运动轨迹建模和机器人操作。结果展示了其鲁棒的性能和对广泛依赖性谱的准确感知。 ## 2 问题陈述 在这项工作中，我们考虑量化两个*多元*随机变量 \(\mathbf{x} \in \mathbb{R}^{d_x}\) 和 \(\mathbf{y} \in \mathbb{R}^{d_y}\) 之间的统计依赖性问题，其中 \(d_x \geq 1\) 且 \(d_y \geq 1\)。 #### 通过互信息衡量依赖性。 互信息（MI）提供了量化多元变量之间统计依赖性的有原则的度量。与仅捕捉线性关系的线性相关系数不同，MI 有效捕捉线性和非线性相关性。形式上，MI 定义为联合分布 \(p_{\mathbf{x},\mathbf{y}}\) 与边缘乘积 \(p_{\mathbf{x}} \otimes p_{\mathbf{y}}\) 之间的 Kullback-Leibler（KL）散度(Kullback, 1997 (https://arxiv.org/html/2606.00241#bib.bib34))： \[ \mathbb{I}(\mathbf{x},\mathbf{y}) = \text{KL}(p_{\mathbf{x},\mathbf{y}} \| p_{\mathbf{x}} \otimes p_{\mathbf{y}}) \tag{1} \] \[ = \int_{\mathcal{Y}} \int_{\mathcal{X}} p_{\mathbf{x},\mathbf{y}}(\mathbf{x},\mathbf{y}) \log\left( \frac{p_{\mathbf{x},\mathbf{y}}(\mathbf{x},\mathbf{y})}{p_{\mathbf{x}}(\mathbf{x}) p_{\mathbf{y}}(\mathbf{y})} \right) d\mathbf{x} d\mathbf{y}. \] 强相关性表现为 \(p(\mathbf{x},\mathbf{y})\) 与 \(p(\mathbf{x})p(\mathbf{y})\) 之间的显著散度，产生较大的 MI，而不相关变量满足 \(p(\mathbf{x},\mathbf{y}) \approx p(\mathbf{x})p(\mathbf{y})\)，导致 MI 接近于零。虽然 MI 提供了一种有原则的依赖性度量，但它很少允许闭式解，除了某些已知分布(Czyż et al., 2023a (https://arxiv.org/html/2606.00241#bib.bib161), b (https://arxiv.org/html/2606.00241#bib.bib193))。因此，实际应用需要从 \(p_{\mathbf{x},\mathbf{y}}\) 中抽取的有限样本 \(\mathcal{D} = \{\mathbf{x}^i, \mathbf{y}^i\}_{i=1}^n\) 进行估计。最近的进展产生了强大的神经估计器(Belghazi et al., 2018 (https://arxiv.org/html/2606.00241#bib.bib170); Duong and Nguyen, 2023 (https://arxiv.org/html/2606.00241#bib.bib27); Franzese et al., 2023 (https://arxiv.org/html/2606.00241#bib.bib185); Tsai et al., 2020 (https://arxiv.org/html/2606.00241#bib.bib28); Poole et al., 2019 (https://arxiv.org/html/2606.00241#bib.bib39); Song and Ermon, 2019 (https://arxiv.org/html/2606.00241#bib.bib187); Letizia et al., 2024 (https://arxiv.org/html/2606.00241#bib.bib26); Tsur et al., 2023a (https://arxiv.org/html/2606.00241#bib.bib188))，其中最突出的利用 Donsker-Varadhan (DV) 表示(Donsker and Varadhan, 1983 (https://arxiv.org/html/2606.00241#bib.bib30))： \[ \mathbb{I}(\mathbf{x},\mathbf{y}) \coloneqq \sup_{\theta} \mathbb{E}_{p_{\mathbf{x},\mathbf{y}}}[\theta] - \log(\mathbb{E}_{p_{\mathbf{x}} \otimes p_{\mathbf{y}}}[e^{\theta}]), \tag{2} \] 其中 \(\theta: \mathcal{X} \times \mathcal{Y} \rightarrow \mathbb{R}\) 是一个批评函数。互信息神经估计（MINE） (Belghazi et al., 2018 (https://arxiv.org/html/2606.00241#bib.bib170)) 将 \(\theta\) 参数化为一个神经网络，并通过基于梯度的优化来逼近上确界。除了 MINE，还存在大量基于不同边界和学习目标的神经估计器；参见 §6 (https://arxiv.org/html/2606.00241#S6)。 #### 实时 MI 估计的挑战。 尽管在理论和算法上存在差异，但所有现有的神经估计器在实践中都共享一个关键的计算瓶颈：它们需要对每个传入的数据集 \(\mathcal{D} = \{\mathbf{x}^i, \mathbf{y}^i\}_{i=1}^n\) 通过梯度下降从头训练一个网络 \(\theta\)： \[ \theta^{t+1} \leftarrow \theta^t - \eta \nabla_{\theta^t} \mathcal{L}(\theta^t), \quad t=1,...,T \tag{3} \] 其中 \(\mathcal{L}(\theta)\) 是特定于估计器的目标函数（例如，对于 MINE 是负 DV 界）。实现准确的 MI 估计通常需要数千步梯度，导致 \(\mathcal{O}(T)\) 的计算复杂度。这种高昂的成本限制了实时应用，如高频金融相关性监控或大规模基因组筛选。最近的 InfoNet (Hu et al., 2024 (https://arxiv.org/html/2606.00241#bib.bib171)) 旨在通过预训练一个网络以通过查找表直接输出最优批评值来解决这种低效问题，从而消除了推理时的优化。然而，InfoNet 从根本上局限于单变量输入，将其扩展到 \(d\) 维变量将要求在其查找表中存储 \(\mathcal{O}(e^d)\) 个值，即使对于 \(d=8\) 也变得不可行，并且它无法处理具有不同数据维度的数据。这些局限性促使我们采用根本不同的方法进行实时统计依赖性测量，其中一个统一模型被开发用于直接处理具有不同维度和样本大小的多元数据。 ## 3 方法 我们提出了InfoAtlas，一种预训练架构，用于解决上述多元随机变量 \(\mathbf{x} \in \mathbb{R}^{d_x}\) 和 \(\mathbf{y} \in \mathbb{R}^{d_y}\) 之间实时相关性估计的挑战。与需要迭代优化的现有神经估计器不同，InfoAtlas 直接在一次前向传播中输出互信息（MI）。这种能力源于两个关键创新：(i) 一种基于双路径注意力超网络的架构，它直接从具有不同大小和维度的观测样本中生成特定于分布的批评参数；(ii) 一种使用多样化合成分布的全面预训练策略，这确保了跨不同应用领域的泛化能力。 ### 3.1 直接最优批评生成 我们的关键创新是将 MI 估计从测试时的优化问题重新表述为借助超网络的直接推理任务。具体来说，给定一个从未知联合分布中抽取的数据集 \(\mathcal{D} = \{(\mathbf{x}^i, \mathbf{y}^i)\}_{i=1}^n\)，InfoAtlas 采用一个基于注意力的超网络 \(\mathcal{H}: \mathcal{D} \mapsto \Theta\)，它通过单次前向传播直接输出 Donsker-Varadhan 表示（公式 (2) (https://arxiv.org/html/2606.00241#S2.E2)）中最优批评网络的完整参数集 \(\theta^*\)¹： \[ \theta^* = \mathcal{H}(\mathcal{D}) = \mathcal{H}(\{(\mathbf{x}^i, \mathbf{y}^i)\}_{i=1}^n) \tag{4} \] 然后通过以下公式获得经验 MI 估计： \[ \hat{\mathbb{I}}_\theta(\mathbf{x},\mathbf{y}) = \frac{1}{n} \sum_{i=1}^n \theta(\mathbf{x}^i, \mathbf{y}^i) - \log\Big( \frac{1}{n} \sum_{j=1}^n e^{\theta(\mathbf{x}^j, \mathbf{y}^{\pi(j)})} \Big), \tag{5} \] 其中 \(\{(\mathbf{x}^j, \mathbf{y}^{\pi(j)})\}_{j=1}^n\) 表示边缘对，\(\pi\) 是索引 \(\{1,...,n\}\) 的随机排列。这消除了神经 MI 估计器所需的迭代梯度更新，同时避免了 InfoNet 查找表方法所需的指数级值存储。这种架构转变从根本上将计算复杂度从 \(\mathcal{O}(T)\) 步梯度（其中 \(T\) 是优化迭代次数）改变为 \(\mathcal{O}(1)\) 次前向传播。 参见图注 图2：InfoAtlas 估计流水线。步骤 1：我们用噪声填充输入维度，以确保所有变量具有相同维度，同时允许灵活的样本大小。步骤 2：一个双路径超网络 \(\mathcal{H}\)——包含联合分支和边缘分支——提取与 D-V 公式（公式 (2) (https://arxiv.org/html/2606.00241#S2.E2)）一致的特征。交叉注意力整合这些特征，然后使用参数生成 MLP 产生批评参数。步骤 3：对联合样本和边缘样本应用经验 D-V 公式（公式 (5) (https://arxiv.org/html/2606.00241#S3.E5)），其中边缘样本通过索引排列获得，以估计 MI。该流水线实现了无需基于梯度优化的单次传递估计。 超网络 \(\mathcal{H}\) 利用注意力机制(Vaswani et al., 2017 (https://arxiv.org/html/2606.00241#bib.bib12))，并由以下关键模块组成： #### 联合分布路径 处理 \(n\) 个成对样本 \(\{(\mathbf{x}^i, \mathbf{y}^i)\}_{i=1}^n\) 以提取相关性模式。