DiScoFormer统一密度与评分估计

关键信息

文章认为，attention 可以推广 kernel density estimation，甚至从理论上说明单个 attention head 能近似高斯核。DiScoFormer 采用共享的 transformer 主干和两个输出头——一个预测密度，一个预测 score——并使用 Gaussian Mixture Models 训练，因为它们能提供精确的密度与 score 监督信号。

资讯摘要

这篇文章把密度估计描述为机器学习和科学计算中的一个通用问题：给定一批样本，重建它们来自的分布。文中解释了密度表示某个区域有多常见，而 score 是对 log-density 求梯度后的结果，它会指向更可能的区域。这样的概念是 diffusion 生成模型的核心，因为这类模型会通过不断沿着 score 方向移动，把随机噪声逐步变成真实样本；它们也用于贝叶斯采样和粒子模拟。文章对比了两类现有方法：kernel density estimation（KDE）实现简单、对分布类型几乎不设限制，但在高维空间里效果会迅速变差；神经网络式的 score-matching 模型在高维下更稳，但通常每遇到一个新分布都要重新训练。DiScoFormer 被提出为一种统一方案：它是一个 transformer，可以在一次前向传播中同时估计密度和 score，而且不需要为不同分布重新训练。

模型使用多层 transformer block 和 cross-attention，因此既能对已有样本位置做估计，也能对任意查询点进行评估。由于密度和 score 之间存在数学关系，作者让两个输出头共享 backbone，并把两者之间的差异作为一种无标签的一致性损失。推理阶段，模型还能在保持上下文不变的情况下，对这个一致性损失做几步梯度更新，从而对分布外输入进行即时适配。文章还指出，attention 在数学上与 KDE 有紧密联系，单个 attention head 就能近似高斯核，因此 KDE 可以看作这种 transformer 形式的特例。训练时，作者为每个 batch 采样一个新的 Gaussian Mixture Model，以便同时获得精确的密度和 score 监督，并让模型接触到近乎无限多的目标分布。

来源与参考

1. 原始链接
2. DiScoFormer: One transformer for density and score, across distributions