PaperlayerView exact revision on arXiv

Source-linked AI summary

Safety Alignment as Continual Learning: Mitigating the Alignment Tax via Orthogonal Gradient Projection

Guanglong Sun, Siyuan Zhang, Liyuan Wang, Jun Zhu, Hang Su, Yi Zhong

arXiv:2602.07892v2cs.LGcs.CL

TL;DR

Safety post-training can improve alignment while reducing general utility, motivating a continual-learning view of gradient interference. OGPSA projects safety gradients away from a low-rank capability subspace and consistently improves the empirical safety–utility trade-off across alignment pipelines.

  • Problem

    The paper studies whether gradient interference during heterogeneous sequential alignment contributes to the alignment tax and capability regression in LLMs.

  • Method

    OGPSA estimates a low-rank general-capability gradient subspace and projects safety gradients onto its orthogonal complement during SFT or DPO.

  • Results

    OGPSA consistently improves the empirical safety–utility trade-off over standard baselines across SFT, DPO, and sequential SFT→DPO pipelines.

  • Takeaways & Limitations

    The results support explicit gradient-interference control as a promising lightweight approach for mitigating capability regression during safety post-training.

  • Takeaways & Limitations

    OGPSA does not guarantee robustness to all jailbreaks, misuse prompts, or deployment distributions and should complement broader safety evaluation and oversight.

Abstract

from arXiv · show

Safety post-training can improve the harmfulness and policy compliance of Large Language Models (LLMs), but it may also reduce general utility, a phenomenon often described as the \emph{alignment tax}. We study this trade-off through the lens of continual learning: sequential alignment stages expose the model to shifted data distributions and objectives, and their gradients may interfere with directions that support previously acquired general capabilities. This view does not claim that all alignment degradation has a single cause; rather, it provides a useful first-order mechanism for mitigating one important source of capability regression. We propose \textbf{O}rthogonal \textbf{G}radient \textbf{P}rojection for \textbf{S}afety \textbf{A}lignment (\textbf{OGPSA}), a lightweight update rule that estimates a low-rank reference subspace from gradients on a small set of general-capability data and removes from each safety gradient the component lying in this subspace. The resulting update is the steepest local safety-descent direction subject to first-order preservation constraints on the reference objectives. OGPSA is compatible with standard post-training pipelines and avoids large-scale replay, although it introduces periodic reference-gradient computation. Across Supervised Fine-Tuning (SFT), Direct Preference Optimization (DPO), and sequential SFT$\rightarrow$DPO settings, OGPSA improves the observed safety--utility trade-off over standard baselines. Under the sequential SFT$\rightarrow$DPO pipeline, the average performance gain increases from 33.98\% to 42.74\% on Qwen2.5-7B-Instruct and from 19.74\% to 32.98\% on Llama3.1-8B-Instruct. We have open sourced our code at https://github.com/SunGL001/OGPSA.

1 Introduction

The introduction frames the alignment tax as partly arising from gradient interference during objective-heterogeneous continual learning and presents OGPSA as a first-order constrained remedy. OGPSA projects safety updates away from a low-rank general-capability gradient subspace and improves the empirical safety–utility trade-off across alignment settings.

  • Motivation: LLMs exhibit strong general capabilities but may produce toxic, biased, misleading, or harmful outputs without explicit alignment.The motivation spans complex reasoning, code generation, open-ended synthesis, and safety risks.
  • Problem: Sequential SFT and preference optimization reduce harmful behaviors but frequently incur an alignment tax through general-capability regressions.The introduction identifies safety post-training as a sequential optimization pipeline after pre-training.
  • Problem: The paper interprets part of the alignment tax as gradient interference in objective-heterogeneous continual learning across stages such as SFT followed by preference optimization.This frames interference as a concrete, testable mechanism rather than claiming a single cause for all alignment degradation.
  • Method: OGPSA estimates a low-rank general-capability gradient subspace from representative general data and projects safety updates onto its orthogonal complement.The method is presented as a lightweight, plug-and-play rule with a first-order feasible-descent characterization.
  • Evaluation: OGPSA consistently improves the empirical safety–utility trade-off over standard baselines across model families and alignment strategies.The evaluation covers multiple alignment strategies while also reporting limitations of the first-order approximation.

