跳转至

Robust Policy Optimization to Prevent Catastrophic Forgetting

会议: NeurIPS 2026
arXiv: 2602.08813
OpenReview: https://openreview.net/forum?id=fg8I6pwObp
代码: https://github.com/Helloworld10011/FRPO
领域: LLM 安全
关键词: 大模型对齐、灾难性遗忘、KL散度邻域、分布鲁棒优化、策略优化

一句话总结

针对下游微调破坏大模型已有安全护栏与特定能力的灾难性遗忘难题,提出微调鲁棒策略优化(FRPO),通过在策略空间以KL球建立最恶最坏奖励极大极小化目标,推导得出无额外计算开销的熵风险等价闭式解与Jackknife偏差校正梯度,在不损害下游适配的前提下显著抑制护栏崩溃。

研究背景与动机

大语言模型通常采用多阶段后训练管线:首先通过基于人类反馈的强化学习(RLHF)赋予模型安全护栏或特定对齐行为,随后在垂直领域开展有监督微调(SFT)或针对特定目标的策略优化。然而,即便下游微调只进行微小的参数更新,模型此前建立的安全防御机制也极易发生崩溃与退化,表现出严重的灾难性遗忘。大量实证研究表明,在看似完全良性的下游数学推理(如 GSM8K)或通用指令数据集(如 Alpaca)上进行标准微调,就会使攻击者针对基础模型的越狱与有害输出诱导成功率急剧攀升。

面对这一困境,现有防御范式主要聚焦于「下游阶段干预」:例如在微调数据集中回放部分对齐样本、采用参数高效微调(LoRA)隔离参数干扰、施加约束正则项限制权重漂移,或者在后处理阶段执行多任务模型合并。然而,这些下游约束手段高度依赖调用者的流程合规性;一旦终端用户采用标准的 SFT 或 GRPO 执行未受约束的常规微调,已有保护便不攻自破。核心矛盾在于:标准 RLHF 算法(如 PPO、GRPO)仅在当前策略单点上贪婪最大化平均奖励,导致策略参数容易收敛至奖励极度陡峭的狭窄尖锐区域(sharp regions),只要下游发生轻微策略扰动,预期奖励便会断崖式下跌。

解决该问题的切入角度是实现「预微调鲁棒性」(pre-finetuning robustness)——既然下游微调在避免严重过拟合时自然受控于距基准模型的平均 KL 邻域内,基座模型在 RLHF 阶段就应当主动避开脆弱的尖锐极值,直接在策略空间探索平坦且抗扰动的平原区域。核心 idea:将未来下游策略偏移显式建模为基准策略周边的平均 KL 有界球,构建最大化最坏情况奖励的极大极小化分布鲁棒目标,并将其对偶等价转化为惩罚低奖励采样的熵风险目标,嵌入 GRPO 架构实现无计算开销的鲁棒策略优化。

方法详解

整体框架

微调鲁棒策略优化(FRPO)围绕三阶段模型演化链 \(\pi_{\mathrm{ref}} \to \pi_\theta \to Q\) 展开:\(\pi_{\mathrm{ref}}\) 为 RLHF 前的参考模型,\(\pi_\theta\) 为待优化的基座模型,而 \(Q\) 代表未来下游微调可能生成的任意派生策略。FRPO 放弃了仅对 \(y \sim \pi_\theta\) 优化期望标量奖励的传统设定,转而在所有满足期望 KL 散度约束 \(\mathbb{E}_{x \sim p}[\mathrm{KL}(Q(\cdot|x) \parallel \pi_\theta(\cdot|x))] \le \rho\) 的邻域策略分布族中,最大化最恶劣情况下的下游奖励。

