DiScoFormer：跨分布的密度和评分一体化Transformer

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

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

从有限样本中提取密度和分数很具挑战性，当今工具在通用性和精度间存在权衡。一个经典方法是核密度估计（KDE），从其周围的数据点计算任何位置的密度：数据点越近、越多，密度越高。它无需训练，适用于任何分布，但精度随维数增长而急剧下降。另一方面，训练来预测分数的神经分数匹配模型即使在高维也能保持精度，但每个都需要学习该分布，对另一个分布必须从头重新训练。

我们推出了一个名为 DiScoFormer（密度和分数 Transformer）的新方案——一个模型，给定一组数据点，可以在单个前向传递中估计分布的密度和分数，无需重新训练。

训练用于密度和分数估计的 Transformer

DiScoFormer 使用堆叠的 Transformer 块将整个样本映射到其背后分布的密度和分数。该模型利用交叉注意力，允许它在任何点上评估密度和分数——不仅仅是有数据的地方。分数和密度共享一个数学关系：分数是密度对数的梯度。我们通过拥有一个共享骨干和两个输出头来利用这一点，一个用于密度，一个用于分数。

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

Transformer 架构适合这项任务有一个数学原因。核密度估计有一个单一的带宽——每个点的影响范围，预先固定并处处相同应用。注意力是它的严格推广：我们通过分析证明单个注意力头的权重几乎是数据上的高斯核，因此一个交叉注意力块就已经可以重现 KDE 的密度和分数。从那里模型进一步发展，同时学习多个这样的尺度并将它们适应到数据。DiScoFormer 不是抛弃经典方法转向黑盒，而是将 KDE 作为特例并在其基础上改进。

我们用什么数据训练 DiScoFormer？我们依靠高斯混合模型有两个主要原因。首先，GMM 是通用密度逼近器——具有足够的分量，它们可以将基本上任何平滑分布逼近到任意小的误差。其次，GMM 具有闭式密度和分数，因此我们总是有确切的监督目标。我们通过为每个批次绘制一个新的 GMM 来利用这两个性质，给模型提供几乎无限的目标分布示例，并根据给定 GMM 的确切密度和分数进行监督。

总体来看，DiScoFormer 在密度和分数估计上都击败了 KDE，差距在 KDE 陷入困境的地方扩大。在 100 维中，差距很明显——对比手调最优的 KDE，它将分数误差降低约 6.5 倍，密度误差降低超过 37 倍，并且随着样本增加继续改进，而 KDE 则耗尽内存。它也远远超出其训练数据的范围，在具有比训练期间见过的更多模式的混合分布上保持准确，以及在拉普拉斯和 Student-t 等非高斯形状上。KDE 的主要优势仍然是速度，特别是当数据集很小时。

我们发现 DiScoFormer 最有前景的部分是，分数估计是许多领域的共同依赖，如生成建模、贝叶斯推断和科学计算。一个预训练的即插即用的估计器，保持在高维上的精度并消除每个问题重新训练的需要，可以同时削减所有这些领域的成本——一个模型，在分数和密度出现的任何地方重复使用。

我们鼓励你阅读我们的技术报告了解更多详情。

更多来自本作者

OlmoEarth 平台：行星规模的地球空间推断

构建 Shippy 教会我们关于构建 Agent 的经验

· 登录或注册评论