2 Related Work

Related work studies LLM safety through test-time interventions and safety-focused post-training, with the latter potentially degrading general capabilities. OGPSA is positioned as a projection-based continual-learning method tailored to heterogeneous alignment objectives and broad capability preservation.

  • LLM Safety Alignment: Safety research uses test-time guards or model steering, but these methods add inference latency and system complexity.Guards identify unsafe responses, while steering adjusts the output distribution.
  • LLM Safety Alignment: Safety-data training can degrade general capabilities, motivating methods that seek to mitigate this regression.The passage contrasts safety post-training with the resulting loss of general capability.
  • Positioning relative to gradient-projection CL: Unlike GEM and GPM, OGPSA targets objective heterogeneity across SFT, DPO, and SFT→DPO while preserving broad LLM capabilities.Its contribution is framed around tailored formulation, subspace construction, and validation for safety alignment rather than the projection operator itself.

3 Preliminaries

This section defines the alignment tax as the evaluation-level utility loss after safety post-training and models its mitigation as preserving general capabilities under sequential objective and distribution shifts. It introduces first-order reference-loss constraints that motivate projecting safety gradients away from a low-rank general-capability subspace.

  • Alignment tax: ∆tax = Φ(θpre; Deval) − Φ(θsafe; Deval) defines the alignment tax as the evaluation-level utility loss after safety alignment.Directly constraining Φ during training is difficult because it may be non-differentiable or expensive.
  • Heterogeneous continual learning: Sequential safety stages such as SFT and DPO shift data distributions and objectives, allowing safety gradients to interfere with capability-supporting directions.This continual-learning-style interference can improve safety behavior while degrading general utility.
  • Capability subspace: The method estimates a low-rank general-capability gradient subspace from small datasets covering facets such as reasoning, coding, and truthfulness.The low-dimensional approximation is motivated by evidence that fine-tuning often operates in effective low-dimensional subspaces.
  • Reference constraints: A sufficient first-order preservation condition is ⟨g(i)(θ), ∆θ⟩ = 0 for every reference capability i.The constraints require the update to have zero inner product with each reference gradient.
  • Projected update: The resulting update lies in Sgen(θ)⊥ by removing from each safety update its component in the local general-capability reference subspace.The next section operationalizes this principle with a low-rank basis and gradient projection.

4 Methodology

OGPSA mitigates first-order interference between safety optimization and selected general-capability objectives by projecting safety gradients onto the orthogonal complement of a periodically refreshed, low-rank reference subspace. This yields the steepest local safety-descent direction under first-order preservation constraints, with periodic reference-gradient computation as the main added overhead.

  • OGPSA framework: OGPSA estimates a low-rank subspace from general-capability gradients and projects each safety gradient onto its orthogonal complement before updating parameters.The reference datasets are small and diverse, targeting facets such as reasoning, coding, truthfulness, helpfulness, and truthfulness.
  • Projected update rule: The projected update θt+1 ← θt −η ˜gsafe lies in Sgen(θ)⊥ up to refresh lag and stochastic gradient noise, reducing first-order interference with selected reference directions.The projected gradient is obtained by removing the safety-gradient component in the current general-capability subspace.
  • Algorithm and complexity: The overhead is M additional reference-gradient computations every K steps plus projection onto a rank-M′ subspace; it is substantially smaller than large-scale replay for small M′ and moderate K, but nonzero.The projection requires M′ inner products and a linear combination, so wall-clock time and token counts should still be reported.
  • Dynamic subspace construction: Every K steps, OGPSA recomputes reference gradients on mini-batches and constructs an orthonormal basis of rank M′ ≤ M, filtering redundant directions with threshold δ.The basis uses Gram–Schmidt, and the rank may be smaller than the number of reference objectives because linearly dependent directions are removed.
  • First-order preservation: Under constraints ⟨g(i)(θ), ∆θ⟩ = 0 for all reference objectives, the projected gradient is the steepest instantaneous safety-descent direction.This is a local first-order statement about the chosen reference losses, not a global guarantee that all downstream capability metrics remain preserved.

