Paper deep dive
Robust Policy Optimization to Prevent Catastrophic Forgetting
Mahdi Sabbaghi, George Pappas, Adel Javanmard, Hamed Hassani
Models: Mistral-v0.1-Instruct, Qwen2.5-7B-Instruct, Qwen2.5-Math-7B
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 94%
Last extracted: 3/12/2026, 5:16:11 PM
Summary
The paper introduces Fine-tuning Robust Policy Optimization (FRPO), a framework designed to prevent catastrophic forgetting in large language models during downstream fine-tuning. By optimizing for reward flatness within a KL-bounded neighborhood of policies, FRPO ensures stability under policy shifts. The method modifies the GRPO algorithm to incorporate a max-min formulation, resulting in an entropic-risk objective that requires no additional computation while significantly improving safety and task performance retention.
Entities (5)
Relation Signals (3)
FRPO → modifies → GRPO
confidence 95% · By modifying GRPO, we develop an algorithm with no extra computation
FRPO → utilizes → KL Divergence
confidence 95% · optimizes reward not only at the current policy, but across a KL-bounded neighborhood of policies
FRPO → prevents → Catastrophic forgetting
confidence 90% · Robust Policy Optimization to Prevent Catastrophic Forgetting
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Large language models are commonly trained through multi-stage post-training: first via RLHF, then fine-tuned for other downstream objectives. Yet even small downstream updates can compromise earlier learned behaviors (e.g., safety), exposing a brittleness known as catastrophic forgetting. This suggests standard RLHF objectives do not guarantee robustness to future adaptation. To address it, most prior work designs downstream-time methods to preserve previously learned behaviors. We argue that preventing this requires pre-finetuning robustness: the base policy should avoid brittle high-reward solutions whose reward drops sharply under standard fine-tuning. We propose Fine-tuning Robust Policy Optimization (FRPO), a robust RLHF framework that optimizes reward not only at the current policy, but across a KL-bounded neighborhood of policies reachable by downstream adaptation. The key idea is to ensure reward stability under policy shifts via a max-min formulation. By modifying GRPO, we develop an algorithm with no extra computation, and empirically show it substantially reduces safety degradation across multiple base models and downstream fine-tuning regimes (SFT and RL) while preserving downstream task performance. We further study a math-focused RL setting, demonstrating that FRPO preserves accuracy under subsequent fine-tuning.
Tags
Links
Trouble viewing inline? Open PDF directly →
Full Text
93,760 characters extracted from source content.
Expand or collapse full text
Robust Policy Optimization to Prevent Catastrophic Forgetting Mahdi Sabbaghi Correspondence can be made to: smahdi@seas.upenn.edu University of Pennsylvania George Pappas University of Pennsylvania Adel Javanmard University of Southern California Hamed Hassani University of Pennsylvania Abstract Large language models are commonly trained through multi-stage post-training: first via RLHF, then fine-tuned for other downstream objectives. Yet even small downstream updates can compromise earlier learned behaviors (e.g., safety), exposing a brittleness known as catastrophic forgetting. This suggests standard RLHF objectives do not guarantee robustness to future adaptation. To address it, most prior work designs downstream-time methods to preserve previously learned behaviors. We argue that preventing this requires pre-finetuning robustness: the base policy should avoid brittle high-reward solutions whose reward drops sharply under standard fine-tuning. We propose Fine-tuning Robust Policy Optimization (FRPO), a robust RLHF framework that optimizes reward not only at the current policy, but across a KL-bounded neighborhood of policies reachable by downstream adaptation. The key idea is to ensure reward stability under policy shifts via a max-min formulation. By modifying GRPO, we develop an algorithm with no extra computation, and empirically show it substantially reduces safety degradation across multiple base models and downstream fine-tuning regimes (SFT and RL) while preserving downstream task performance. We further study a math-focused RL setting, demonstrating that FRPO preserves accuracy under subsequent fine-tuning. https://github.com/Helloworld10011/FRPO 1 Introduction Large language models (LLMs) are becoming the driving force of agentic systems, ranging from everyday chat and mathematical reasoning to robotics (Ahn et al., 2022; Schick et al., 2023; Driess et al., 2023; Achiam et al., 2023). Such broad deployment requires training a single base model capable of supporting diverse tasks, and, in some cases, fine-tuning for specific downstream applications. However, optimizing for a downstream objective can compromise other capabilities developed during earlier training (de Masson D’Autume et al., 2019; Sun et al., 2019). Addressing this trade-off is the core challenge of continual learning: training models on new tasks without losing prior knowledge (De Lange et al., 2021; Wang et al., 2022b, a). We study a two-stage continual learning setting in LLM pipelines: (1) the model acquires a behavior such as safety guardrails or a specific capability; and then (2) it is adapted to a downstream task while maintaining the earlier behavior. Recent studies (Qi et al., 2023; Zhan et al., 2023; Qi et al., 2024) show that this often fails: downstream fine-tuning can degrade previously learned behaviors, a phenomenon known as catastrophic forgetting. Figure˜1(right) illustrates this failure mode: fine-tuning on a math task such as GSM8K inadvertently leads to “forgetting” the safety guardrails. To mitigate forgetting, many studies focus on developing methods at downstream time to preserve the model’s capabilities. Rehearsal methods augment downstream training with a subset of previous data (Rolnick et al., 2019; Sun et al., 2019; Scialom et al., 2022; Huang et al., 2024). Parameter-efficient fine-tuning (PEFT) constrains downstream updates into separate modules (e.g., LoRA), thereby reducing parameter interference between the objectives (Hu et al., 2022; Hsu et al., 2024; Qiao and Mahdavi, 2024). Regularization strategies constrain optimization through penalties that prevent the policy from drifting excessively (Li and Hoiem, 2017; Schulman et al., 2017; Lee et al., 2019; Kirkpatrick et al., 2017). Finally, model merging methods aim to combine task-specific models post-hoc to retain all capabilities without expensive retraining (Ilharco et al., 2022; Wortsman et al., 2022; Yi et al., 2024; Djuhera et al., 2025). Figure 1: Illustration of FRPO. Standard RLHF finds high-reward policies that may lie in sharp regions, whereas our method optimizes for reward-flatness within a KL neighborhood, finding policies that maintain high reward after downstream adaptation. Downstream-time methods are effective only when used within their intended fine-tuning procedures. Consequently, any later “ordinary” fine-tuning that deviates from these procedures can still compromise previously learned behaviors. Much of the existing literature continues to rely on standard approaches—supervised fine-tuning (SFT) or RL-based methods such as GRPO (Ouyang et al., 2022; Shao et al., 2024)—which focus on optimizing the immediate downstream objective. However, the central issue remains: “The base model must be made robust to future fine-tuning, regardless of the downstream task or algorithm.” To this end, we propose Fine-tuning Robust Policy Optimization (FRPO), a new algorithm aimed at making the policy robust to downstream fine-tuning. Since downstream adaptations typically remain within a neighborhood of the current policy (e.g., bounded in KL divergence), achieving robustness requires avoiding high-reward solutions that reside in sharp regions of the policy space, where even small updates can lead to substantial reward degradation. Rather than maximizing reward solely at the current policy, our objective explicitly considers a neighborhood around the policy and seeks flatter regions, leading to policy stability. Formally, to account for future downstream adaptations, we consider the set of all policies reachable within a KL-bounded neighborhood around the current policy as illustrated in Figure˜1. This formulation is agnostic to the downstream fine-tuning, assuming only that the policy remains within this KL neighborhood—reflecting standard practices like KL regularization in RLHF (Ouyang et al., 2022) and implicit constraints from limited learning rates or LoRA. We then propose an objective that maximizes the minimum reward over this set. By deriving the dual form, we show that this robustness criterion is equivalent to penalizing low-reward rollouts under the current policy. This leads to an entropic-risk objective that emphasizes low-reward trajectories and discourages policies with high reward variance. Finally, we derive FRPO to optimize this objective within the standard GRPO framework, with no additional computation. Our contributions are summarized as follows: • RLHF framework for fine-tuning robustness. We consider a max-min optimization that accounts for all downstream shifts within a KL ball centered at the base policy. We show that this formulation is equivalent to optimizing an entropic-risk objective, with a tunable parameter λ controlling sensitivity to low-reward trajectories. • FRPO with no extra computation. We present FRPO, a policy gradient method that integrates seamlessly with the GRPO framework. FRPO derives from a closed-form solution to the max-min problem, includes a baseline for stable optimization, and recovers GRPO as λ→∞λ →∞. • Experimental results. We evaluate several fine-tuning schemes, including instruction-following and math (see Figures˜3 and 4) to demonstrate that FRPO-trained models are significantly more effective at preserving safety guardrails and prior capabilities compared to other methods. Lastly, we fine-tune math-trained models on code generation to show that FRPO maintains 22% higher accuracy on MATH than GRPO. 1.1 Related Work RLHF. RLHF optimizes a KL-regularized objective to align LLMs (Christiano et al., 2017; Ouyang et al., 2022; Bai et al., 2022; Stiennon et al., 2020). PPO (Schulman et al., 2017) is the standard choice but requires a learned critic, whereas GRPO and RLOO (Shao et al., 2024; Ahmadian et al., 2024) remove this requirement via group-based advantages. DPO (Rafailov et al., 2023) bypasses reward modeling entirely. We modify GRPO to incorporate robustness while preserving its computational simplicity. Distributionally robust optimization and robust RL. Distributionally robust optimization (DRO) optimizes worst-case performance over uncertainty sets (Shapiro and Kleywegt, 2002; Kuhn et al., 2019; Duchi and Namkoong, 2021; Shapiro, 2017). Applications include group distributionally robustness and domain adaptation (Sagawa et al., 2019; Oren et al., 2019; Sinha et al., 2017). In RL, robust MDPs consider adversarial dynamics (Vinitsky et al., 2020; Pinto et al., 2017), while risk-sensitive RL optimizes CVaR (Chow et al., 2015; Tamar et al., 2015) or entropic risk (Osogami, 2012; Fei et al., 2021b, a). Unlike robust MDPs (Smirnova et al., 2019; Eysenbach and Levine, 2021; Zhang et al., 2020; Derman and Mannor, 2020) that perturb environment dynamics and demand robust decisions, our method perturbs the policy itself in an RLHF setting to stay robust to downstream fine-tuning. We apply DRO to the policy space with KL constraints, yielding an entropic risk objective over trajectories. Flat minima and sharpness-aware optimization. Flat minima in the loss landscape correlate with better generalization (Keskar et al., 2016; Jiang et al., 2019; Cha et al., 2021). This motivates sharpness-aware minimization (SAM) (Foret et al., 2020) and variants (Kwon et al., 2021; Zhuang et al., 2022; Kim et al., 2022). SAM considers a min-max loss in the parameter space, seeking flat regions. Our method is analogous but operates in policy space on the reward landscape: we seek “reward flatness” over a KL neighborhood. This yields policies stable under downstream perturbations—a different notion from parameter-space flatness. Adversarial training for alignment. Several methods train models to preserve alignment after adversarial fine-tuning. TAR (Tamirisa et al., 2024) performs iterative adversarial fine-tuning to identify vulnerable configurations; Circuit Breaking (Zou et al., 2024) projects unsafe inputs to incoherent outputs, and Representation Noising (Rosati et al., 2024) drives harmful representations toward random noise. These methods involve expensive inner-loop optimization (TAR) or modify internal representations (Representation Rerouting) that is primarily applicable to safety rather than a general objective. In contrast, we derive a closed-form solution for optimizing a general-purpose reward and demonstrate broader effectiveness on mathematical reasoning and continual learning as well. Additional related work is provided in Appendix˜A. 2 Preliminaries: RLHF and Fine-tuning We study the scenario where a base model is fine-tuned on a downstream task. Let x∼p(x)x p(x) denote a prompt and yi=(yi,1,…,yi,|yi|)∼π(⋅∣x)y_i=(y_i,1,…,y_i,|y_i|) π(· x) a full generated sample, where yi,t∼π(⋅∣[x,yi,<t])y_i,t π(· [x,y_i,<t]). We denote the reference policy with πref _ref as the model before the RLHF stage. At training time, we aim to optimize the policy πθ _θ as the base policy we want to be robust to downstream fine-tuning. At the downstream stage, we wish to optimize a policy Q initialized from πθ _θ. Therefore, the chain of models is: πref→πθ→Q _ref→ _θ→ Q. RLHF objective and policy optimization. Following standard RLHF and given a reward signal, we define the trajectory return r(x,y)=∑t=1|yi|rt(x,y≤t)r(x,y)= _t=1^|y_i|r_t(x,y_≤ t) (often implemented as an outcome reward, i.e., one reward for the entire response). The RLHF objective then maximizes (Ouyang et al., 2022): maxθx∼py∼πθ(⋅|x)[r(x,y)]−βx∼p[KL(πθ(⋅|x)∥πref(⋅|x))] _θ\;\; E_x p E_y _θ(·|x)\! [r(x,y) ]\;-\;β\, E_x p\! [ KL( _θ(·|x)\,\|\, _ref(·|x)) ] In practice, this objective is optimized with PPO-style policy gradients using on-policy samples from a lagged policy πold _old (Schulman et al., 2017); we adopt GRPO (Shao et al., 2024) as an efficient algorithm for optimizing this objective, which we summarize next. GRPO for policy optimization. For each prompt x, GRPO samples a group of responses yi|i=1G∼πold(⋅∣x)y_i |_i=1^G _old(· x). It then updates the trainable policy πθ _θ using a PPO-style clipped importance ratio. Without a learned critic, advantages are obtained from within-group reward statistics, and a KL penalty to a fixed reference policy is added directly to the loss: J(θ)= J(θ)= x∼p[(1G∑i=1G1|yi|∑t=1|yi|minπθ(yi,t|x,yi,<t)πold(yi,t|x,yi,<t)Ai,t,[πθ(yi,t|x,yi,<t)πold(yi,t|x,yi,<t)]1−ϵ1+ϵAi,t)] E_x p\! [ ( 1G _i=1^G 1|y_i| _t=1^|y_i| \ _θ(y_i,t|x,y_i,<t) _old(y_i,t|x,y_i,<t)A_i,t, [ _θ(y_i,t|x,y_i,<t) _old(y_i,t|x,y_i,<t) ]_1-ε^1+εA_i,t \ ) ] −βKL(πθ∥πref) 224.2479pt-β KL( _θ\| _ref) (2.1) where Ai,t≡Ai=ri−1G∑j=1GrjA_i,t≡ A_i=r_i- 1G _j=1^Gr_j as we only work with an outcome-based reward. Unlike the approach of (Shao et al., 2024), we do not normalize the advantages by the standard deviation to avoid the difficulty bias reported in (Liu et al., 2025). The KL divergence is estimated with the following approximately unbiased estimator (Ouyang et al., 2022): KL(πθ(⋅∣x)∥πref(⋅∣x))≈1G∑i=1G1|yi|∑t=1|yi|(πref(yi,t|x,yi,<t)πθ(yi,t|x,yi,<t)−logπref(yi,t|x,yi,<t)πθ(yi,t|x,yi,<t)−1)KL( _θ(· x)\,\|\, _ref(· x))≈ 1G _i=1^G 1|y_i| _t=1^|y_i| ( _ref(y_i,t|x,y_i,<t) _θ(y_i,t|x,y_i,<t)- _ref(y_i,t|x,y_i,<t) _θ(y_i,t|x,y_i,<t)-1 ) (2.2) KL-bounded fine-tuning. In standard downstream adaptation, the fine-tuned policy Q does not move arbitrarily far from its initialization πθ _θ. In RL-based methods, this is enforced explicitly via trust-region updates (Schulman et al., 2015) or the KL regularizer (the RLHF penalty to a reference policy) (Schulman et al., 2017; Ouyang et al., 2022). In SFT, this is controlled implicitly through conservative optimization choices like small learning rates and early stopping (Mosbach et al., 2020; Dodge et al., 2020), or parameter-efficient updates such as LoRA (Hu et al., 2022). This is validated in Figure˜2, where the KL grows controllably during SFT. In contrast, aggressive fine-tuning with large learning rates or extensive epochs often leads to overfitting; this is illustrated in Section˜F.1, where fine-tuning on Alpaca (Taori et al., 2023) with a high learning rate causes a significant harm to the general capabilities. Motivated by this, we model downstream adaptation as a policy Q that remains within a KL neighborhood of πθ _θ. 3 Problem Formulation: Robust RLHF Our primary objective is to train a policy πθ _θ that preserves its expected reward across all fine-tuned variants Q lying within a specified distance of πθ _θ. More concretely, let p(x)p(x) be a distribution over contexts, let πθ(y∣x) _θ(y x) denote the trainable base policy, let πref(y∣x) _ref(y x) be a fixed reference policy (e.g., the pre-RLHF model), and let r(x,y)r(x,y) be a scalar reward. We consider any arbitrary Q(y|x)Q(y|x) that is an adaptation of πθ(⋅|x) _θ(·|x) in downstream fine-tuning, but that must stay within an average KL KL ball of radius ρ>0ρ>0: maxπθ _ _θ infQx∼py∼Q(⋅∣x)[r(x,y)]⏟robust reward−βKL(πθ∥πref) \; _Q\; E_x p E_y Q(· x)[\,r(x,y)\,]_robust reward-β\, KL\! ( _θ\| _ref ) (3.1) s.t. x∼p[KL(Q(⋅∣x)∥πθ(⋅∣x))]≤ρ,∀x:∫Q(dy∣x)=1. \ \ E_x p [ KL\! (Q(· x)\, \|\, _θ(· x) ) ]\ ≤\ ρ, ∀ x: Q(dy x)=1. Crucially, the constraint is imposed in expectation over the prompt distribution p(x)p(x), rather than point-wise for each x. This allows the downstream policy to change significantly on task-relevant prompts (where fine-tuning occurs), but requires it to stay close on average. We use the following lemma from (Shapiro, 2017; Duchi and Namkoong, 2021) modified to handle the additional expectation over x. The proof can be found in Appendix˜C. Lemma 3.1 (Inner optimization under a general f–divergence). Let f:ℝ→ℝ+∪+∞f:R\! _+∪\+∞\ be a convex function with f(1)=0f(1)=0. Define the likelihood ratio L(y∣x):=Q(y∣x)π(y∣x)L(y x):= Q(y x)π(y x). Then, the inner problem in Equation˜3.1 under the average constraint: x∼py∼π(⋅∣x)[f(L(y∣x))]≤ρ,y∼π(⋅∣x)[L(y∣x)]=1∀x E_x p E_y π(· x)[\,f(L(y x))\,]\ ≤\ ρ, E_y π(· x)[\,L(y x)\,]=1\ \ ∀ x admits the following dual form: infQx∼py∼Q(⋅∣x)[r(x,y)]=supλ≥0,η:→ℝ _Q\ E_x p E_y Q(· x)[\,r(x,y)\,]= _λ≥ 0,\ η:X \ −λρ−x∼pη(x) -λρ- E_x pη(x) −λx∼py∼π(⋅∣x)[f∗(−r(x,y)−η(x)λ)], -λ\, E_x p E_y π(· x)\! [f^*\! ( -\,r(x,y)-η(x)λ ) ] \, (3.2) where f∗(s):=supt≥0st−f(t)f^*(s):= _t≥ 0\st-f(t)\ is the Fenchel conjugate. In addition, if the supremum is finite, it is attained at some (λ∗,η∗)(λ^*,η^*). In our setting, we take f(x)=xlog(x)f(x)=x (x) that leads to the forward KL divergence between Q and πθ _θ. Then, it is easy to see that the conjugate is f∗(y)=exp(y−1)f^*(y)= (y-1), and the corresponding optimal likelihood ratio is L(y|x)=exp(−r(x,y)/λ)y∼πθ(⋅∣x)[e−r(x,y)/λ]L(y|x)= (-r(x,y)/λ) E_y _θ(· x)\! [e^-\,r(x,y)/λ ]. The optimization problem thus becomes: supλ,ηx∼py∼πθ(⋅∣x)[−λexp(−1λ(r(x,y)+η(x)+λ))−η(x)]−λρ _λ,\ η \ E_x p E_y _θ(· x) [-λ\, (- 1λ(r(x,y)+η(x)+λ) )-η(x) ]-λρ \ For each x, Equation˜3.2 can be optimized w.r.t. η(x)η(x) separately. After plugging in, we obtain: minLmaxλ≥0,ηℒ=maxλ≥0maxη(x)minLℒ(L,λ,η)= _L _λ≥ 0,ηL= _λ≥ 0 _η(x) _LL(L,λ,η)= maxλ≥0−x∼p[λlogZ(x)]−λρ, _λ≥ 0 \-\, E_x p\! [λ Z(x) ]-λρ \, (3.3) Z(x):=y∼πθ(⋅∣x)[e−r(x,y)/λ] Z(x):= E_y _θ(· x)\! [e^-\,r(x,y)/λ ] Combining with Equation˜3.1, the saddle problem reduces to the following max–max program where we can swap the order of the two maxes to obtain the following: maxλ≥0maxπθ−x∼p[λlog(y∼πθ(⋅∣x)e−r(x,y)/λ)]−λρ−βKL(πθ∥πref) \ _λ≥ 0 _ _θ \- E_x p\! [λ ( E_y _θ(· x)e^-\,r(x,y)/λ ) ]-λρ-β\, KL\! ( _θ\| _ref ) \ (3.4) Eventually, we view λ as a tunable hyperparameter in the policy optimization. Therefore, the final optimization is: maxπθ−x∼p[λlog(y∼πθ(⋅∣x)e−r(x,y)/λ)]−βKL(πθ∥πref)⏟Jλ(θ) _ _θ \ - E_x p\! [λ ( E_y _θ(· x)e^-\,r(x,y)/λ ) ]-β\, KL\! ( _θ\| _ref )_J_λ(θ) \ (3.5) Remark 3.2. If ρ=0ρ=0 (unperturbed objective), then the optimal λ is λ∗=∞λ^*=∞, which recovers GRPO and the term [r(x,y)] E[r(x,y)] in the objective as λ→∞λ→∞. To see this, we start from ρ=0ρ=0 and write (3.3) as −minλ≥0[x∼p[λlogZ(x)]+λρ]- _λ≥ 0[ E_x p\! [λ Z(x) ]+λρ]. Now, by the Jensen’s inequality and the concavity of the log function, we have logZ(x)≥y∼πθ(⋅|x)[−r(x,y)/λ] Z(x)≥ E_y _θ(·|x)[-r(x,y)/λ], and also x∼p[λlogZ(x)]≥[−r(x,y)] E_x p\! [λ Z(x) ]≥ E[-r(x,y)]. Furthermore, by taking λ→∞λ→∞ we can achieve this minimum. Therefore, the optimal value is λ∗=∞λ^*=∞. Remark 3.3. The max-min formulation yields the entropic risk −λlogπθe−r/λ-λ E_ _θe^-r/λ, where smaller λ yields a more risk-averse objective. For large λ, by a Taylor expansion in 1/λ1/λ and keeping O(1/λ)O(1/λ), the first two terms in (3.4) become: maxλ≥0maxπθ[r(x,y)]−12λVar(r(x,y))−λρ=maxπθ[r(x,y)]−2ρVar(r(x,y)) _λ≥ 0 _ _θ \ E[r(x,y)]- 12λ Var(r(x,y))-λρ \= _ _θ \ E[r(x,y)]- 2ρ Var(r(x,y)) \ This shows a trade-off between mean reward and variance in the objective, favoring more consistent policies. Then, this reward consistency in the dual problem transfers to policy robustness in the primal problem. 4 FRPO Algorithm Algorithm 1 Fine-tuning Robust Policy Optimization (FRPO) 1:Initialize policy parameters θ. 2:for iteration k=1,2,…k=1,2,… do 3: πold←πθ _old← _θ ⊳ lagged sampling policy; updated every iteration (or every K iterations) 4: for all x in minibatch do 5: Sample yii=1G∼πold(⋅∣x)\y_i\_i=1^G _old(· x) 6: ri←r(x,yi)r_i← r(x,y_i); Ai←ri−1G∑j=1GrjA_i← r_i- 1G _j=1^Gr_j ⊳ centered group advantages 7: Compute Z^λ(x) Z_λ(x) using clipped ratios πθπold _θ _old (Eq. 4.1) ⊳ partition function 8: Compute leave-one-out Z^λ,−j(x) Z_λ,-j(x) for j=1,…,Gj=1,…,G 9: logZ~λ(x)←GlogZ^λ(x)−G−1G∑j=1GlogZ^λ,−j(x) Z_λ(x)← G Z_λ(x)- G-1G _j=1^G Z_λ,-j(x) ⊳ jackknife: reduces O(1/G)O(1/G) bias 10: J(θ)←x[−λlogZ~λ(x)−βKL(πθ∥πref)]+baseline (Eq. 4.2)J(θ) _x [-λ Z_λ(x)-β KL( _θ\| _ref) ]+baseline (Eq.~ eq:low_var) ⊳ add the baseline and the KL 11: Update θ←θ+η∇θJ(θ)θ←θ+η _θJ(θ). We derived the robust objective Jλ(θ)J_λ(θ) in Equation˜3.5. Following the discussion in Section˜2, we focus on the outcome-based setting where ri,t=ri=r(x,yi)r_i,t=r_i=r(x,y_i), and use the centered advantages: Ai=ri−1G∑j=1GrjA_i=r_i- 1G _j=1^Gr_j. We now explain how to optimize this objective using a GRPO-style policy gradient. First, we rewrite the inner expectation under πold _old using importance sampling, as samples are drawn from the lagged policy. The objective becomes: Jλ(θ)=x∼p[−λlog(y∼πold(⋅∣x)πθ(y|x)πold(y|x)e−r(x,y)/λ)]−βKL(πθ∥πref) J_λ(θ)= E_x p\! [-λ ( E_y _old(· x) _θ(y|x) _old(y|x)e^-\,r(x,y)/λ ) ]-β\, KL\! ( _θ\| _ref ) We replace rewards with advantages in the equation above; we show in Appendix˜D that subtracting the group average does not affect the gradient. Using token-wise ratios and clipping in Section˜2, the Monte Carlo estimate of the objective becomes: Jλ(θ)=x∼p[−λlog(Z^λ(x))−KL(πθ∥πref)], J_λ(θ)= E_x p\! [-λ ( Z_λ(x) )- KL( _θ\,\|\, _ref) ], (4.1) Z^λ(x):=1G∑i=1G1|yi|∑t=1|yi|e−Ai,t/λ⌈πθ(yi,t|x,yi,<t)πold(yi,t|x,yi,<t)⌉1+ϵ Z_λ(x):= 1G _i=1^G 1|y_i| _t=1^|y_i|e^-A_i,t/λ _θ(y_i,t|x,y_i,<t) _old(y_i,t|x,y_i,<t) ^1+ε Here we use that eA/λ≥0e^A/λ≥ 0, and minu,u|1−ϵ1+ϵ=u|1+ϵ \u,u |_1-ε^1+ε\=u |^1+ε for u≥0u≥ 0, and KL is plugged in from Equation˜2.2. Baseline. In Appendix˜D, we show that when λ is large, the gradient of the log-partition function estimate log(Z) (Z) has a high variance, leading to convergence issues as discussed in several prior works (Chung et al., 2021; Mei et al., 2022). We address this by adding a baseline to the objective that completely cancels the leading drift term in the gradient for large λ, but does not change the expected gradient. This extra term makes our algorithm converge to GRPO as λ→∞λ→∞ (see Remark˜3.2). Jlow-variance(θ)=J(θ)+λx∼p(1G∑i=1G1|yi|∑t=1|yi|[πθ(yi,t|x,yi,<t)πold(yi,t|x,yi,<t)]1+ϵ)⏟baseline splitJ_low-variance(θ)=J(θ)+ λ\;E_x p ( 1G _i=1^G 1|y_i| _t=1^|y_i| [ _θ(y_i,t|x,y_i,<t) _old(y_i,t|x,y_i,<t) ]^1+ε )_baseline split (4.2) The term inside the sum is independent of the advantages and only adds a constant to the objective in expectation. Bias reduction. In Appendix˜E, we show that the Monte Carlo approximation of logZ(x) Z(x) introduces a bias of order O(1/G)O(1/G) in ∇θJλ(θ) _θJ_λ(θ). This bias is problematic when the group size G is small (the standard group size is typically between 4 and 16). Then, we show that the jackknife technique reduces this bias to O(1/G2)O(1/G^2). Skipping the KL regularizer and the baseline term for readability, we write the corrected objective as: Jcorrected(θ)=x∼p[ J_corrected(θ)= E_x p\! [ −λGlog(Z^λ(x))+λG−1G∑j=1Glog(Z^λ,−j(x))] -λ G\, \! ( Z_λ(x) )+λ G-1G _j=1^G \! ( Z_λ,-j(x) ) ] (4.3) where Z^λ,−j(x) Z_λ,-j(x) is the leave-one-out estimate with index j removed. The full algorithm with the baseline and the jackknife correction is summarized in Algorithm˜1. 5 Experiments Our experimental investigation focuses on two scenarios: (i) Safety training with RLHF, where we aim to mitigate the catastrophic forgetting of safety guardrails; and (i) Math training with RL, which represents a broader non-safety application where we study the preservation of mathematical accuracy after fine-tuning on code generation. 5.1 Robustness in Safety Training We train two model families: Mistral-v0.1-Instruct (Jiang et al., 2023) and Qwen2.5-7B-Instruct (Qwen et al., 2025). We include Mistral because its lack of built-in guardrails makes the contrast between our method and GRPO clearer. Additionally, Qwen is a more capable model and a standard baseline for GRPO in the literature and does not have strong guardrails. Figure 2: (left/middle) The safety reward for Mistral and Qwen as the policy moves from the base by increasing the KL during fine-tuning, when sweeping λ, and evaluated on a split of the safety prompts; λ=0.2λ=0.2 better preserves the safety reward for both models and yields the most flat landscape. (right) KL is the average sequence-level on safety prompts which increases under a constant-lr schedule. Training Results. Our safety dataset contains 1000 harmful and 1000 harmless prompts (to avoid over-refusal). We use separate rewards: a safety score for harmful prompts, and a helpfulness reward for harmless prompts. Details are in Section˜B.1. We show that FRPO achieves similar training results as GRPO (Figure˜6) with no difference in safety evaluations (see below). Additionally, in Section˜B.3, we show that allowing for a larger KL by reducing β enables FRPO to find more reward-flat solutions that better preserve the reward under fine-tuning. Evaluations. we report the refusal rate on HarmBench and the harmfulness score from StrongREJECT (Mazeika et al., 2024; Souly et al., 2024) under three downstream fine-tuning settings: (i) SFT on Alpaca (Taori et al., 2023), which prior work (Qi et al., 2023) reports can harm the guardrails; (i) SFT on GSM8K (Cobbe et al., 2021), testing vulnerability under math fine-tuning similar to (Qi et al., 2024); (i) RL on UltraFeedback (Cui et al., 2023), probing the helpfulness-harmlessness tradeoff (Bai et al., 2022; Tan et al., 2025). All results are averaged over 3 runs. Selection of λ. For FRPO, we sweep values in [0.05, 5] and select λ based on the safety reward after downstream fine-tuning. Importantly, we perform this selection only once using Alpaca and keep the same λ for all other downstream tasks. This is justified because the optimal λ in Equation˜3.3 depends on the KL range and reward statistics, not the specific downstream task. Baselines. Our primary baselines are GRPO-trained models using the same base model and safety dataset, which ensures a controlled comparison. For broader comparison, we include: (i) Llama-3.1-8B (Meta AI, 2024), a model with standard guardrails; (i) Llama-3-TAR-refusal (Tamirisa et al., 2024), which runs iterative adversarial fine-tuning to find vulnerable policies; (i) Llama-3-R and Mistral-R (Zou et al., 2024) models, which use an unlearning approach; (iv) Llama-3-derta (Yuan et al., 2025), which improves robustness to jailbreaking attacks. 5.1.1 SFT on Alpaca We evaluate the robustness of the guardrails to instruction-tuning by SFT on Alpaca (Taori et al., 2023), a chat-based fine-tuning dataset, following the setting of (Qi et al., 2023). We use a small learning rate (lr=10−6lr=10^-6) to better preserve general capabilities; in Section˜F.1, we show that larger learning rates degrade both general capabilities and safety reward. We fine-tune the models for 1 epoch, and use a constant learning rate (with a warm-up) so that the KL from the base model increases steadily (see Figure˜2, right). We first study how the value of λ impacts robustness, and find that λ=0.2λ=0.2 is optimal for both Mistral and Qwen models. We then compare the results with those of other baselines. Role of λ in the reward landscape. We fine-tune all models and measure the safety reward on the safety-training data. For different values of λ, Figure˜2 (left/middle) shows the trajectory of the safety reward on the landscape during fine-tuning as the KL increases. When decreasing λ from GRPO (λ→∞λ→∞), the final reward improves up to the optimal point and decreases afterwards. Early in fine-tuning, when the KL is small, GRPO has the highest reward because it optimizes the average reward. But as KL grows, its reward falls below other λ values. For Mistral, λ∈0.2,0.5λ∈\0.2,0.5\ perform best in the explored KL range. The model with λ=0.2λ=0.2 starts slightly below λ=0.5λ=0.5 for small KL, but surpasses it later as KL increases (larger ρ in Equation˜3.1 leads to a smaller λ in the dual problem). We therefore use λ=0.2λ=0.2 in the main comparisons. For Qwen, the separation is larger and λ=0.2λ=0.2 better preserves the reward. (a) (b) (c) (d) (e) (f) Figure 3: Safety evaluation after Alpaca SFT on HarmBench (↑ is better) and StrongREJECT (↓ is better). (a,d) and (b,e) compare models from the same reference (Mistral and Qwen), showing consistent improvements of our method over GRPO baselines. (c,f) Broader comparison with other safety-focused methods; our approach achieves the highest refusal rates on HarmBench and competitive StrongREJECT scores. Main results. We compare the results on HarmBench and StrongREJECT in Figure˜3. The first two columns compare methods starting from the same base model, showing gains over GRPO baselines and Mistral-R, even though Mistral-R effectively unlearns unsafe content. The right column compares against other baselines. After downstream fine-tuning, models trained with FRPO and λ=0.2λ=0.2 maintain the highest refusal rates on HarmBench and match Llama-3-R on StrongREJECT. As a sanity check, we also verify that improved safety is not due to overfitting or performing worse on the downstream task. This is studied in Section˜F.2: we show that the results on Alpaca are similar to those of GRPO and that general capabilities, measured by the helpfulness reward on UltraFeedback, remain similar across models. 5.1.2 SFT on GSM8K We test whether safety alignment is preserved under math fine-tuning with GSM8K (Cobbe et al., 2021). This is motivated by Qi et al. (2024) that show fine-tuning on GSM8K increases the attack success rate on Llama-2 models. We fine-tune the safety-trained models with SFT on GSM8K solutions for 1 epoch with lr=6×10−6lr=6× 10^-6. We compare only with the Mistral- and Qwen-based baselines here, and a broader comparison is deferred to Section˜F.3. Safety during fine-tuning. The safety evaluations on HarmBench and StrongREJECT can be seen in Figure˜4. The Qwen model trained with our method and λ=0.2λ=0.2 better preserves the safety on both benchmarks. The gap with GRPO is larger for Mistral, and our method significantly slows the safety degradation. In general, we observe a sharper drop in safety for Mistral models than for Qwen and other baselines under GSM8K SFT. We hypothesize that this is due to Mistral’s weaker GSM8K performance, leading to stronger training signals during fine-tuning. Is improved safety an artifact of reduced downstream adaptation? We analyze whether the preservation of safety guardrails could reflect a weaker downstream adaptation (i.e., the model changes less and therefore forgets less). We evaluate the models before and after GSM8K fine-tuning on the following benchmarks: GSM8K, MMLU (general capabilities), IFEval (instruction-following), and HumanEval (coding). As Table˜1 shows, post fine-tuning metrics on these benchmarks are similar for our method and GRPO. This indicates that the improved safety is not explained by a failure to learn the downstream task, but rather by finding updates that better preserve safety while maintaining comparable downstream performance. Figure 4: Safety metrics during GSM8k SFT for Mistral and Qwen models. Our method maintains higher refusal rates (left, ↑ is better) and better StrongREJECT scores (right, ↓ is better) compared to GRPO baselines. GSM8K (maj@8) MMLU (5-shot) IFEval (pass@1) HumanEval Base Models Base FT Base FT Base FT Base FT Mistral-v0.1-Instruct 50.5 57.4 55.0 53.7 39.6 37.2 28.7 23.2 Mistral-GRPO 48.4 55.5 54.9 54.1 34.4 38.0 32.3 24.4 Mistral-FRPO(λ=0.2) 48.7 56.2 55.2 54.2 36.6 37.0 31.1 24.4 Table 1: Downstream task performance before (Base) and after (FT) GSM8K fine-tuning. FRPO achieves similar scores to GRPO across all benchmarks. This confirms that the improved safety retention shown in Figure˜4 is not due to weaker downstream adaptation. 5.1.3 RL on Helpfulness We consider RL fine-tuning and test if the robustness observed under SFT persists. Prior work shows that harmlessness and helpfulness often conflict (Bai et al., 2022; Tan et al., 2025). We study this tradeoff by fine-tuning on 15k UltraFeedback prompts (Cui et al., 2023) with GRPO, using Llama-3.1-8B-Instruct-RM-RB2 (Malik et al., 2025) as the reward model. We again sweep λ and tune β to keep all models at roughly the same KL distance from their base policy. Helpfulness-safety trade-off under RL. UltraFeedback consists of complex prompts (e.g., multi-step coding tasks) that require long responses. Accordingly, the average response length increases during training (see Figure˜5, left). This affects the responses to unsafe prompts, and models tend to generate more detailed harmful content. Figure˜5 (right) shows StrongREJECT scores declining over training across all models (initial values shown in Figure˜2, left). We observe a clear tradeoff: smaller λ better preserves safety (StrongREJECT), while larger λ slightly improves helpfulness at the cost of safety. Overall, GRPO and λ∈2.0,1.0,0.5,0.2λ∈\2.0,1.0,0.5,0.2\ lie on a Pareto frontier. This behavior matches the intuition in Figure˜1: within a specific KL ball, the best λ yields a more reward-flat landscape, thus reducing drift toward helpfulness as the two objectives conflict. 5.2 Robustness in Math Training We test whether our method reduces forgetting of math capabilities after further fine-tuning. We start from Qwen2.5-Math-7B and train on MATH (Hendrycks et al., 2021) (levels 3–5) using the Qwen-Math template, following (Liu et al., 2025). We use an outcome-only 0–1 reward that verifies the final answer, with a small bonus when a final answer is provided (details in Section˜B.2). We train with GRPO and FRPO for λ∈[0.1,10]λ∈[0.1,10]. We first show that all models reach a similar accuracy on MATH500. We then fine-tune all math-trained models with SFT on 25k samples from "nvidia/OpenCodeInstruct" (Ahmad et al., 2025), an instruction-tuning dataset for code generation. The goal is to improve performance on coding benchmarks such as MBPP+ (Austin et al., 2021; Liu et al., 2023) while avoiding an accuracy drop on MATH500, thereby demonstrating improved continual learning across math and coding. Math training. Training results are provided in Section˜B.2. All the models behave similarly during training and converge to a similar accuracy: the top row of Table˜2 shows that all models reach roughly 73%73\% on MATH500, improving over Qwen2.5-Math-7B (56.6%). We also report MBPP+ results after math training in the third row of Table˜2. All models remain close to the base MBPP+ accuracy (57.4%). Figure 5: (left) Fine-tuning the models on UltraFeedback with GRPO leads to a significant increase in the average response length, inducing more detailed answers to harmful demands. (right) Helpfulness vs. Safety score (1 −- StrongREJECT score) for Mistral models after GRPO on UltraFeedback. λ=0.5λ=0.5 has better safety score but also lower helpfulness. λ=2.0λ=2.0 and GRPO have the higher helpfulness score. Code fine-tuning We fine-tune all math-trained models on "OpenCodeInstruct" using SFT. In initial experiments, full-parameter fine-tuning without LoRA caused a substantial drop in MATH500 accuracy, so we use LoRA with rank r=16r=16 and lr=5×10−5lr=5× 10^-5 for all models. Results on MATH500 and MBPP+ are summarized in Table˜2. FRPO with λ=2.0λ=2.0 preserves MATH500 accuracy best across models, outperforming GRPO by 22%. On MBPP+, all the models improve similarly, with a 44–5%5\% gain. The optimal λ here differs from the optimal λ in safety training. This is because the optimal λ must be derived from Equation˜3.3, and depends on training reward statistics; in particular, it is related to the reward standard deviation. In the math setting, the final standard deviation is roughly 10×10× larger than in safety training (∼0.5 0.5 vs. ∼0.05 0.05). GRPO (λ→∞λ →∞) FRPO (λ=10) FRPO (λ=4) FRPO (λ=2) FRPO (λ=1) FRPO (λ=0.5) FRPO (λ=0.2) FRPO (λ=0.1) MATH500 Base 73.0 73.0 73.2 73.0 73.0 73.4 73.4 71.6 FT 42.3 (± 2) 44.5 (± 3) 61.0 (± 3) 64.5 (± 2) 63.1(± 2) 53.6(± 2) 59.3 (± 4) 54.6 (± 3) MBPP+ Base 57.1 57.9 57.4 57.9 58.2 57.9 57.1 57.7 FT 62.5 62.2 61.2 62.4 62.0 61.8 61.7 61.7 Table 2: Performance before (Base) and after (FT) code fine-tuning all the math-trained model on "OpenCodeInstruct". FRPO with λ=2.0λ=2.0 best preserves math accuracy while achieving comparable coding performance. Results are averaged over 3 fine-tuning seeds. 6 Conclusion We proposed that robustness to downstream fine-tuning should be incorporated directly into the base policy during RLHF, rather than relying on later interventions. Our approach begins by optimizing reward stability within a KL-bounded neighborhood of policies. Solving this resulting max–min formulation yields FRPO, a robust policy gradient method that identifies reward-flat regions in policy space. Experiments demonstrate that this form of robustness transfers across diverse domains, including safety alignment and mathematical reasoning. Notably, FRPO maintained up to 22% higher mathematical accuracy under code fine-tuning. This work opens several promising directions. The principle of optimizing for robustness to future adaptation may extend beyond RLHF to pretraining and supervised fine-tuning. Moreover, understanding which capabilities are inherently easier, or harder, to make robust remains an open question, with implications for alignment and continual learning. Acknowledgment This research has been supported by Coefficient Giving and the UK AI Security Institute. AJ was supported in part by the Sloan fellowship in mathematics, the NSF Award DMS-2311024, an Amazon Faculty Research Award, an Adobe Faculty Research Award and an iORB grant form USC Marshall School of Business. References J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al. (2023) Gpt-4 technical report. arXiv preprint arXiv:2303.08774. Cited by: §1. A. Aghajanyan, A. Shrivastava, A. Gupta, N. Goyal, L. Zettlemoyer, and S. Gupta (2020) Better fine-tuning by reducing representational collapse. arXiv preprint arXiv:2008.03156. Cited by: Appendix A. W. U. Ahmad, A. Ficek, M. Samadi, J. Huang, V. Noroozi, S. Majumdar, and B. Ginsburg (2025) OpenCodeInstruct: a large-scale instruction tuning dataset for code llms. arXiv preprint arXiv:2504.04030. Cited by: §5.2. A. Ahmadian, C. Cremer, M. Gallé, M. Fadaee, J. Kreutzer, O. Pietquin, A. Üstün, and S. Hooker (2024) Back to basics: revisiting reinforce style optimization for learning from human feedback in llms. arXiv preprint arXiv:2402.14740. Cited by: §1.1. M. Ahn, A. Brohan, N. Brown, Y. Chebotar, O. Cortes, B. David, C. Finn, C. Fu, K. Gopalakrishnan, K. Hausman, et al. (2022) Do as i can, not as i say: grounding language in robotic affordances. arXiv preprint arXiv:2204.01691. Cited by: §1. J. Austin, A. Odena, M. Nye, M. Bosma, H. Michalewski, D. Dohan, E. Jiang, C. Cai, M. Terry, Q. Le, et al. (2021) Program synthesis with large language models. arXiv preprint arXiv:2108.07732. Cited by: §5.2. Y. Bai, A. Jones, K. Ndousse, A. Askell, A. Chen, N. DasSarma, D. Drain, S. Fort, D. Ganguli, T. Henighan, et al. (2022) Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862. Cited by: §1.1, §5.1, §5.1.3. D. Biderman, J. Portes, J. J. G. Ortiz, M. Paul, P. Greengard, C. Jennings, D. King, S. Havens, V. Chiley, J. Frankle, et al. (2024) Lora learns less and forgets less. arXiv preprint arXiv:2405.09673. Cited by: Appendix A. J. Cha, S. Chun, K. Lee, H. Cho, S. Park, Y. Lee, and S. Park (2021) Swad: domain generalization by seeking flat minima. Advances in Neural Information Processing Systems 34, p. 22405–22418. Cited by: §1.1. Y. Chow, A. Tamar, S. Mannor, and M. Pavone (2015) Risk-sensitive and robust decision-making: a cvar optimization approach. In Advances in Neural Information Processing Systems, C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (Eds.), Vol. 28, p. . External Links: Link Cited by: §1.1. P. F. Christiano, J. Leike, T. Brown, M. Martic, S. Legg, and D. Amodei (2017) Deep reinforcement learning from human preferences. Advances in neural information processing systems 30. Cited by: §1.1. W. Chung, V. Thomas, M. C. Machado, and N. Le Roux (2021) Beyond variance reduction: understanding the true impact of baselines on policy optimization. In International conference on machine learning, p. 1999–2009. Cited by: Appendix D, §4. K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, et al. (2021) Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168. Cited by: §5.1, §5.1.2. G. Cui, L. Yuan, N. Ding, G. Yao, B. He, W. Zhu, Y. Ni, G. Xie, R. Xie, Y. Lin, et al. (2023) Ultrafeedback: boosting language models with scaled ai feedback. arXiv preprint arXiv:2310.01377. Cited by: §5.1, §5.1.3. J. Cui, W. Chiang, I. Stoica, and C. Hsieh (2024) Or-bench: an over-refusal benchmark for large language models. arXiv preprint arXiv:2405.20947. Cited by: §B.1. M. De Lange, R. Aljundi, M. Masana, S. Parisot, X. Jia, A. Leonardis, G. Slabaugh, and T. Tuytelaars (2021) A continual learning survey: defying forgetting in classification tasks. IEEE transactions on pattern analysis and machine intelligence 44 (7), p. 3366–3385. Cited by: Appendix A, §1. C. de Masson D’Autume, S. Ruder, L. Kong, and D. Yogatama (2019) Episodic memory in lifelong language learning. Advances in Neural Information Processing Systems 32. Cited by: §1. E. Derman and S. Mannor (2020) Distributional robustness and regularization in reinforcement learning. arXiv preprint arXiv:2003.02894. Cited by: §1.1. A. Djuhera, S. R. Kadhe, F. Ahmed, S. Zawad, and H. Boche (2025) SafeMERGE: preserving safety alignment in fine-tuned large language models via selective layer-wise model merging. arXiv preprint arXiv:2503.17239. Cited by: Appendix A, §1. J. Dodge, G. Ilharco, R. Schwartz, A. Farhadi, H. Hajishirzi, and N. Smith (2020) Fine-tuning pretrained language models: weight initializations, data orders, and early stopping. arXiv preprint arXiv:2002.06305. Cited by: §2. D. Driess, F. Xia, M. S. Sajjadi, C. Lynch, A. Chowdhery, A. Wahid, J. Tompson, Q. Vuong, T. Yu, W. Huang, et al. (2023) Palm-e: an embodied multimodal language model. arXiv preprint arXiv:2303.03378. Cited by: §1. J. C. Duchi and H. Namkoong (2021) Learning models with uniform performance via distributionally robust optimization. The Annals of Statistics 49 (3), p. 1378–1406. Cited by: §1.1, §3. B. Eysenbach and S. Levine (2021) Maximum entropy rl (provably) solves some robust rl problems. arXiv preprint arXiv:2103.06257. Cited by: §1.1. Y. Fei, Z. Yang, Y. Chen, and Z. Wang (2021a) Exponential bellman equation and improved regret bounds for risk-sensitive reinforcement learning. Advances in neural information processing systems 34, p. 20436–20446. Cited by: §1.1. Y. Fei, Z. Yang, and Z. Wang (2021b) Risk-sensitive reinforcement learning with function approximation: a debiasing approach. In International Conference on Machine Learning, p. 3198–3207. Cited by: §1.1. P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur (2020) Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412. Cited by: §1.1. R. M. French (1999) Catastrophic forgetting in connectionist networks. Trends in cognitive sciences 3 (4), p. 128–135. Cited by: Appendix A. S. Ghosh, P. Varshney, M. N. Sreedhar, A. Padmakumar, T. Rebedea, J. R. Varghese, and C. Parisien (2025) Aegis2. 0: a diverse ai safety dataset and risks taxonomy for alignment of llm guardrails. arXiv preprint arXiv:2501.09004. Cited by: §B.1. I. J. Goodfellow, M. Mirza, D. Xiao, A. Courville, and Y. Bengio (2013) An empirical investigation of catastrophic forgetting in gradient-based neural networks. arXiv preprint arXiv:1312.6211. Cited by: Appendix A. T. L. Hayes, N. D. Cahill, and C. Kanan (2019) Memory efficient experience replay for streaming learning. In 2019 International Conference on Robotics and Automation (ICRA), p. 9769–9776. Cited by: Appendix A. D. Hendrycks, C. Burns, S. Kadavath, A. Arora, S. Basart, E. Tang, D. Song, and J. Steinhardt (2021) Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874. Cited by: §5.2. C. Hsu, Y. Tsai, C. Lin, P. Chen, C. Yu, and C. Huang (2024) Safe lora: the silver lining of reducing safety risks when finetuning large language models. Advances in Neural Information Processing Systems 37, p. 65072–65094. Cited by: Appendix A, §1. Y. Hsu, Y. Liu, A. Ramasamy, and Z. Kira (2018) Re-evaluating continual learning scenarios: a categorization and case for strong baselines. arXiv preprint arXiv:1810.12488. Cited by: Appendix A. E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, W. Chen, et al. (2022) Lora: low-rank adaptation of large language models.. ICLR 1 (2), p. 3. Cited by: Appendix A, §B.1, §1, §2. J. Huang, L. Cui, A. Wang, C. Yang, X. Liao, L. Song, J. Yao, and J. Su (2024) Mitigating catastrophic forgetting in large language models with self-synthesized rehearsal. arXiv preprint arXiv:2403.01244. Cited by: Appendix A, §1. G. Ilharco, M. T. Ribeiro, M. Wortsman, S. Gururangan, L. Schmidt, H. Hajishirzi, and A. Farhadi (2022) Editing models with task arithmetic. arXiv preprint arXiv:2212.04089. Cited by: Appendix A, §1. A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. S. Chaplot, D. de las Casas, F. Bressand, G. Lengyel, G. Lample, L. Saulnier, L. R. Lavaud, M. Lachaux, P. Stock, T. L. Scao, T. Lavril, T. Wang, T. Lacroix, and W. E. Sayed (2023) Mistral 7b. External Links: 2310.06825, Link Cited by: §5.1. L. Jiang, K. Rao, S. Han, A. Ettinger, F. Brahman, S. Kumar, N. Mireshghallah, X. Lu, M. Sap, Y. Choi, et al. (2024) Wildteaming at scale: from in-the-wild jailbreaks to (adversarially) safer language models. Advances in Neural Information Processing Systems 37, p. 47094–47165. Cited by: §B.1. Y. Jiang, B. Neyshabur, H. Mobahi, D. Krishnan, and S. Bengio (2019) Fantastic generalization measures and where to find them. arXiv preprint arXiv:1912.02178. Cited by: §1.1. J. Jiao and Y. Han (2020) Bias correction with jackknife, bootstrap, and taylor series. IEEE Transactions on Information Theory 66 (7), p. 4392–4418. Cited by: Appendix E. N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang (2016) On large-batch training for deep learning: generalization gap and sharp minima. arXiv preprint arXiv:1609.04836. Cited by: §1.1. M. Kim, D. Li, S. X. Hu, and T. Hospedales (2022) Fisher sam: information geometry and sharpness aware minimisation. In International Conference on Machine Learning, p. 11148–11161. Cited by: §1.1. T. Kim, F. Tajwar, A. Raghunathan, and A. Kumar (2025) Reasoning as an adaptive defense for safety. arXiv preprint arXiv:2507.00971. Cited by: §B.1. J. Kirkpatrick, R. Pascanu, N. Rabinowitz, J. Veness, G. Desjardins, A. A. Rusu, K. Milan, J. Quan, T. Ramalho, A. Grabska-Barwinska, et al. (2017) Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences 114 (13), p. 3521–3526. Cited by: Appendix A, Appendix A, §1. S. Kotha, J. M. Springer, and A. Raghunathan (2023) Understanding catastrophic forgetting in language models via implicit inference. arXiv preprint arXiv:2309.10105. Cited by: Appendix A. D. Kuhn, P. M. Esfahani, V. A. Nguyen, and S. Shafieezadeh-Abadeh (2019) Wasserstein distributionally robust optimization: theory and applications in machine learning. In Operations research & management science in the age of analytics, p. 130–166. Cited by: §1.1. J. Kwon, J. Kim, H. Park, and I. K. Choi (2021) Asam: adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In International conference on machine learning, p. 5905–5914. Cited by: §1.1. C. Lee, K. Cho, and W. Kang (2019) Mixout: effective regularization to finetune large-scale pretrained language models. arXiv preprint arXiv:1909.11299. Cited by: Appendix A, §1. Z. Li and D. Hoiem (2017) Learning without forgetting. IEEE transactions on pattern analysis and machine intelligence 40 (12), p. 2935–2947. Cited by: Appendix A, §1. J. Liu, C. S. Xia, Y. Wang, and L. Zhang (2023) Is your code generated by chatgpt really correct? rigorous evaluation of large language models for code generation. Advances in Neural Information Processing Systems 36, p. 21558–21572. Cited by: §5.2. Z. Liu, C. Chen, W. Li, P. Qi, T. Pang, C. Du, W. S. Lee, and M. Lin (2025) Understanding r1-zero-like training: a critical perspective. arXiv preprint arXiv:2503.20783. Cited by: §B.2, §2, §5.2. Y. Luo, Z. Yang, F. Meng, Y. Li, J. Zhou, and Y. Zhang (2025) An empirical study of catastrophic forgetting in large language models during continual fine-tuning. IEEE Transactions on Audio, Speech and Language Processing. Cited by: Appendix A. S. Malik, V. Pyatkin, S. Land, J. Morrison, N. A. Smith, H. Hajishirzi, and N. Lambert (2025) RewardBench 2: advancing reward model evaluation. arXiv preprint arXiv:2506.01937. Cited by: §B.1, §F.1, §5.1.3. T. Markov, C. Zhang, S. Agarwal, F. E. Nekoul, T. Lee, S. Adler, A. Jiang, and L. Weng (2023) A holistic approach to undesired content detection in the real world. In Proceedings of the AAAI conference on artificial intelligence, Vol. 37(12), p. 15009–15018. Cited by: §B.1. M. Mazeika, L. Phan, X. Yin, A. Zou, Z. Wang, N. Mu, E. Sakhaee, N. Li, S. Basart, B. Li, et al. (2024) Harmbench: a standardized evaluation framework for automated red teaming and robust refusal. arXiv preprint arXiv:2402.04249. Cited by: §5.1. M. McCloskey and N. J. Cohen (1989) Catastrophic interference in connectionist networks: the sequential learning problem. In Psychology of learning and motivation, Vol. 24, p. 109–165. Cited by: Appendix A. J. Mei, W. Chung, V. Thomas, B. Dai, C. Szepesvari, and D. Schuurmans (2022) The role of baselines in policy gradient optimization. Advances in Neural Information Processing Systems 35, p. 17818–17830. Cited by: Appendix D, §4. Meta AI (2024) The llama 3 herd of models. External Links: arXiv:2407.21783, Link Cited by: §5.1. R. G. Miller (1974) The jackknife–a review. Biometrika 61 (1), p. 1–15. External Links: ISSN 00063444, 14643510, Link Cited by: Appendix E. M. Mosbach, M. Andriushchenko, and D. Klakow (2020) On the stability of fine-tuning bert: misconceptions, explanations, and strong baselines. arXiv preprint arXiv:2006.04884. Cited by: §2. A. H. Nobari, K. Alim, A. ArjomandBigdeli, A. Srivastava, F. Ahmed, and N. Azizan (2025) Activation-informed merging of large language models. External Links: 2502.02421, Link Cited by: Appendix A. Y. Oren, S. Sagawa, T. B. Hashimoto, and P. Liang (2019) Distributionally robust language modeling. arXiv preprint arXiv:1909.02060. Cited by: §1.1. T. Osogami (2012) Robustness and risk-sensitivity in markov decision processes. Advances in neural information processing systems 25. Cited by: §1.1. L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. (2022) Training language models to follow instructions with human feedback. Advances in neural information processing systems 35, p. 27730–27744. Cited by: §1.1, §1, §1, §2, §2, §2. A. B. Owen (2013) Monte carlo theory, methods and examples. https://artowen.su.domains/mc/. Cited by: Appendix E. R. Pan, X. Liu, S. Diao, R. Pi, J. Zhang, C. Han, and T. Zhang (2024) Lisa: layerwise importance sampling for memory-efficient large language model fine-tuning. Advances in Neural Information Processing Systems 37, p. 57018–57049. Cited by: Appendix A. L. Pinto, J. Davidson, R. Sukthankar, and A. Gupta (2017) Robust adversarial reinforcement learning. In International conference on machine learning, p. 2817–2826. Cited by: §1.1. A. Prabhu, P. H. Torr, and P. K. Dokania (2020) Gdumb: a simple approach that questions our progress in continual learning. In European conference on computer vision, p. 524–540. Cited by: Appendix A. X. Qi, A. Panda, K. Lyu, X. Ma, S. Roy, A. Beirami, P. Mittal, and P. Henderson (2024) Safety alignment should be made more than just a few tokens deep. arXiv preprint arXiv:2406.05946. Cited by: Appendix A, §1, §5.1, §5.1.2. X. Qi, Y. Zeng, T. Xie, P. Chen, R. Jia, P. Mittal, and P. Henderson (2023) Fine-tuning aligned language models compromises safety, even when users do not intend to!. arXiv preprint arXiv:2310.03693. Cited by: Appendix A, §1, §5.1, §5.1.1. F. Qiao and M. Mahdavi (2024) Learn more, but bother less: parameter efficient continual learning. In Advances in Neural Information Processing Systems, A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang (Eds.), Vol. 37, p. 97476–97498. External Links: Document, Link Cited by: Appendix A, §1. Qwen, :, A. Yang, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Li, D. Liu, F. Huang, H. Wei, H. Lin, J. Yang, J. Tu, J. Zhang, J. Yang, J. Yang, J. Zhou, J. Lin, K. Dang, K. Lu, K. Bao, K. Yang, L. Yu, M. Li, M. Xue, P. Zhang, Q. Zhu, R. Men, R. Lin, T. Li, T. Tang, T. Xia, X. Ren, X. Ren, Y. Fan, Y. Su, Y. Zhang, Y. Wan, Y. Liu, Z. Cui, Z. Zhang, and Z. Qiu (2025) Qwen2.5 technical report. External Links: 2412.15115, Link Cited by: §5.1. R. Rafailov, A. Sharma, E. Mitchell, C. D. Manning, S. Ermon, and C. Finn (2023) Direct preference optimization: your language model is secretly a reward model. Advances in neural information processing systems 36, p. 53728–53741. Cited by: §1.1. R. Ratcliff (1990) Connectionist models of recognition memory: constraints imposed by learning and forgetting functions.. Psychological review 97 (2), p. 285. Cited by: Appendix A. D. Rolnick, A. Ahuja, J. Schwarz, T. Lillicrap, and G. Wayne (2019) Experience replay for continual learning. Advances in neural information processing systems 32. Cited by: Appendix A, §1. D. Rosati, J. Wehner, K. Williams, L. Bartoszcze, R. Gonzales, S. Majumdar, H. Sajjad, F. Rudzicz, et al. (2024) Representation noising: a defence mechanism against harmful finetuning. Advances in Neural Information Processing Systems 37, p. 12636–12676. Cited by: §1.1. A. A. Rusu, N. C. Rabinowitz, G. Desjardins, H. Soyer, J. Kirkpatrick, K. Kavukcuoglu, R. Pascanu, and R. Hadsell (2016) Progressive neural networks. arXiv preprint arXiv:1606.04671. Cited by: Appendix A. S. Sagawa, P. W. Koh, T. B. Hashimoto, and P. Liang (2019) Distributionally robust neural networks for group shifts: on the importance of regularization for worst-case generalization. arXiv preprint arXiv:1911.08731. Cited by: §1.1. M. Samvelyan, S. C. Raparthy, A. Lupu, E. Hambro, A. Markosyan, M. Bhatt, Y. Mao, M. Jiang, J. Parker-Holder, J. Foerster, et al. (2024) Rainbow teaming: open-ended generation of diverse adversarial prompts. Advances in Neural Information Processing Systems 37, p. 69747–69786. Cited by: §B.1. T. Schick, J. Dwivedi-Yu, R. Dessì, R. Raileanu, M. Lomeli, E. Hambro, L. Zettlemoyer, N. Cancedda, and T. Scialom (2023) Toolformer: language models can teach themselves to use tools. Advances in Neural Information Processing Systems 36, p. 68539–68551. Cited by: §1. J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz (2015) Trust region policy optimization. In International conference on machine learning, p. 1889–1897. Cited by: §2. J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov (2017) Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. Cited by: Appendix A, §1.1, §1, §2, §2. T. Scialom, T. Chakrabarty, and S. Muresan (2022) Fine-tuned language models are continual learners. arXiv preprint arXiv:2205.12393. Cited by: Appendix A, §1. Z. Shao, P. Wang, Q. Zhu, R. Xu, J. Song, X. Bi, H. Zhang, M. Zhang, Y. Li, Y. Wu, et al. (2024) Deepseekmath: pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300. Cited by: Appendix D, §1.1, §1, §2, §2. A. Shapiro and A. Kleywegt (2002) Minimax analysis of stochastic problems. Optimization Methods and Software 17 (3), p. 523–542. External Links: Document, https://doi.org/10.1080/1055678021000034008 Cited by: §1.1. A. Shapiro (2017) Distributionally robust stochastic programming. SIAM Journal on Optimization 27 (4), p. 2258–2275. External Links: Document, Link, https://doi.org/10.1137/16M1058297 Cited by: §1.1, §3. A. Sinha, H. Namkoong, R. Volpi, and J. Duchi (2017) Certifying some distributional robustness with principled adversarial training. arXiv preprint arXiv:1710.10571. Cited by: §1.1. E. Smirnova, E. Dohmatob, and J. Mary (2019) Distributionally robust reinforcement learning. arXiv preprint arXiv:1902.08708. Cited by: §1.1. A. Souly, Q. Lu, D. Bowen, T. Trinh, E. Hsieh, S. Pandey, P. Abbeel, J. Svegliato, S. Emmons, O. Watkins, et al. (2024) A strongreject for empty jailbreaks. Advances in Neural Information Processing Systems 37, p. 125416–125440. Cited by: §B.1, §5.1. N. Stiennon, L. Ouyang, J. Wu, D. Ziegler, R. Lowe, C. Voss, A. Radford, D. Amodei, and P. F. Christiano (2020) Learning to summarize with human feedback. Advances in neural information processing systems 33, p. 3008–3021. Cited by: §1.1. F. Sun, C. Ho, and H. Lee (2019) Lamol: language modeling for lifelong language learning. arXiv preprint arXiv:1909.03329. Cited by: Appendix A, §1, §1. A. Tamar, Y. Glassner, and S. Mannor (2015) Optimizing the cvar via sampling. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 29(1). Cited by: §1.1. R. Tamirisa, B. Bharathi, L. Phan, A. Zhou, A. Gatti, T. Suresh, M. Lin, J. Wang, R. Wang, R. Arel, et al. (2024) Tamper-resistant safeguards for open-weight llms. arXiv preprint arXiv:2408.00761. Cited by: §F.2, §1.1, §5.1. Y. Tan, Y. Jiang, Y. Li, J. Liu, X. Bu, W. Su, X. Yue, X. Zhu, and B. Zheng (2025) Equilibrate rlhf: towards balancing helpfulness-safety trade-off in large language models. arXiv preprint arXiv:2502.11555. Cited by: §5.1, §5.1.3. R. Taori, I. Gulrajani, T. Zhang, Y. Dubois, X. Li, C. Guestrin, P. Liang, and T. B. Hashimoto (2023) Stanford alpaca: an instruction-following llama model. GitHub. Note: https://github.com/tatsu-lab/stanford_alpaca Cited by: §2, §5.1, §5.1.1. E. Vinitsky, Y. Du, K. Parvate, K. Jang, P. Abbeel, and A. Bayen (2020) Robust reinforcement learning using adversarial populations. arXiv preprint arXiv:2008.01825. Cited by: §1.1. L. von Werra, Y. Belkada, L. Tunstall, E. Beeching, T. Thrush, N. Lambert, S. Huang, K. Rasul, and Q. Gallouédec (2020) TRL: transformer reinforcement learning. GitHub. Note: https://github.com/huggingface/trl Cited by: §B.1. L. Wang, X. Zhang, H. Su, and J. Zhu (2024) A comprehensive survey of continual learning: theory, method and application. IEEE transactions on pattern analysis and machine intelligence 46 (8), p. 5362–5383. Cited by: Appendix A. X. Wang, T. Chen, Q. Ge, H. Xia, R. Bao, R. Zheng, Q. Zhang, T. Gui, and X. Huang (2023) Orthogonal subspace learning for language model continual learning. In Findings of the Association for Computational Linguistics: EMNLP 2023, p. 10658–10671. Cited by: Appendix A. Z. Wang, Z. Zhang, S. Ebrahimi, R. Sun, H. Zhang, C. Lee, X. Ren, G. Su, V. Perot, J. Dy, et al. (2022a) Dualprompt: complementary prompting for rehearsal-free continual learning. In European conference on computer vision, p. 631–648. Cited by: §1. Z. Wang, Z. Zhang, C. Lee, H. Zhang, R. Sun, X. Ren, G. Su, V. Perot, J. Dy, and T. Pfister (2022b) Learning to prompt for continual learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, p. 139–149. Cited by: §1. M. Wortsman, G. Ilharco, S. Y. Gadre, R. Roelofs, R. Gontijo-Lopes, A. S. Morcos, H. Namkoong, A. Farhadi, Y. Carmon, S. Kornblith, et al. (2022) Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International conference on machine learning, p. 23965–23998. Cited by: Appendix A, §1. X. Yang, X. Wang, Q. Zhang, L. Petzold, W. Y. Wang, X. Zhao, and D. Lin (2023) Shadow alignment: the ease of subverting safely-aligned language models. arXiv preprint arXiv:2310.02949. Cited by: Appendix A. X. Yi, S. Zheng, L. Wang, X. Wang, and L. He (2024) A safety realignment framework via subspace-oriented model fusion for large language models. Knowledge-Based Systems 306, p. 112701. Cited by: Appendix A, §1. Y. Yuan, W. Jiao, W. Wang, J. Huang, J. Xu, T. Liang, P. He, and Z. Tu (2025) Refuse whenever you feel unsafe: improving safety in llms via decoupled refusal training. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), p. 3149–3167. Cited by: §F.2, §5.1. F. Zenke, B. Poole, and S. Ganguli (2017) Continual learning through synaptic intelligence. In International conference on machine learning, p. 3987–3995. Cited by: Appendix A. Q. Zhan, R. Fang, R. Bindu, A. Gupta, T. Hashimoto, and D. Kang (2023) Removing rlhf protections in gpt-4 via fine-tuning. arXiv preprint arXiv:2311.05553. Cited by: Appendix A, §1. H. Zhang, H. Chen, C. Xiao, B. Li, M. Liu, D. Boning, and C. Hsieh (2020) Robust deep reinforcement learning against adversarial perturbations on state observations. Advances in neural information processing systems 33, p. 21024–21037. Cited by: §1.1. J. Zhuang, B. Gong, L. Yuan, Y. Cui, H. Adam, N. Dvornek, S. Tatikonda, J. Duncan, and T. Liu (2022) Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.08065. Cited by: §1.1. A. Zou, L. Phan, J. Wang, D. Duenas, M. Lin, M. Andriushchenko, J. Z. Kolter, M. Fredrikson, and D. Hendrycks (2024) Improving alignment and robustness with circuit breakers. Advances in Neural Information Processing Systems 37, p. 83345–83373. Cited by: §F.3, §1.1, §5.1. Appendix A Additional Related Work Catastrophic forgetting. Catastrophic forgetting—the abrupt loss of previously learned knowledge when training on new tasks—has been studied since early neural network research (McCloskey and Cohen, 1989; Ratcliff, 1990; French, 1999). The phenomenon arises because gradient updates for new objectives overwrite parameters critical to earlier tasks (Goodfellow et al., 2013; Kirkpatrick et al., 2017). In LLMs, fine-tuning degrades capabilities acquired during pretraining or alignment (Kotha et al., 2023; Luo et al., 2025). Recent work shows that even benign fine-tuning can remove safety guardrails (Qi et al., 2023; Yang et al., 2023), and task-specific adaptation (e.g., math) degrades safety (Qi et al., 2024; Zhan et al., 2023). This motivates building robustness into the pre-fine-tuning policy rather than relying solely on downstream interventions. Continual learning. Continual learning methods aim to learn new tasks while preserving prior knowledge (De Lange et al., 2021; Wang et al., 2024). Regularization-based approaches constrain updates: EWC (Kirkpatrick et al., 2017) uses Fisher information, synaptic intelligence (SI) (Zenke et al., 2017) accumulates importance online, and LwF (Li and Hoiem, 2017) applies knowledge distillation. For LLMs, MixOut (Lee et al., 2019) stochastically resets toward pretrained weights, while other methods prevent the policy from drifting too far (Li and Hoiem, 2017; Schulman et al., 2017; Lee et al., 2019), or ensure that the model’s latent representations are preserved (Kirkpatrick et al., 2017; Aghajanyan et al., 2020; Pan et al., 2024). Rehearsal methods augment downstream training with a subset of previous data or synthetic data generated by the model (Rolnick et al., 2019; Sun et al., 2019; Scialom et al., 2022; Huang et al., 2024). Parameter-efficient fine-tuning (PEFT) constrain downstream updates to separate modules (e.g., LoRA (Hu et al., 2022) and progressive networks (Rusu et al., 2016)), thereby reducing parameter interference between the objectives (Hsu et al., 2024; Qiao and Mahdavi, 2024), or aiming to keep the features orthogonal (Wang et al., 2023; Qiao and Mahdavi, 2024). Finally, recent model merging methods aim to mathematically combine task-specific models post-hoc to retain all capabilities without expensive retraining (Ilharco et al., 2022; Wortsman et al., 2022; Yi et al., 2024; Djuhera et al., 2025; Nobari et al., 2025). All these methods operate at downstream time; our approach instead builds robustness into the upstream policy, making it agnostic to the downstream fine-tuning protocol. Limitations of downstream methods. Downstream methods have several limitations. Replay-based methods require curating samples that cover all capabilities prone to being forgotten. Also, buffer size and sample selection impact both effectiveness and training efficiency, and poor choices lead to continued forgetting (Hayes et al., 2019; Prabhu et al., 2020). Regularization approaches like EWC and SI have been shown to fail when task similarity is low (Hsu et al., 2018). Parameter-efficient methods, despite widespread adoption, do not reliably prevent forgetting—LoRA fine-tuning still degrades prior capabilities, sometimes comparably to full fine-tuning (Biderman et al., 2024). Most critically, all these methods assume the downstream optimization follows a prescribed recipe. Our upstream approach avoids these issues by building robustness directly into the base policy, making it agnostic to how downstream adaptation is performed. Appendix B Training Details B.1 Safety Training In order to fine-tune the models with our algorithm, we modified TRL’s implementation of GRPO (von Werra et al., 2020). We used 8xH200s for training our models. Dataset. The safety dataset, adapted from (Kim et al., 2025), contains 1000 harmful and 1000 harmless samples. Harmful prompts are from WildJailbreak, Aegis AI Content Safety Dataset 2.0, and RainbowTeaming jailbreaking prompts (Jiang et al., 2024; Ghosh et al., 2025; Samvelyan et al., 2024). Harmless prompts come from OR-Bench (Cui et al., 2024) and are used to preserve general instruction-following behavior and mitigate over-refusal. Reward models. We use separate rewards for harmful vs. harmless subsets. For harmful prompts, the reward is 1−(s1+s2)/2,1-(s_1+s_2)/2, where s1,s2s_1,s_2 are scores from OpenAI Moderation API and StrongREJECT judge (Markov et al., 2023; Souly et al., 2024). For harmless prompts, we use an off-the-shelf helpfulness reward model: “Llama-3.1-8B-Instruct-RM-RB2” (Malik et al., 2025). This component is only added to avoid over-refusal, and our focus is not on improving helpfulness. Hyper-parameters. To avoid overfitting in both GRPO and our algorithm, we use a relatively large β and tune it to keep the “per-token KL” near 0.10.1 for all models. This makes comparisons meaningful: each method is effectively searching for the best policy within a similar KL ball around the same reference model. We used group size G=8G=8 for all the experiments. As noted before, we compute the advantages by subtracting the group average reward in Equation˜4.1 but do not divide them by the standard deviation. Instead, we ensure the rewards are scaled between 0 and 1—the safety reward is naturally between 0 and 1, the normal reward model’s output is passed through a sigmoid function—to keep both the safety and normal signals relevant. We use LoRA (Hu et al., 2022) with r=64r=64 and α=64α=64 for safety training of all the models. The learning-rate is lr=10−5lr=10^-5 for the Mistral models and lr=3×10−5lr=3× 10^-5 for Qwen models. In order to keep the gradient norm consistent across values of λ, we omitted the λ factor in Equation˜4.1 and tuned β to keep the final KL bounded, rather than changing the learning-rate for each λ. We used 2 epochs on 2000 samples of the training data for training the Mistral models and 3 epochs for the Qwen models. We found that the Qwen models need more steps for convergence as they begin with higher rewards and a smaller standard deviation, leading to smaller gradient signals. Results. The training results are shown in Figure˜6, where the safety rewards converge for all the models with negligibly higher rewards for GRPO. As demonstrated in Figure˜6(b) and (e), the models trained with our algorithm achieve higher rewards on harmless prompts; as λ decreases, our algorithm becomes more sensitive to a response with a low reward (near-zero reward for over-refusal) when the group average is high, leading to a large signal in Equation˜4.1. Finally, Figure˜6(c) and (f) show that we keep the KL to the reference models bounded after training. The values of the allowed KL are determined so that the normal reward and the policy’s entropy are not affected. (a) (b) (c) (d) (e) (f) Figure 6: Safety training curves for Mistral and Qwen with GRPO and FRPO for λ∈0.5,0.2,0.1λ∈\0.5,0.2,0.1\. (a, d) show that all Mistral and Qwen models converge in the safety score. (b,e) show that the Helpfulness (Normal) reward for the models slightly improve on OR-Bench, meaning that over-refusals are avoided. (c, f) We tune β such that all models converge to roughly the same per-token-KL to ensure a controlled comparison. B.2 Math Training As described in Section˜5.2, we use a final answer verifier as the 0-1 reward for the RL training. We also add a small format reward if there is only one final answer within " ". We average the two rewards with weights [0.8,0.2][0.8,0.2]. The Qwen-Math system prompt is: "Please reason step by step, and put your final answer within .". We use lr=6×10−6lr=6× 10^-6 for both GRPO and FRPO with λ∈10.0,4.0,2.0,1.0λ∈\10.0,4.0,2.0,1.0\ and lr=5×10−6lr=5× 10^-6 for λ∈0.5,0.2,0.1λ∈\0.5,0.2,0.1\ due to larger gradient norms for smaller λ. We choose a small β=10−4β=10^-4 for all the models, consistent with other math training settings (Liu et al., 2025). We use LoRA adapters with r=64r=64 and α=128α=128. We train all the models for 3 epochs on the MATH training dataset (Levels 3-5). Results. The results of the training are presented in Figure˜7, where all the models roughly converge to the same point. Figure˜7 (left) shows that the format reward converges to ≈0.9≈ 0.9 for all models. Figure˜7 (right) shows that the initial training accuracy is 44% for all the models, and all of them reach around 70% by the end of training. The first row of Table˜2 shows the final accuracy on MATH500 for all the models. Figure 7: We have two rewards for math training: (left) shows that the format reward (whether the response contains any final answer) increases to ∼ 0.9 for all the training models; (right) shows that the correctness reward increases similarly for all model, and roughly reaches 0.7. B.3 Ablation: Effect of Training KL Budget on Downstream Robustness During safety training, we tune β to control the KL divergence to the reference model. Specifically, β is chosen such that the resulting KL preserves the helpfulness reward while maintaining roughly constant policy entropy (noting that smaller β typically reduces output entropy). Among such KL values below a given threshold, Figure˜8 shows that larger KL (smaller β) yields policies with higher safety reward after downstream fine-tuning. This occurs because Equation˜3.5 defines optimization over a KL ball around πref _ref whose radius shrinks with larger β; a smaller ball overly constrains the reward-flatness objective in the first term −x∼p[λlog(y∼πθ(⋅∣x)e−r(x,y)/λ)]- E_x p\! [λ \! ( E_y _θ(· x)e^-r(x,y)/λ ) ], hindering its optimization. Figure 8: During training, among KL values that preserve the helpfulness reward and avoid reducing the policy entropy, using a larger KL (i.e., smaller β) yields a more robust solution with flatter rewards. Appendix C Proof of Lemma 3.1 Denoting the likelihood ratio of Q over πθ _θ as L(y∣x):=Q(y∣x)πθ(y∣x)L(y x)\ :=\ Q(y x) _θ(y x), we can rewrite the KL term as: KL(Q(⋅∣x)∥π(⋅∣x))=∫f(Q(y∣x)π(y∣x))π(dy∣x)=y∼π(⋅∣x)[f(L(y∣x))] KL\! (Q(· x)\, \|\,π(· x) )= f ( Q(y x)π(y x) )π(dy x)= E_y π(· x) [f(L(y x)) ] where f(x)=xlog(x)f(x)=x (x). Moreover, using the Radon–Nikodym derivative, we parametrize the Lagrangian for solving the infimum over Q as a function of L(y∣x)L(y x), a global multiplier λ≥0λ≥ 0 for the average KL constraint and a per-x multiplier η(x)η(x) for normalization (π[L∣x]=1 E_π[L x]=1): ℒ(L,λ,η)= (L,λ,η)= x∼py∼πθ(⋅∣x)[Lr(x,y)]+λ(x∼py∼πθ(⋅∣x)[f(L(y∣x)]−ρ) E_x p E_y _θ(· x) [\,L\,r(x,y) ]+λ ( E_x p E_y _θ(· x)[f(L(y x)]-ρ ) +∫η(x)(L(y∣x)−1)π(dy∣x)x + η(x) (L(y x)-1 )π(dy x)dx If we redefine η(x)←η(x)p(x)η(x)← η(x)p(x) as the new multiplier: ℒ(L,λ,η) (L,λ,η) =x∼py∼πθ(⋅∣x)[Lr(x,y)+λf(L)+η(x)(L−1)]−λρ = E_x p E_y _θ(· x) [\,L\,r(x,y)+λ f(L)+η(x) (L-1 )\, ]-λρ It is easy to see that the primal problem is convex since f(x)f(x) is a convex function, which results in strong duality minLmaxλ≥0,ηℒ=maxλ≥0,ηminLℒ _L _λ≥ 0,ηL= _λ≥ 0,η _LL. Thus, minimizing over L in terms of the Fenchel conjugate function gives: minLℒ= _LL= minLx∼py∼πθ(⋅∣x)[Lr(x,y)+λf(L)+η(x)(L)]−η(x)−λρ _L \ E_x p E_y _θ(· x) [\,L\,r(x,y)+λ f(L)+η(x) (L )\, ] \- Eη(x)-λρ = = x∼py∼πθ(⋅∣x)[minL(r(x,y)+η(x))L+λf(L)]−η(x)−λρ E_x p E_y _θ(· x) [\, _L \(r(x,y)+η(x))\,L+λ f(L) \ ]- Eη(x)-λρ = = x∼py∼πθ(⋅∣x)[−λf∗(−r(x,y)−η(x)λ)−η(x)]−λρ E_x p E_y _θ(· x) [-λ\,f^* ( -r(x,y)-η(x)λ )-η(x) ]-λρ (C.1) Appendix D Gradient of the Objective Figure 9: The policy collapses in the absence of the derived baseline in Equation˜4.2 for FRPO; i.e., the token distribution becomes random and the policy entropy explodes, and the helpfulness (Normal) reward collapses as the responses become gibberish. Figure˜9 shows that omitting the baseline in FRPO causes the policy to collapse, resulting in a completely random token distribution (maximum entropy in Figure˜9, right). This occurs because without a baseline, the variance is high. Furthermore, all trajectories are suppressed because the gradient assigns negative coefficients to all of them as we will show in Equation˜D.1. This convergence issue is discussed in prior work (Chung et al., 2021; Mei et al., 2022), which demonstrates that the baseline’s role extends beyond variance reduction and can determine the convergence point. We compute the gradient of J(θ)J(θ) defined in Equation˜4.1. We drop the thresholding for simplicity. Define: u(θ):=1G∑i=1G1|yi|∑t=1|yi|πθ(yi,t|x,yi,<t)πold(yi,t|x,yi,<t)e−Ai,t/λ.u(θ):= 1G _i=1^G 1|y_i| _t=1^|y_i| _θ(y_i,t|x,y_i,<t) _old(y_i,t|x,y_i,<t)e^-A_i,t/λ\,. So, ∇J(θ) ∇ J(θ) =x∼p[−λ∇u(θ)u(θ)−β1G∑i=1G1|yi|∑t=1|yi|(−πref(yi,t|x,yi,<t)πθ(yi,t|x,yi,<t)2∇πθ(yi,t|x,yi,<t)+∇logπθ(yi,t|x,yi,<t))] = E_x p\! [-λ ∇ u(θ)u(θ)-β 1G _i=1^G 1|y_i| _t=1^|y_i| (- _ref(y_i,t|x,y_i,<t) _θ(y_i,t|x,y_i,<t)^2∇ _θ(y_i,t|x,y_i,<t)+∇ _θ(y_i,t|x,y_i,<t) ) ] =x∼p[−λ∇u(θ)u(θ)−β1G∑i=1G1|yi|∑t=1|yi|(−πref(yi,t|x,yi,<t)πθ(yi,t|x,yi,<t)∇logπθ(yi,t|x,yi,<t)+∇logπθ(yi,t|x,yi,<t))] = E_x p\! [-λ ∇ u(θ)u(θ)-β 1G _i=1^G 1|y_i| _t=1^|y_i| (- _ref(y_i,t|x,y_i,<t) _θ(y_i,t|x,y_i,<t)∇ _θ(y_i,t|x,y_i,<t)+∇ _θ(y_i,t|x,y_i,<t) ) ] In addition, ∇u(θ)=1G∑i=1G1|yi|∑t=1|yi|∇πθ(yi,t|x,yi,<t)πold(yi,t|x,yi,<t)e−Ai,t/λ=1G∑i=1G1|yi|∑t=1|yi|πθ(yi,t|x,yi,<t)πold(yi,t|x,yi,<t)e−Ai,t/λ∇logπθ(yi,t|x,yi,<t)∇ u(θ)= 1G _i=1^G 1|y_i| _t=1^|y_i| ∇ _θ(y_i,t|x,y_i,<t) _old(y_i,t|x,y_i,<t)e^-A_i,t/λ= 1G _i=1^G 1|y_i| _t=1^|y_i| _θ(y_i,t|x,y_i,<t) _old(y_i,t|x,y_i,<t)e^-A_i,t/λ∇ _θ(y_i,t|x,y_i,<t) where we use the log-derivative trick to derive the gradient as in other policy gradient algorithms. Putting it back into ∇J(θ)∇ J(θ), we obtain: ∇J(θ) ∇ J(θ) =x∼p[1G∑i=1G1|yi|∑t=1|yi|−πθ(yi,t|x,yi,<t)πold(yi,t|x,yi,<t)λe−Ai,t/λu(θ)+βπref(yi,t|x,yi,<t)πθ(yi,t|x,yi,<t)−β∇logπθ(yi,t|x,yi,<t)] = E_x p\! [ 1G _i=1^G 1|y_i| _t=1^|y_i| \- _θ(y_i,t|x,y_i,<t) _old(y_i,t|x,y_i,<t) λ e^-A_i,t/λu(θ)+β _ref(y_i,t|x,y_i,<t) _θ(y_i,t|x,y_i,<t)-β \∇ _θ(y_i,t|x,y_i,<t) ] Now, similar to GRPO (Shao et al., 2024), we simplify the analysis by assuming that the model only has a single update following each exploration stage, thereby ensuring that πold=πθ _old= _θ. Doing this, we can write u(θ)=1G∑i=1G1|yi|∑t=1|yi|e−Ai,t/λ,u(θ)= 1G _i=1^G 1|y_i| _t=1^|y_i|e^-A_i,t/λ\,, and ∇J(θ) ∇ J(θ) =x∼p[1G∑i=1G1|yi|∑t=1|yi|−λe−Ai,t/λu(θ)+βπref(yi,t|x,yi,<t)πθ(yi,t|x,yi,<t)−β∇logπθ(yi,t|x,yi,<t)] = E_x p\! [ 1G _i=1^G 1|y_i| _t=1^|y_i| \- λ e^-A_i,t/λu(θ)+β _ref(y_i,t|x,y_i,<t) _θ(y_i,t|x,y_i,<t)-β \∇ _θ(y_i,t|x,y_i,<t) ] (D.1) The term above shows that adding a constant to the advantages of a group does not change the gradient, and it cancels out from the numerator and denominator of the first term in Equation˜D.1. Therefore, rewards can be replaced with advantages in Equation˜4.1. Baseline derivation. Note that if λ≫1λ 1, we can apply the Taylor expansion for the terms inside Equation˜D.1: ∇J(θ) ∇ J(θ) ≈x∼p[1G∑i=1G1|yi|∑t=1|yi|−λ(1−Ai,t/λ+O(1λ2))1−⟨Ai,t⟩/λ+O(1λ2)+βπref(yi,t|x,yi,<t)πθ(yi,t|x,yi,<t)−β∇logπθ(yi,t|x,yi,<t)] ≈ E_x p\! [ 1G _i=1^G 1|y_i| _t=1^|y_i| \-λ\, (1-A_i,t/λ+O( 1λ^2) )1- A_i,t /λ+O( 1λ^2)+β _ref(y_i,t|x,y_i,<t) _θ(y_i,t|x,y_i,<t)-β \∇ _θ(y_i,t|x,y_i,<t) ] =x∼p[1G∑i=1G1|yi|∑t=1|yi|(−λ+Ai,t+O(1λ))+βπref(yi,t|x,yi,<t)πθ(yi,t|x,yi,<t)−β∇logπθ(yi,t|x,yi,<t)] = E_x p\! [ 1G _i=1^G 1|y_i| _t=1^|y_i| \ (-λ+A_i,t+O( 1λ) )+β _ref(y_i,t|x,y_i,<t) _θ(y_i,t|x,y_i,<t)-β \∇ _θ(y_i,t|x,y_i,<t) ] In the second line, we note that ⟨Ai,t⟩=0 A_i,t =0 because we subtract the group average from each reward, similar to GRPO. Then, the λ term can be removed since y∼πθ(⋅|x)[∇θlogπθ(y|x)]=0 E_y _θ(·|x) [ _θ _θ(y|x) ]=0: any baseline b(x)b(x) that does not depend on y can be subtracted without changing the expectation of the gradient. This recovers the gradient from GRPO. This term would cause a large variance in cases where λ is large. Appendix E Bias of the Partition Function Figure 10: (left) The bias in the gradient estimator does not show itself in the safety training curves and the training without the jackknife trick looks similar. (right) However, at the downstream time, the model that is trained without the jackknife trick is less robust and the StrongREJECT score (↓ is better) grows more compared to the model with the Jackknife trick. Unlike the baseline discussed in Appendix˜D, the gradient bias arising from the Monte Carlo estimation of the log partition function in Equation˜E.1 does not impact the training curves, as shown in Figure˜10 (left). However, Figure˜10 (right) reveals that when the model is fine-tuned downstream, the reward drops more sharply and the harmfulness score increases relative to the model trained with the jackknife trick. We now explain why this bias appears and how we address it. Starting from the gradient derived in Equation˜D.1, consider the first term: ∇θJ=−x∼p[λ∑ie−A(x,yi)/λ∑ie−A(x,yi)/λ∇θlogπθ(yi,j∣xi)] _θJ=- E_x p\! [ λ _ie^-\,A(x,y_i)/λ _ie^-\,A(x,y_i)/λ _θ _θ(y_i,j x_i) ] (E.1) While this estimator converges to the expected value for large G, it is a biased estimator. This term is in fact self-normalized importance sampling (SNIS). SNIS is known to have a bias of order O(1/G)O(1/G) (Owen, 2013, § 9). More explicitly, the gradient term has the form of a ratio estimator: g^=∑i=1Gwiϕi∑i=1Gwi,wi:=e−Ai/λ,ϕi:=∇θlogπθ(yi∣x), g\;=\; _i=1^Gw_i\, _i _i=1^Gw_i, w_i:=e^-A_i/λ,\;\; _i:= _θ _θ(y_i x), which is exactly SNIS for expectations under Q(y∣x)∝πθ(y∣x)e−A/λQ(y x) _θ(y x)e^-A/λ. Even if both ∑iwiϕi _iw_i _i and ∑iwi _iw_i are unbiased for their respective expectations, their ratio is not: by a first-order Taylor (i.e., delta-method) expansion around ([wϕ],[w])( E[wφ], E[w]), one obtains a bias term proportional to Cov(wϕ,w)/([w])2Cov(wφ,w)/( E[w])^2, yielding the standard O(1/G)O(1/G) bias for SNIS. This effect is most pronounced when G is small or when the weights wiw_i are heavy-tailed (which happens for small λ). The jackknife technique. Consider an estimator g^(X) g(X) for the underlying parameter g, computed from a sample size |X|=k|X|=k. Let X−iX_-i denote the sample vector in which the i-th sample is deleted, and g^(X−i) g(X_-i) denote the estimator in the absence of the i-th sample. Then, the Jackknife estimator is: g~(X)=kg^(X)−k−1k∑ig^(X−i) g(X)=k g(X)- k-1k _i g(X_-i) It can be easily seen that if the bias of the original estimator is eg^=c/k+O(1/k2)e_ g=c/k+O(1/k^2), then the jackknife technique reduces this to eg~=O(1/k2)e_ g=O(1/k^2) (Miller, 1974; Jiao and Han, 2020). The jackknife targets this bias by using the leave-one-out terms g^(X−j) g(X_-j) to estimate the leading 1/k1/k term in the Taylor expansion of the bias, and subtracting it via the linear combination in g~(X) g(X). As a result, the O(1/k)O(1/k) term cancels and the remaining bias is O(1/k2)O(1/k^2). Appendix F Additional Experiments F.1 High learning-rate SFT without LoRA Degrades the Capabilities In this paper we posit that the KL constraint is implicitly satisfied in SFT when LoRA is deployed, or when the learning rate and training duration are moderate; we verified this empirically in Figure˜3 (right). Here, we show that in a high-learning-rate regime, even though safety degrades faster, general capabilities—measured by the general reward model “Llama-3.1-8B-Instruct-RM-RB2” (Malik et al., 2025)—deteriorate significantly as well. We repeat the SFT on Alpaca experiment for Mistral model as described in Section˜5.1.1, but with lr=1e−5lr=1e-5 rather than lr=1e−6lr=1e-6. Figure˜11 (left) shows that the higher learning rate results in a slightly lower safety reward. This corresponds to a slightly higher KL with the base model in Figure˜11 (right), a level that the lower learning rate would eventually reach given more training steps. However, as Figure˜11 (middle) shows, the helpfulness score is highly compromised. This indicates that higher learning rates are more prone to overfitting in general, suggesting that our KL constraint remains applicable as long as the fine-tuning scheme avoids overfitting. Figure 11: The model trained with FRPO and λ=0.5λ=0.5 fine-tuned on Alpaca; (left/middle) we consider a higher choice of learning rate for SFT on Alpaca (described in Section˜5.1.1) to show that it slightly degrades the safety but the helpfulness score is highly impacted. (right) This shows that a higher learning rate still changes the KL controllably but with higher slope, confirming our constraint in Section˜3. F.2 Evaluations on Alpaca It is important to demonstrate that the safety robustness conferred by FRPO does not come at the expense of downstream performance. Figure˜3 (left) shows that the SFT loss of the FRPO-trained model aligns closely with that of GRPO and Llama-3.1. In contrast, models such as Llama3-TAR (Tamirisa et al., 2024) and Llama3-derta (Yuan et al., 2025) exhibit an initial jump in loss. These models are adversarially trained to resist adaptation, preventing them from fitting the downstream data as effectively as models trained with other methods. We also track the helpfulness reward during SFT on Alpaca in Figure˜12 (right). The results indicate that the general capabilities of the models, as measured by helpfulness score, follow similar trajectories. Notably, FRPO maintains a higher reward than GRPO by the end of fine-tuning. This suggests that the superior safety score is not an artifact of overfitting. Figure 12: (left) We measure the SFT loss during fine-tuning to show that the FRPO-trained model fits the downstream task as well as other models. (right) FRPO also keeps the general capabilities higher that GRPO after fine-tuning. F.3 Full Comparison on GSM8K We compare the Qwen model trained with FRPO and λ=0.2λ=0.2 against other baselines on HarmBench and StrongREJECT after fine-tuning on GSM8K (we already compared with Qwen and Mistral models in Section˜5.1.2). As Figure˜13 shows, our model outperforms others except for Llama-3-R (Zou et al., 2024); this model has effectively unlearned unsafe responses and is not directly comparable to our robust RLHF method. Nevertheless, we show in Section˜5.1.1 that the model trained with FRPO outperforms Llama-3-R after fine-tuning on Alpaca. It must be noted that the drop in the StrongREJECT score of Llama-3-TAR is due to a significant drop in the helpfulness score (down to ∼0.1 0.1). As mentioned earlier in Section˜F.2, this model strictly resists fine-tuning, even if it comes at the cost of losing its general capabilities. In contrast, we show in Table˜1 that our models maintain their general capabilities and improve on GSM8K. Figure 13: The refusal rate and the StrongREJECT score of the models during fine-tuning on GSM8K. Llama-3-TAR loses its general capabilities after fine-tuning, and thus the harmfulness score drops as well.