[Paper] 分类重新参数化与去噪扩散模型

Source: arXiv - 2601.00781v1

概述

本文提出了一种新方法，用于训练涉及分类（即离散）变量的模型，且无需使用嘈杂的得分函数估计器或有偏的连续松弛。通过利用去噪扩散过程，作者推导出针对分类分布的闭式“软”重新参数化，可以直接进行反向传播，为广泛的机器学习流水线中的基于梯度的优化提供了实用的替代方案。

关键贡献

• 基于扩散的软重新参数化 用于分类变量，扩展了连续松弛的家族。
• 闭式去噪器 在高斯噪声过程下用于分类分布，消除了对昂贵的扩散模型训练的需求。
• 免训练扩散采样器 提供路径梯度，使得反向传播变得直接。
• 实证验证显示，在标准基准上 竞争或更优的性能，相较于经典的得分函数估计器和流行的 Gumbel‑Softmax 松弛。

方法论

1. 对 One‑Hot 向量进行高斯噪声处理 – 作者从类别变量的 one‑hot 表示出发，加入各向同性的高斯噪声，将离散点转化为连续向量。
2. 解析去噪器 – 对于这种特定的噪声模型，最优去噪器（即在噪声观测下原始 one‑hot 向量的条件期望）可以使用类似 softmax 的运算以闭式形式表达。
3. 扩散采样作为重参数化 – 通过将扩散过程反向（去噪）从高斯样本回到原始类别空间，得到一个从标准正态变量到“软”类别样本的可微映射。该映射充当 重参数化技巧：随机性被隔离在高斯种子中，其余计算是确定且可微的。
4. 梯度传播 – 由于去噪器是解析的，梯度可以在整个扩散轨迹上无缝传播，无需学习的去噪网络，从而避免额外的训练开销。

结果与发现

• 与 score‑function 估计器相比，基于 diffusion 的方法始终 降低梯度方差。
• 与温度相关的松弛方法（例如 Gumbel‑Softmax）不同，该方法 不需要调节温度调度；diffusion 时间起到类似作用，但具有原则性的解释。
• 训练时间开销极小，因为去噪器是解析的；额外成本仅为每次前向传播的少量矩阵‑向量运算。

实际意义

• Deep generative models（VAEs、normalizing flows）在需要离散潜变量的情况下，现在可以使用低方差、无偏的梯度估计器，而不牺牲模型保真度。
• Reinforcement learning 代理在离散动作空间下可以受益于更平滑的策略梯度，可能加速在探索成本高的环境中的收敛。
• Structured prediction 任务（例如 parsing、sequence labeling）传统上依赖 REINFORCE，现在可以用即插即用的 diffusion 重新参数化来替代，从而减少围绕方差降低技巧的工程工作量。
• 因为该方法是 training‑free，只需几行代码即可嵌入现有的 PyTorch/TensorFlow 流程，使其对快速原型开发和生产系统具有吸引力。

限制与未来工作

• 当前的公式假设 独立的分类变量；将扩散去噪器扩展以捕获依赖关系（例如，分类马尔可夫链）仍是一个未解决的挑战。
• 虽然去噪器对高斯噪声是解析的，但其他噪声族（例如 Laplace）可能更适合某些硬件约束，需要新的推导。
• 论文主要在中等规模基准上进行评估；扩展到 大词表（例如，拥有数万 token 的语言模型）可能会暴露需要优化实现的计算瓶颈。
• 未来的工作可以探索 自适应扩散调度，自动平衡 bias‑variance 权衡，或将该方法与学习型去噪器结合，以实现更丰富的后验近似。

作者

• Samson Gourevitch
• Alain Durmus
• Eric Moulines
• Jimmy Olsson
• Yazid Janati

论文信息

• arXiv ID: 2601.00781v1
• 分类: cs.LG, stat.ML
• 发表时间: 2026年1月2日
• PDF: 下载 PDF