5 Experiments

Experiments across two instruct models and three alignment paradigms show that OGPSA improves the observed safety–utility trade-off over standard and established mitigation baselines. Ablations further indicate that subspace diversity, data-efficient reference gradients, periodic updates, and model-scale evaluation affect its performance and practicality.

  • Experimental setup: OGPSA is evaluated on LLaMA3.1-8B-Instruct and Qwen2.5-7B-Instruct across SFT, DPO, and sequential SFT-DPO.The study uses a 10k-sample PKU-SafeRLHF-derived dataset and compares against naive tuning, +Merged, +LoRA, and +General Data.
  • Standard baselines: Standard alignment improves safety metrics but often reduces general utility, while existing mitigations only partially recover the alignment tax.General-data replay and parameter averaging can dilute the safety signal in some settings.
  • Main results: OGPSA obtains the highest Avg. Gain across both models and all three alignment stages in this evaluation.In sequential SFT-DPO, Avg. Gain rises from 33.98% to 42.74% on Qwen2.5-7B-Instruct and from 19.74% to 32.98% on Llama3.1-8B-Instruct.
  • Jailbreak resistance: OGPSA reduces Attack Success Rate relative to corresponding SFT and DPO baselines under I-GCG while keeping optimization difficulty comparable or higher.This stress test does not show an immediate vulnerability to the evaluated optimization-based jailbreak.
  • Ablations: Single-domain reference subspaces target different capabilities: UltraFeedback (Helpful) yields HHH 88.74%, whereas HaluEval (Truthful) restores SimpleQA 3.17%.The passage reports limited truthfulness protection with UltraFeedback, at SimpleQA 1.94%, and lower IFEval and HHH with HaluEval.
  • Efficiency and optimization: OGPSA needs only 100–200 samples per reference dimension to recover several capability metrics more effectively than 10K-sample replay, with periodic updates outperforming no updating.The experiments re-estimate the subspace every 5 steps in DPO and 30 steps in SFT.
  • Scalability: Across Qwen2.5 models from 0.5B to 7B parameters, OGPSA generally improves retained general-capability metrics under SFT and DPO while maintaining competitive safety scores.On Qwen2.5-0.5B, SimpleQA after SFT increases from 0.09% to 0.88%, and Stereotype increases from 79.31% to 96.74%.

6 Conclusion and Limitations

The paper frames the alignment tax as heterogeneous continual learning and proposes OGPSA, a lightweight gradient-projection method that mitigates capability regression while improving the safety–utility trade-off across common post-training pipelines.

  • The work frames the alignment tax as a heterogeneous continual learning problem.
  • OGPSA mitigates capability regression by updating within the orthogonal complement of a learned low-rank capability subspace.
  • OGPSA is a simple, plug-and-play method for SFT, DPO, and sequential SFT→DPO pipelines that consistently improves the safety–utility trade-off.

7 Impact Statement

OGPSA is intended to improve the safety–utility trade-off of LLM post-training by reducing unnecessary capability degradation in benign settings. However, higher post-training utility does not guarantee robustness to jailbreaks, misuse prompts, or deployment distributions.

  • Potential positive impact: OGPSA aims to reduce unnecessary capability degradation from safety alignment, potentially improving aligned systems’ usefulness in benign settings.This is identified as a potential positive impact of the work.
  • Potential risk: Improved utility after safety training could be mistaken for comprehensive safety.The paper explicitly distinguishes utility gains from complete safety guarantees.
  • Limitations: OGPSA does not guarantee robustness to all jailbreaks, misuse prompts, or deployment distributions.The limitation concerns safety behavior beyond the settings evaluated by the method.