通过拉格朗日强对偶性推导,这一极大极小化问题在数学上严格等价于一个单阶段的熵风险(entropic risk)最大化目标。在实现层面,算法复用 GRPO 的分组采样机制与近端裁剪梯度,引入抵消漂移项的零阶基线确保训练收敛,并结合 Jackknife 重抽样消除小样本配分函数对数估计的代数偏差,实现完全不依赖额外 Critic 网络、不增加前向反向计算量的策略梯度更新。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入 Prompt 批次 x ~ p(x)"] --> B["分组响应采样与居中优势计算<br/>从旧策略 π_old 采样 G 条轨迹计算优势 A_i"]
    B --> C["配分函数截断估计<br/>指数加权低奖励采样并引入比率裁剪"]
    C --> D["Jackknife 留一重抽样偏差校正<br/>消除 log 变换产生的 O(1/G) 估计偏差"]
    D --> E["低方差基线补偿与策略梯度更新<br/>消除高 λ 漂移项并添加参考模型 KL 约束"]
    E --> F["输出下游微调鲁棒策略 π_θ"]

关键设计

1. 策略空间极大极小化到熵风险的对偶转化:解决下游盲微调的抗扰动建模

传统 RLHF 将奖励最大化限制在当前策略点,忽视了参数微调引入的分布偏移。FRPO 将下游策略 \(Q\) 视为在全 prompt 分布期望下距离 \(\pi_\theta\) 不超过半径 \(\rho\) 的对抗扰动分布: $\(\max_{\pi_\theta} \left\{ \inf_{Q: \mathbb{E}_x[\mathrm{KL}(Q \parallel \pi_\theta)] \le \rho} \mathbb{E}_{x \sim p, y \sim Q(\cdot|x)}[r(x,y)] - \beta \mathrm{KL}(\pi_\theta \parallel \pi_{\mathrm{ref}}) \right\}\)$ 借助 \(f\)-散度的 Fenchel 共轭性质(\(f(t)=t\log t\) 对应 \(f^*(s)=\exp(s-1)\)),内层下确界可解析求得最优似然比 \(L^*(y|x) \propto \exp(-r(x,y)/\lambda)\)。将其代入对偶问题后,最坏情况优化完全转化为最小化加权配分函数的对数期望,即熵风险形式: $\(J_\lambda(\theta) = -\mathbb{E}_{x \sim p}\left[\lambda \log \left(\mathbb{E}_{y \sim \pi_\theta(\cdot|x)} \left[e^{-r(x,y)/\lambda}\right]\right)\right] - \beta \mathrm{KL}(\pi_\theta \parallel \pi_{\mathrm{ref}})\)$ 超参数 \(\lambda > 0\) 控制着策略的风险规避敏感度。当 \(\lambda \to \infty\)(对应扰动半径 \(\rho = 0\))时,目标退化为标准 GRPO;而在有限 \(\lambda\) 下,一阶泰勒展开表明该目标显式等价于在均值奖励基础上惩罚奖励方差(\(\mathbb{E}[r] - \frac{1}{2\lambda}\mathrm{Var}(r) - \lambda\rho\)),从而迫使优化器压制低得分轨迹,优先收敛至平坦宽阔的高奖励平台。

2. 组级配分函数与单侧截断加权:在 GRPO 架构下实现无缝嵌入

为了在无 Critic 架构中直接计算上述熵风险,FRPO 借鉴 GRPO 的采样机制,对每个 prompt \(x\) 生成 \(G\) 个候选回答 \(\{y_i\}_{i=1}^G\),计算其去均值后的标量优势 \(A_i = r_i - \frac{1}{G}\sum_j r_j\)。将策略权重比与加权优势相结合,构建经验配分函数估计量: $\(\hat{Z}_\lambda(x) = \frac{1}{G}\sum_{i=1}^G \frac{1}{|y_i|}\sum_{t=1}^{|y_i|} e^{-A_{i,t}/\lambda} \left\lceil \frac{\pi_\theta(y_{i,t}|x,y_{i,<t})}{\pi_{\mathrm{old}}(y_{i,t}|x,y_{i,<t})} \right\rceil^{1+\epsilon}\)$ 此处设计的精巧之处在于利用了 \(e^{-A/\lambda} > 0\) 的非负性,将 PPO 的双向裁剪简化为单侧上限裁剪 \(\lceil u \rceil^{1+\epsilon}\)。当样本表现极差(\(A_i\) 严重为负)时,上限截断防止重要性比率发散造成梯度爆炸;同时下界 \(1-\epsilon\) 被自然省去,使得有害或劣质采样的概率能够不受阻碍地被压低至零。

3. 零阶低方差基线:消除高 \(\lambda\) 下的梯度漂移与策略崩塌

