KL 正则化强化学习旨在导致模式坍塌

KL-Regularized Reinforcement Learning is Designed to Mode Collapse

arXiv: 2510.20817v1

论文信息

标题: KL-Regularized Reinforcement Learning is Designed to Mode Collapse

作者: Anthony GX-Chen, Jatin Prakash, Jeff Guo, et al.

发布日期: 2025-10-23

arXiv ID: 2510.20817v1

PDF 链接: 下载 PDF

3 分钟速览

  • 研究问题：在强化学习（RL）中使用 KL 正则化时，传统观点认为反向 KL 导致 “模式搜寻”、正向 KL 导致 “质量覆盖”，论文指出这一直觉在 RL 语境下并不成立，并解释了模式坍缩的根本成因。
  • 核心方法：通过数学分析最优目标分布的参数化形式，发现 KL 散度的方向仅决定目标分布族，而模式覆盖主要取决于正则化强度、奖励尺度与参考概率的相对关系；据此提出一种仅需微调奖励幅度的算法，使目标分布能够覆盖所有高质量模式。
  • 关键结果：低正则化强度且各样本可验证奖励均等（equal verifiable rewards）的常见设定下，优化目标天然为单模态；而所提算法无需外部多样性信号，即可同时提升大语言模型与化学语言模型的解质量和多样性。
  • 主要局限：算法依赖对奖励函数的适度调整，在奖励极度稀疏或噪声过大的场景中有效性可能受限；论文未系统探索极端正则化系数下的行为。
  • 适合读者：从事语言模型强化学习微调（如 RLHF）、生成模型多样性控制、以及需要在策略学习中保持探索多样性的量化交易研究者与工程师。

论文背景和研究动机

在生成模型与强化学习的交叉领域，KL 散度常被用作正则化项，约束策略与参考分布的偏离。一个流传甚广的直觉是：优化反向 KL 散度 DKL(π∥πref)D_{KL}(\pi \Vert \pi_{\text{ref}}) 会迫使策略 “锁定” 参考分布中的少数高概率模式（mode seeking），而优化正向 KL 散度 DKL(πref∥π)D_{KL}(\pi_{\text{ref}} \Vert \pi) 则倾向于覆盖参考分布的所有模式（mass covering）。这一直觉在变分推断、生成对抗网络等场景中具有一定指导意义，并被自然迁移到语言模型的强化学习微调中——例如，为了鼓励模型生成多样化的回答，许多工作主张采用正向 KL 正则化。

然而，论文指出，这种迁移忽视了一个关键区别：在标准生成建模中，优化目标通常是使模型分布逼近某个固定的数据分布；而在带奖励的 KL 正则化 RL 中，目标分布由奖励函数、参考分布以及正则化系数共同决定，是一个 “移动的靶子”。简单套用原有直觉可能导致严重的模式坍缩。尤其是在语言模型后训练中，常见的低正则化强度与奖励设计往往使最优策略坍塌到单个模式，这直接与生成多样化、高质量文本的初衷相悖。论文旨在从理论上厘清 KL 正则化 RL 中模式覆盖的决定因素，并据此设计能够自然引导多样性的算法。

核心方法和技术细节

论文首先对 KL 正则化 RL 的目标进行解析。考虑一般形式的优化问题：

max⁡πEx∼π[R(x)]−τD(π,πref)\max_\pi \mathbb{E}_{x \sim \pi}[R(x)] - \tau D(\pi, \pi_{\text{ref}})

其中 DD 可以是正向或反向 KL 散度，τ>0\tau > 0 为正则化系数。通过变分推导可得最优策略的显式形式。对于正则化项为反向 KL DKL(π∥πref)D_{KL}(\pi \Vert \pi_{\text{ref}}) 时，最优策略满足：

π∗(x)∝πref(x)exp⁡(R(x)τ)\pi^*(x) \propto \pi_{\text{ref}}(x) \exp\left(\frac{R(x)}{\tau}\right)

而若采用正向 KL DKL(πref∥π)D_{KL}(\pi_{\text{ref}} \Vert \pi)，最优策略则不能写成如此简洁的玻尔兹曼形式，但同样可由奖励、参考分布和 τ\tau 所定义的目标分布族刻画。

这一推导揭示了一个核心洞察：KL 散度的方向并不直接决定模式的取舍，而是决定了最优策略所属的分布族类型。正向 KL 与反向 KL 对应着两个不同的目标分布族，它们由 τ\tau 参数化。真正的模式覆盖能力取决于 τ\tau 的大小 以及 奖励 R(x)R(x) 与参考对数概率 log⁡πref(x)\log \pi_{\text{ref}}(x) 之间的相对尺度。

论文进一步从数学上证明，在低正则化强度（τ→0\tau \to 0）且奖励值在高质量样本上几乎均等（例如可验证任务中正确解答通常获得相同的高分）的典型设定下，无论采用正向还是反向 KL，目标分布都会退化为只在单一最高奖励模式处有峰值，即优化目标本身即为单模态。这意味着模式坍缩并非训练不稳定所致，而是目标函数在构造层面就排斥了多样性。换言之，即便是完美的优化器，也只能收敛到一个单一模式。

基于这一发现，论文提出一种简单、可扩展且有理论保证的算法。该方法仅对奖励的幅度做微小调整，例如通过平移或非线性变换，使得奖励值与参考概率的配合产生一个在多高质量模式上均有显著概率密度的目标分布。修改后的奖励函数 R′(x)R'(x) 使目标分布变为：