8 Large language model assistance

Large language models assisted with manuscript polishing, while the authors reviewed and edited the content and retained responsibility for the final paper.

  • Large language models were used to polish the manuscript.
  • The authors reviewed and edited the content.
  • The authors take responsibility for the final paper.

A Detailed Experimental Setups · B Pseudo-code

The experiments evaluate OGPSA on two instruction-tuned LLMs using safety, truthfulness, and utility data, standard alignment baselines, and a ten-benchmark evaluation suite. The pseudo-code periodically constructs an orthonormal reference-gradient subspace and projects safety gradients away from it before updating the model.

  • A Detailed Experimental Setups: Experiments use LLaMA3.1-8B-Instruct and Qwen2.5-7B-Instruct, with a 10k-sample PKU-SafeRLHF safety dataset.The safety labels are refusal responses generated by gpt-4omini.
  • A Detailed Experimental Setups: Truthfulness data use correct answers as SFT and DPO chosen labels and hallucinated answers as DPO rejected labels.Hallucinated answers are reformatted to match correct-answer formatting, highlighting data efficiency relative to replay-based methods.
  • A Detailed Experimental Setups: The study compares SFT, DPO, and sequential SFT-DPO against OGPSA augmented with regularization and replay strategies.The listed comparison strategies include +Merged, described as linear interpolation between pretrained and safety-aligned parameters.
  • A Detailed Experimental Setups: Training uses 3 SFT epochs, 1 DPO epoch, learning rate 1e-6, DPO β 0.2, batch size 128, zero weight decay, and cosine scheduling with 0.1 warm-up.All LLM training uses LLaMA-Factory.
  • A Detailed Experimental Setups: The evaluation suite contains 10 benchmarks covering safety and general utility.Safety evaluation includes StrongReject, XSTest, toxic WildChat, and stereotype Do-Not-Answer splits.
  • B Pseudo-code: OGPSA takes pretrained parameters, a safety loss, reference datasets, refresh period K, and learning rate η, and returns aligned parameters.The algorithm initializes an empty basis U and iterates for T optimization steps.
  • B Pseudo-code: Every K steps, OGPSA samples reference batches, computes reference gradients, and applies Gram-Schmidt to construct an orthonormal basis U.This is labeled dynamic subspace construction in the pseudo-code.
  • B Pseudo-code: OGPSA removes the safety-gradient component in U using ˜gsafe = gsafe − U(U^⊤gsafe), then updates parameters with learning rate η.The pseudo-code labels this stage projected safety optimization and identifies the subtraction as removing conflicting components.

C Theoretical Derivations · C.1 Steepest Descent Direction in a Linear Subspace

The section establishes that, for differentiable objectives, steepest descent constrained to a linear subspace is the negative orthogonal projection of the gradient. Applied to OGPSA, this result identifies the update as the optimal safety-descent direction within the non-forgetting subspace.

  • C Theoretical Derivations: The section provides the mathematical foundation for Proposition 1 by proving the projected-gradient characterization for any differentiable function.The constrained steepest-descent direction equals the negative gradient projected onto the chosen linear subspace.
  • C.1 Steepest Descent Direction in a Linear Subspace: The formal setup models the LLM parameter space as a d-dimensional Euclidean space V = R^d with a differentiable loss f and gradient g = ∇f(θ).A linear subspace S defines allowable optimization directions, with P_S denoting orthogonal projection onto S.
  • C.1 Steepest Descent Direction in a Linear Subspace: For unit vectors v ∈ S, the objective is to minimize the directional derivative, provided P_S(g) ≠ 0.This defines the steepest descent problem restricted to the subspace.
  • C.1 Steepest Descent Direction in a Linear Subspace: Thus, the steepest descent direction within a subspace is the negative of the orthogonally projected gradient.This is the central conclusion of Lemma C.1.
  • C.1 Steepest Descent Direction in a Linear Subspace: Orthogonal decomposition separates the gradient into g_S = P_S(g) within S and g_⊥ ∈ S^⊥, with g_⊥ orthogonal to every vector in S.Therefore, the orthogonal component contributes zero to directional derivatives along feasible directions.
  • C.1 Steepest Descent Direction in a Linear Subspace: The constrained objective reduces to minimizing ⟨g_S, v⟩ over unit vectors in S, whose lower bound is −∥g_S∥.Cauchy–Schwarz shows that maximum descent occurs when v is collinear with g_S and points oppositely.
  • C.1 Steepest Descent Direction in a Linear Subspace: In OGPSA, S is the null space of general capabilities S_⊥gen and f is the safety loss L_safe.The lemma establishes that the OGPSA update follows the optimal safety-optimization path within the non-forgetting zone.

