DiScoFormer: 一个用于密度和分数的Transformer，跨分布适用 | Ai2

机器学习和科学中的许多问题都可以归结为同一项任务：你有一组数据点，想要恢复它们所来自的分布——哪些值常见，哪些值罕见。确定这个分布意味着估计两个量：分布的密度，以及随着维度增加而更有用的得分（score）。密度是直方图的平滑版本——在点聚集的地方高，在点稀少的地方低。得分——对数密度的梯度——指向密度上升最快的方向：沿着得分移动一个点，它会朝向更可能的区域。

基于扩散的生成模型（如 Stable Diffusion 和 DALL-E 等 AI 图像生成器背后的技术）从随机噪声开始，反复跟随得分，将噪声转化为逼真的图像。同样的得分驱动着贝叶斯采样和用于模拟等离子体等系统的粒子模拟。

从有限样本中提取密度和得分具有挑战性，当今的工具迫使人们在泛化能力和准确性之间做出权衡。一种经典方法，核密度估计（KDE），根据周围的数据点计算任意位置的密度：数据点越近、越多，密度就越高。它无需训练，适用于任何分布，但其准确性会随着维度增加而急剧下降。另一种方法是神经得分匹配模型，它经过训练可以预测得分，即使在高维空间中也能保持准确性，但每个模型都需要学习分布，并且必须为另一个分布从头开始重新训练。

我们提出了一种名为 DiScoFormer（密度与得分变换器） 的新解决方案——给定一组数据点，该模型可以在一次前向传播中估计分布的密度和得分，而无需重新训练。

训练用于密度和得分估计的变换器

DiScoFormer 使用堆叠的 transformer 块层，将整个样本映射到其背后分布的密度和得分。该模型利用交叉注意力（cross-attention），使其能够在任意点（而不仅仅是有数据的地方）评估密度和得分。得分和密度之间存在数学关系：得分是密度对数的梯度。我们通过一个共享主干网络和两个输出头（一个用于密度，一个用于得分）来利用这一点。

这种耦合不仅仅节省了参数。得分头必须在每个查询点匹配对数密度头的梯度，因此它们之间的任何差距都构成一个无标签的一致性损失。我们在推理时利用这一点——固定上下文，对该一致性损失进行几步梯度更新，DiScoFormer 就能即时适应分布外的输入，无需真实密度或得分。

transformer 架构适合这项任务有其数学原因。核密度估计只有一个带宽——每个点的影响范围，预先固定并在各处相同地应用。注意力机制是它的严格泛化：我们分析性地证明，单个注意力头的权重近似于数据上的高斯核，因此一个交叉注意力块已经可以复现 KDE 的密度和得分。在此基础上，模型更进一步，同时学习多个这样的尺度，并使它们适应数据。DiScoFormer 并没有将经典方法丢弃为黑箱，而是将 KDE 作为特例包含在内并加以改进。

我们用什么数据来训练 DiScoFormer？我们主要依赖高斯混合模型（GMM），原因有二。首先，GMM 是通用密度逼近器——只要有足够多的分量，它们就能以任意小的误差匹配任何平滑分布。其次，GMM 具有闭式密度和得分，因此我们总能获得精确的目标进行监督。我们利用这两个特性，为每个批次生成一个新的 GMM，为模型提供几乎无限的目标分布示例，并根据给定 GMM 的精确密度和得分进行监督。

性能

总体而言，DiScoFormer 在密度和得分估计方面都优于 KDE，而且差距恰恰在 KDE 表现不佳的地方扩大。在 100 维空间中，差距明显——与最佳手动调参的 KDE 相比，它将得分误差降低了约 6.5 倍，密度误差降低了超过 37 倍，并且随着样本增加，性能持续提升，而 KDE 则会出现内存不足。它还能很好地泛化到训练数据之外，在比训练时见过的模式更多的混合分布以及非高斯形状（如拉普拉斯分布和学生 t 分布）上保持准确性。KDE 的主要优势仍然是速度，尤其是在数据集较小的情况下。

我们认为 DiScoFormer 最有前景的部分在于，得分估计是许多领域的共同依赖，例如生成建模、贝叶斯推理和科学计算。一个预训练的、即插即用的估计器，能在高维空间中保持准确性，并消除每个问题重新训练的需要，可以同时降低所有这些领域的成本——一个模型，在得分和密度出现的任何地方重复使用。

我们鼓励您阅读我们的技术报告以获取更多详细信息。

订阅以每月接收关于 Ai2 最新动态的更新。