DiScoFormer：一个用于密度与得分的Transformer，跨分布通用

](https://huggingface.co/Ai2Comms)

• 训练用于密度和分数估计的 Transformer
• 性能表现

📄 技术报告：arxiv.org/abs/2511.05924

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

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

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

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

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

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

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

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

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

性能表现

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

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

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