D Appendix Results … D.5 Impact of Subspace Update Frequency using SFT

The appendix evaluates OGPSA across overall Llama performance, subspace composition, mathematical generalization, reference-set swapping, gradient-estimation budgets, and subspace-update frequency. These analyses examine robustness and optimization behavior across SFT and DPO settings.

  • D.1 Overall Performance of Lllama: Overall performance is assessed on Llama3.1-8B-Instruct using aggregate Safety Score and General Capacity Score metrics.Safety Score averages four datasets, while General Capacity Score averages six datasets.
  • D.2 Impact of Subspace Dimensionality and Diversity using SFT: The Qwen SFT study tests how Helpfulness-versus-Truthfulness reference-data diversity and one-versus-two-dimensional constraint subspaces affect alignment outcomes.The comparison reports the best results in bold.
  • D.3 Generalization to the Mathematical Domain: Mathematical generalization is evaluated by comparing Qwen2.5-7B-Instruct performance on mathematical benchmarks across alignment categories.Best performance within each alignment category is marked in bold.
  • D.5 Impact of Subspace Update Frequency using SFT: Reference-set robustness is tested by replacing the SFT reference sets with 200 GSM8K samples from a math-only dataset.The experiment evaluates whether alignment outcomes remain robust under this complete reference-set swap.
  • D.4 Impact of Sample Size Budgets for Gradient Estimation: DPO gradient estimation is evaluated as the reference-gradient sample budget increases, with the method remaining effective under limited data budgets.The study measures performance stability across sample sizes on Qwen.
  • D.5 Impact of Subspace Update Frequency using SFT: Qwen SFT optimization dynamics are compared for static subspaces and dynamic subspace updates at varying intervals.The comparison identifies the best results in bold.

NeurIPS Paper Checklist

The checklist reports that the paper’s claims, assumptions, proofs, reproducibility information, code and data access, and experimental details are adequately disclosed. It also identifies missing error bars as a statistical-significance limitation.

  • Claims: The abstract and introduction accurately reflect the paper’s contributions and scope.The checklist also emphasizes that claims should match the theoretical and experimental results and state important assumptions and limitations.
  • Theoretical analysis: The paper provides the full set of assumptions and complete, correct proofs for each theoretical result.The checklist recommends numbering and cross-referencing theorems, formulas, and proofs and stating assumptions clearly.
  • Reproducibility: The paper fully discloses the information needed to reproduce its main experimental results.The checklist separately states that the paper provides sufficient information on the computer resources required for reproduction.
  • Open access: The paper provides open access to the data and code with sufficient reproduction instructions.This checklist item concerns faithful reproduction of the main experimental results, including the supplemental material.
  • Experimental setting: The paper specifies the training and test details necessary to understand its results.Examples listed include data splits, hyperparameters, hyperparameter selection, and optimizer type.
  • Statistical significance: The paper does not report error bars for its results.The checklist notes that appropriate error bars, confidence intervals, or statistical significance tests should accompany experiments supporting the main claims.
Loading 2602.07892v2…