直接对 \(-\lambda \log \hat{Z}_\lambda(x)\) 求导会暴露出严峻的数值不稳定性。在理论分析中,当对数配分函数的梯度在 \(\lambda \gg 1\) 进行展开时,会产生一个由常数比率主导的 \(-\lambda\) 漂移项。由于蒙特卡洛抽样的随机涨落,该漂移项会导致所有采样的有效更新系数均呈现负值,模型熵值迅速爆炸、生成退化为乱码。FRPO 构造了不依赖于优势值的解析低方差基线: $\(J_{\text{baseline}}(\theta) = \lambda \mathbb{E}_{x \sim p}\left( \frac{1}{G}\sum_{i=1}^G \frac{1}{|y_i|}\sum_{t=1}^{|y_i|} \left[ \frac{\pi_\theta(y_{i,t}|x,y_{i,<t})}{\pi_{\mathrm{old}}(y_{i,t}|x,y_{i,<t})} \right]^{1+\epsilon} \right)\)$ 在数学期望层面,该项对梯度的贡献恒等于零,却能精确对消梯度中 \(-\lambda\) 阶的虚假负偏置,使梯度的渐近主导项严格回归至真实优势 \(A_{i,t}\),确保了算法在全参数区间内的稳定收敛,且在 \(\lambda \to \infty\) 时平滑收敛到标准 GRPO。

4. Jackknife 留一重抽样:校准小样本自归一化重要性采样偏差

在实践中,由于显存与计算预算限制,分组候选数通常较小(例如 \(G \in [4, 16]\))。经验对数配分函数本质上属于自归一化重要性采样(SNIS),非线性算子 \(\log(\cdot)\) 会在梯度中引入 \(\mathcal{O}(1/G)\) 阶的系统性统计偏差。尽管该偏差在训练初期对损失曲线影响隐蔽,但会严重削弱下游微调时的抗遗忘能力。FRPO 引入非参数 Jackknife 校正技术,计算剔除第 \(j\) 个样本后的留一配分估计量 \(\hat{Z}_{\lambda,-j}(x)\),构造出无偏化对数估计: $\(\log \tilde{Z}_\lambda(x) = G \log \hat{Z}_\lambda(x) - \frac{G-1}{G}\sum_{j=1}^G \log \hat{Z}_{\lambda,-j}(x)\)$ 该变换通过留一组合精确抵消了泰勒展开中的首阶协方差偏差,将梯度误差的代数阶降至 \(\mathcal{O}(1/G^2)\),以微秒级的张量重运算代价确保了分布鲁棒目标的保真度。

损失函数 / 训练策略

FRPO 的综合优化目标由偏差校正后的熵风险项、低方差稳定基线与相对预训练参考模型的 KL 散度惩罚共同组成: $\(J(\theta) = \mathbb{E}_{x \sim p}\left[ -\lambda \log \tilde{Z}_\lambda(x) + \lambda \cdot \text{baseline}(x) - \beta \mathrm{KL}(\pi_\theta(\cdot|x) \parallel \pi_{\mathrm{ref}}(\cdot|x)) \right]\)$ 其中参考模型散度采用无偏估计器 \(\mathrm{KL} \approx \frac{1}{G}\sum_i \frac{1}{|y_i|}\sum_t \left( \frac{\pi_{\mathrm{ref}}}{\pi_\theta} - \log \frac{\pi_{\mathrm{ref}}}{\pi_\theta} - 1 \right)\) 计算。

在训练参数配置上,安全后训练采用 LoRA 架构(\(r=64, \alpha=64\)),每批包含 1,000 条有害提示与 1,000 条无害提示。有害提示由 OpenAI 审核与 StrongREJECT 评分联合确定惩罚(\(r = 1 - (s_1+s_2)/2\)),无害提示则引入 Llama-3.1 奖励模型引导,避免过度拒答。经验标准差 \(\mathrm{std}(r) \approx 0.07\) 结合实测有效 KL 邻域 \(\rho \in [0.02, 0.3]\),通过理论公式 \(\lambda^* \approx \frac{\mathrm{std}(r)}{\sqrt{2\rho}}\) 指导确定安全最优超参数 \(\lambda = 0.2\)。

