DiScoFormer: One transformer for density and score, across distributions

DiScoFormer：一个跨分布的密度和分数变换器

返回文章列表

[0

[-1

企业

]

文章

发布于2026年6月29日

点赞

4

[

Kyle Wiggers

Ai2Comms

关注

allenai

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

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

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

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

我们提出了一种新方案——DiScoFormer（密度和分数变换器），这是一个模型，给定一组数据点，可以在单次前向传递中同时估计分布的密度和分数，而无需重新训练。

为密度和分数估计训练变换器

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

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

Transformer 架构之所以适合这项任务，背后有数学原理支撑。核密度估计（KDE）使用单一的带宽参数——每个点的影响范围在提前固定后统一应用于所有位置。注意力机制是其严格的推广形式：我们通过数学分析证明，单个注意力头的权重几乎等同于数据上的高斯核函数，因此单个交叉注意力模块已能复现 KDE 的密度估计和得分计算。在此基础上，模型更进一步，能够同时学习多个尺度并根据数据进行自适应调整。DiScoFormer 并未用黑箱方法取代经典方法，而是将 KDE 作为特例纳入其中并加以改进。

我们使用了哪些数据来训练 DiScoFormer？我们选择高斯混合模型（GMM）作为训练数据主要基于两个原因。首先，GMM 是通用密度逼近器——当组件数量足够时，它能以任意小的误差拟合几乎所有平滑分布。其次，GMM 具有闭合形式的密度函数和得分函数，因此我们始终拥有精确的目标用于监督学习。我们通过为每个批次生成新的 GMM 实现了这两种特性，使模型获得几乎无限的目标分布示例，并针对每个示例的精确密度和得分进行监督。

性能表现

在所有测试场景中，DiScoFormer 在密度估计和得分估计上均优于 KDE，且在 KDE 表现不佳的区域差距更加显著。在 100 维数据中，DiScoFormer 与最佳手动调参的 KDE 相比，得分误差降低约 6.5 倍，密度误差降低超过 37 倍，并且随着样本数量增加持续提升性能，而 KDE 则会因内存不足而无法处理。DiScoFormer 还能有效泛化到训练数据之外的区域，保持对训练中从未见过的更多模态混合数据的准确性，以及对拉普拉斯分布和学生 t 分布等非高斯形状数据的准确建模。KDE 的主要优势仍是速度，特别是在数据集较小时。

我们认为 DiScoFormer 最具前景的特性是，分数估计是生成建模、贝叶斯推断和科学计算等多个领域共享的核心依赖。一个经过预训练的即插即用估计器，能够在高维空间保持准确性并消除针对每个问题重新训练的需求，这可能一次性降低所有相关领域的计算成本——只需一个模型，即可在所有需要分数和密度估计的场景中复用。

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

更多来自该作者的文章

混合模型在哪些 token 上预测效果更好？

8

2026 年 6 月 25 日

MolmoMotion：语言引导的三维运动预测

10

2026 年 6 月 17 日

社区

编辑

预览

通过拖拽、粘贴或点击此处上传图片、音频和视频

点击此处上传图片

评论

· 注册或登录以发表评论