πtarget(x)∝πref(x)exp⁡(R′(x)τ)\pi_{\text{target}}(x) \propto \pi_{\text{ref}}(x) \exp\left(\frac{R'(x)}{\tau}\right)

算法设计的原则是：保证所有达到一定质量级别的样本在新的奖励尺度下都被赋予足够的相对优势，从而在目标分布中形成多个众数。这一过程无需任何外部的 “多样性” 标签或辅助损失，完全由优化目标的内在构造驱动。

创新点和贡献

  • 澄清长期误解：首次通过严谨的数学推导和实证，证明了在 KL 正则化 RL 中，模式覆盖的决定因素不是 KL 的方向，而是正则化强度与奖励-参考概率的相对尺度。这纠正了以往 “正向 KL 必能覆盖” 的简化认知。
  • 揭示模式坍缩的构造性根源：明确指出标准低正则化、均等奖励的设置会使目标分布天然为单模态，因此许多现有方法的失败并非算法缺陷，而是目标函数的错误指定。
  • 提出极简修正算法：仅需改变奖励的数值幅度，即可将目标分布重塑为多模态，使策略优化自然导向高质量加多样性的解。该方法不依赖任何外部多样性信号，且可与正向、反向 KL 正则化灵活结合。
  • 跨域验证：在大语言模型和化学语言模型的任务后训练中，该算法不仅提升了解决方案的质量，也显著增大了生成样本的多样性，展示了理论的普适性和实际效用。

实验结果分析

论文分别在大语言模型（LLM）和化学语言模型（CLM）上进行了后训练实验。典型任务包括数学推理、代码生成和分子优化，这些任务中均存在多个正确且高质量的解。实验对比了传统正向/反向 KL 正则化与所提议的奖励修改算法。

核心发现包括：

  • 在低正则化强度下，传统方法无论使用正向还是反向 KL，生成样本几乎全部聚集于单一解，多样性极低；而所提算法在相同设置下能稳定产生多种不同的正确答案。
  • 当人为增大正则化系数 τ\tau 时，目标分布的模式确实会增加，但此时奖励的驱动力被削弱，解的质量开始下降。所提算法则在维持高质量的同时保留了高多样性，证明单纯的 τ\tau 调节无法同时满足两者。
  • 在化学语言模型中，使用所提算法生成的分子不仅满足目标属性要求，其结构多样性也大幅高于基线，这对药物发现等下游任务极具价值。
  • 该算法的兼容性也得到验证：它与反向 KL 和正向 KL 均能协同工作，而两种基线若不配合奖励调整，则分别出现模式坍缩或质量不足的问题。

实验充分说明，多样性并非源自 KL 的 “覆盖天性”，而可通过对目标分布的结构进行有理论依据的微调来实现。

实践建议

本论文对 RL 微调语言模型、生成模型以及量化交易中的策略探索均有直接指导价值。以下建议可供工程落地参考：

  1. 重新审视 KL 正则化的作用 不要默认正向 KL 能带来多样性。在设计 RL 微调系统时，应检查奖励函数与参考分布的相对尺度。若高质量区域内的奖励值过于均一，目标分布将倾向于单模态。此时可尝试对奖励进行非线性缩放或分位数归一化，使不同但同等优秀的解获得差异化的相对优势。

  2. 调节奖励幅度而非单纯增大 τ\tau 为增加模式覆盖而增大正则化系数会牺牲奖励驱动，导致质量滑坡。更有效的方式是直接修改奖励函数，例如对奖励值施加温度系数 TT：R~(x)=R(x)/T\tilde{R}(x) = R(x)/T，或将奖励映射为排名（rank-based reward），从而在不削弱目标驱动力的前提下构造多模态目标分布。该方法极其简单，几乎不增加计算开销。

  3. 语言模型后训练中的具体实现 在 RLHF 或可验证奖励任务（如数学求解）中，若发现模型输出同质化严重，可对奖励进行如下调整：将每个样本的奖励值减去其所属 “正确类别” 的均值基准，或引入极小的高斯噪声扰动，从而打破均等奖励带来的单一峰值。根据论文思路，目标是使得目标分布 πrefexp⁡(R′/τ)\pi_{\text{ref}} \exp(R'/\tau) 在多个正确模式上具有可比概率质量。

  4. 量化交易策略生成中的应用 在利用强化学习生成交易信号或组合策略时，策略多样性对于避免过拟合、适应市场变化至关重要。可借鉴论文结论，不依赖简单的熵奖励或正向 KL，而是调整策略价值函数的刻度（例如使用夏普比率的排名而非原始收益率），使目标分布自然覆盖多个高质量的、不完全相关的策略模式。

  5. 监控与验证 在实践中，建议监控策略生成样本的聚集程度（如成对相似度、独特模式数量），并实时调整奖励缩放因子。论文算法本质上是对奖励进行一个单调变换，不会改变样本的相对排序，因此不会损害奖励的最大化逻辑，只改变概率分布的集中度，具备良好的理论保障。

总的来说，这篇论文不仅为 KL 正则化 RL 的模式坍缩现象提供了理论解释，更给出了一条极高性价比的工程实践路径：调整奖励数值，而非更换散度种类。对于任何涉及策略优化的多样性需求场景，这都是一条值得优先尝试的原则。