实验关键数据

主实验

评估涵盖 Mistral-v0.1-Instruct (7B) 与 Qwen2.5-7B-Instruct 两个模型底座,下游分别施加 Alpaca 通用指令微调(SFT, \(\mathrm{lr}=10^{-6}\))与 GSM8K 数学微调(SFT, \(\mathrm{lr}=6\times 10^{-6}\)),在 HarmBench 拒答率(越高越安全)与 StrongREJECT 有害性得分(越低越安全)上检验对齐保持能力。同时追踪下游主任务(GSM8K、MMLU、IFEval、HumanEval)的性能迁移,排除「通过拒绝学习下游任务来假性维持安全」的干扰。

模型与配置 下游微调方案 HarmBench 拒答率 (%) ↑ StrongREJECT 有害得分 ↓ 下游 GSM8K maj@8 下游 MMLU (5-shot)
Mistral (GRPO Baseline) Alpaca SFT ~18.5 ~0.62 55.5 54.1
Mistral (FRPO, \(\lambda=0.2\)) Alpaca SFT 64.2 0.23 56.2 54.2
Qwen2.5 (GRPO Baseline) Alpaca SFT ~42.0 ~0.45 68.4 69.8
Qwen2.5 (FRPO, \(\lambda=0.2\)) Alpaca SFT 78.6 0.18 68.9 70.1
Mistral (GRPO Baseline) GSM8K SFT ~12.0 ~0.74 55.5 54.1
Mistral (FRPO, \(\lambda=0.2\)) GSM8K SFT 48.5 0.36 56.2 54.2
Qwen2.5 (GRPO Baseline) GSM8K SFT ~38.0 ~0.49 74.2 71.0
Qwen2.5 (FRPO, \(\lambda=0.2\)) GSM8K SFT 69.4 0.25 74.6 71.2

消融实验

在持续学习任务中,以 Qwen2.5-Math-7B 为基础先进行 MATH 强化学习对齐(初始 MATH500 精度 73%),随后采用 OpenCodeInstruct(25k 代码样本)进行编程 SFT 微调。对比不同 \(\lambda\) 配置下的抗遗忘能力,并横向比较参数空间锐度感知最小化(SAM)与算法核心组件的作用。

算法 / 策略变体 消融变量与机制说明 MATH500 微调后精度 (%) ↑ 下游 MBPP+ 编程精度 (%) ↑ 安全遗忘相对改善
FRPO (\(\lambda=2.0\), 理论最优) 全量机制(基线 + Jackknife + \(\lambda=2.0\)) 64.5 ± 2 62.4 基准最优
FRPO (\(\lambda=10.0\)) 偏向大 \(\lambda\)(方差惩罚偏弱) 44.5 ± 3 62.2 遗忘显著加剧
FRPO (\(\lambda=0.2\)) 采用安全场景超参(过小,过度保守) 59.3 ± 4 61.7 次优
GRPO Baseline (\(\lambda \to \infty\)) 无鲁棒加权的标准强化学习 42.3 ± 2 62.5 遗忘幅度高达 -30.7%
GRPO + SAM (\(\rho=0.05\)) 参数空间 \(L_2\) 扰动平坦优化 51.2 ± 3 62.0 优于GRPO但弱于FRPO
FRPO w/o Baseline 移除低方差补偿项式 (4.2) 训练崩溃(熵爆炸) - 无法收敛
FRPO w/o Jackknife 移除留一重抽样偏差校正 57.8 ± 3 61.9 下游抗遗忘性削弱

