DiScoFormer：一个Transformer同时处理密度与分数，适用于多种分布

来源：https://huggingface.co/blog/allenai/discoformer 返回文章列表 (https://huggingface.co/blog)

Kyle Wiggers 的头像 (https://huggingface.co/Ai2Comms)

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

📄 技术报告： arxiv.org/abs/2511.05924 (https://arxiv.org/abs/2511.05924)

DiScoFormer：一个适用于多种分布的密度与分数Transformer - 谷歌 - 图片1 (https://cdn-uploads.huggingface.co/production/uploads/638e39b249de7ae552d977b5/RsD3JLEW1EABr7i2WCKlF.png)

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

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

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

我们提出了一种名为 DiScoFormer（密度与分数Transformer）的新解决方案 (https://arxiv.org/pdf/2511.05924)——给定一组数据点，该模型只需一次前向传播就能同时估计分布的密度和分数，无需重新训练。

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

DiScoFormer：一个适用于多种分布的密度与分数Transformer - 谷歌 - 图片2 (https://cdn-uploads.huggingface.co/production/uploads/638e39b249de7ae552d977b5/tNRWwYaFG4jyW5f2xWdoY.png)

DiScoFormer 使用堆叠的Transformer块层，将整个样本映射到其背后分布的密度和分数。该模型利用交叉注意力，使其能够评估任意点（而不仅仅是已有数据点）上的密度和分数。分数和密度之间存在数学关系：分数是密度对数的梯度。我们利用这一点，设计了一个共享主干网络，并配备两个输出头，一个用于密度，一个用于分数。

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

从数学上讲，Transformer架构非常适合这个任务。核密度估计只有一个带宽——每个点的影响范围是固定且在全局一致的。而注意力机制是它的严格推广：我们通过分析证明，单个注意力头的权重几乎相当于数据上的高斯核，因此一个交叉注意力块已经可以复现KDE的密度和分数。在此基础上，模型进一步学习多个这样的尺度并使其适应数据。DiScoFormer 并不是将经典方法当作黑箱丢弃，而是将KDE作为一个特例包含进来，并加以改进。

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

性能表现

总体而言，DiScoFormer 在密度和分数估计上均优于KDE，而且差距在KDE表现不佳的地方更加明显。在100维情况下，差距不是一点点——与手动调优的最佳KDE相比，它使分数误差降低约6.5倍，密度误差降低超过37倍，并且随着样本数量增加仍在持续改善，而KDE则因内存不足而失效。DiScoFormer 还能远超出其训练数据的范围，对于训练期间从未见过的、模态更多的混合分布，以及拉普拉斯分布和t分布等非高斯形状，它仍能保持准确。KDE的主要优势仍然是速度，尤其是在数据集较小的时候。

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

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