关键发现

  • 对偶参数 \(\lambda\) 的跨任务预测性:最佳超参数完全由底层奖励方差与微调允许的有效 KL 球半径决定(\(\lambda^* \approx \mathrm{std}(r)/\sqrt{2\rho}\))。安全对齐中奖励标准差极小(\(\sim 0.07\)),最优 \(\lambda=0.2\);数学强化学习任务中验证奖励标准差放大 10 倍(\(\sim 0.5\)),最优 \(\lambda=2.0\) 完美契合实测峰值。
  • 策略空间与参数空间平坦化的本质差异:GRPO+SAM 通过参数空间 \(L_2\) 球扰动平坦化虽能略微缓解遗忘,但在微调后期显著落后于 FRPO。原因在于参数 \(L_2\) 球混杂了大量任务无关或冗余维度的扰动,而策略空间的 KL 球直接约束语义概率分布偏移,贴合下游微调的真实动力学。
  • 下游可塑性与抗遗忘的双赢:不同于对抗性鲁棒对齐(如 TAR、derta)通过人为增加微调损失阻碍模型适配下游任务,FRPO 在 GSM8K、MMLU、IFEval 与 HumanEval 上的微调收敛速率与最终得分均与 GRPO 保持严格一致,表明其抗遗忘性源于在保持下游学习能力的同时锁定了低风险流形。

亮点与洞察

  • 凸对偶闭式解彻底卸去鲁棒对抗训练的额外计算开销:传统分布鲁棒或对抗训练往往依赖代价高昂的双层内外循环交替迭代,FRPO 通过凸共轭推导出单阶段熵风险闭式解,使得鲁棒性收益可以以零 FLOPs 增量的形式直接在基础 GRPO 梯度管线中兑现。
  • 极小化极大解构重塑了对强化学习样本利用的认知:传统强化学习过度追逐高正向优势的极端样本(容易落入过拟合尖峰),FRPO 的对偶机制表明,对抗未知微调的根本抓手在于对组内低奖励劣质响应施加指数级的负反馈重赋权,以全局方差收缩换取下游扰动下的性能下限保障。
  • 可迁移至多阶段后训练泛化体系:该方法中建立的「有界 KL 球内极小化极大稳定化」设计,不仅适用于安全对齐与数学推理,还能无缝泛化至大模型指令微调、跨任务持续学习以及代理系统的长期策略固化。

局限与展望

  • 假设下游偏移局限于适度 KL 球半径:理论推导立足于下游微调不发生灾难性过拟合的假设(per-token KL \(\le 0.3\))。当用户采用极端高学习率且无正则的持续暴力训练时,策略将不可避免地脱离初始 KL 保护球。
  • 缺乏细粒度令牌级或动态多目标方差自适应机制:当前算法将超参数 \(\lambda\) 作为离散常量进行标定。针对奖励信号极度稀疏、多目标(如 Helpful-Harmless)剧烈冲突的任务,全局静态的 \(\lambda\) 可能无法在单条序列的不同决策步实现精细化权衡。
  • 未来方向:探索将预微调鲁棒思想前移至大规模自监督预训练与指令微调阶段,并研究特定认知能力(如概念知识 vs 过程推理)在策略平坦化表征上的本质难度差异。

相关工作与启发

  • vs GRPO / PPO (Ouyang et al., 2022; Shao et al., 2024):GRPO 仅优化经验期望奖励,收敛于尖锐解;FRPO 在其基础上通过对数配分函数重权化,引入风险规避机制,并在 \(\lambda \to \infty\) 时无缝退化为 GRPO。
  • vs SAM (Foret et al., 2020; Watts et al., 2026):SAM 在权重参数空间优化 \(L_2\) 球最坏损失,受参数参数化形式与尺度缩放影响严重;FRPO 在概率输出分布空间基于 KL 散度寻找平坦解,几何意义更为本质且能精确捕捉文本序列的微调扰动。
  • vs 对抗训练与遗忘消除 (Tamirisa et al., 2024; Zou et al., 2024):TAR 采用迭代式对抗微调,内部循环训练耗时极其昂贵且损伤下游任务拟合能力;Circuit Breaking / Representation Rerouting 依赖表征层硬性几何重定向,通用性受限。FRPO 提供通用标量奖励框架,无内循环计算且兼顾下游任务适应性。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从策略空间分布鲁棒优化角度形式化解决微调灾难性遗忘,严格导出解析熵风险解。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖安全(Mistral/Qwen)、数学到代码持续学习双重场景,包含 48 组超参微调扰动分析及消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰,闭式解与 GRPO 衔接自然,问题动机与实证分析逻辑严密。
  • 价值: ⭐⭐⭐⭐⭐ 完美兼容工业界流行的大模型后训练流水线,零计算代价大幅提升开源模型防恶意微调的内生抗性。