LoRA is All You Need for Safety Alignment of Reasoning LLMs

Reasoning-capable LLMs have achieved major breakthroughs in solving complex problems, but recent work shows that acquiring and deploying strong reasoning can introduce significant safety risks. A common mitigation is to apply a secondary safety-alignment phase after reasoning is learned; however, safety alignment often degrades reasoning performance–a phenomenon known as the “Safety Tax”. In this work, we show that a simple approach can largely bypass this trade-off: applying LoRA during SFT on refusal datasets. Despite its simplicity, this recipe achieves safety comparable to full-model alignment while preserving reasoning performance close to the original reasoning-tuned model, and the result holds across multiple model sizes and architectures, two safety benchmarks, and four reasoning benchmarks spanning mathematics, science, and code generation. We further ablate LoRA configurations and find that (1) rank-1 updates are sufficient to achieve the best safety-reasoning trade-off, (2) applying LoRA only to the MLP up-projection layers can outperform updating the full MLP, and (3) updating middle layers is more effective than updating early or late layers. Finally, we provide a theoretical analysis that helps understand when and why LoRA works, revealing that overshooting the rank budget (using a larger rank than needed for the finetuning task) induces base-task degradation at a rate inversely proportional to the intrinsic dimensionality of the base task. This suggests LoRA is most effective when the finetuning task is low-rank and the base capability is high-rank.

💡 Research Summary

The paper tackles a pressing problem in modern large language models (LLMs): the “Safety Tax” that arises when safety‑alignment is performed after a model has been fine‑tuned for advanced reasoning. Existing pipelines typically first endow a model with strong chain‑of‑thought or other reasoning capabilities and then apply a secondary safety‑alignment phase (often full‑model supervised fine‑tuning (SFT) or reinforcement learning). While this second phase can dramatically improve safety, it usually degrades the model’s reasoning performance, creating an undesirable trade‑off.

The authors hypothesize that safety‑related behavior in LLMs lives in a low‑dimensional subspace of either activation space (e.g., steering vectors, refusal features) or weight space (previous work shows safety‑critical weights occupy a low‑rank subspace). In contrast, full‑model safety alignment changes many directions (high stable rank), unnecessarily perturbing the reasoning circuitry. If safety can be achieved by modifying only a few directions, then restricting updates to a low‑rank subspace should preserve reasoning.

To test this, they adopt Low‑Rank Adaptation (LoRA), a parameter‑efficient fine‑tuning technique that injects trainable low‑rank matrices (ΔW = α·B·A) into frozen pretrained weights. By applying LoRA during SFT on a straightforward “direct‑refusal” safety dataset, they aim to align the model without touching the high‑rank components that support reasoning.

Experimental setup

• Models: Three open‑source reasoning‑capable LLMs: DeepSeek‑R1‑Distill‑Qwen‑7B, DeepSeek‑R1‑Distill‑Qwen‑14B, and DeepSeek‑R1‑Distill‑Llama‑8B (two architectures, three sizes).
• Safety data: DirectRefusal (harmful query → refusal response). Evaluation on StrongREJECT (310 policy‑violating queries) and BeaverTails (14 harm categories).
• Reasoning benchmarks: AIME (math), GPQA (science), HumanEval+ and MBPP+ (code generation). Metric: Pass@1 with 8 samples per query.
• Training regimes: Full‑model SFT for 5 epochs (baseline) vs. LoRA SFT for 10 epochs (default rank = 1, applied only to MLP up‑projection layers). Additional ablations vary rank, target modules, and layer positions.

Key findings

1. Safety vs. reasoning trade‑off eliminated: LoRA achieves safety scores comparable to full‑model alignment (≈ same reduction in harmful responses) while preserving reasoning performance at 98‑100 % of the original reasoning‑tuned model. Full‑model SFT, by contrast, drops Pass@1 by 5‑10 % across all benchmarks.
2. Rank‑1 suffices: Increasing the LoRA rank beyond 1 yields diminishing returns and can even cause “overshoot” – a phenomenon where excess capacity harms the base task. Theoretical analysis shows that overshoot degrades the base task at a rate inversely proportional to the intrinsic dimensionality of that task.
3. Targeted modules matter: Updating only the up‑projection part of the MLP yields a better safety‑reasoning balance than updating the full MLP, the down‑projection, or the gating component. This suggests that modest changes to the output side of the feed‑forward network are enough to encode refusal behavior without disrupting the internal reasoning pathways.
4. Middle layers are most effective: Applying LoRA to a contiguous block of middle transformer layers (e.g., 16‑32 out of 40) is sufficient; early or late layers produce weaker safety gains and larger reasoning drops. This aligns with prior observations that harmful representations tend to emerge in middle layers.
5. Theoretical justification: The authors model the situation as a linear regression with a high‑dimensional “base” task (reasoning) and a low‑dimensional “fine‑tuning” task (safety). Full‑model fine‑tuning can fit the safety objective but at the cost of erasing the base solution. LoRA, constrained to rank r, can fit the safety task only if r matches the task’s intrinsic dimension; otherwise, overshoot occurs, but the degradation scales with 1/(intrinsic dimension of base). Hence, when the base capability is high‑dimensional (as with reasoning) and the safety task is low‑rank, LoRA is theoretically optimal. Empirical validation confirms that safety‑fine‑tuning an instruction‑tuned model (where the base task is low‑dimensional) does not benefit from LoRA, matching the theory.

Implications

• Practical workflow: For developers building reasoning‑capable LLMs, a single LoRA‑based safety fine‑tuning step can replace the traditional two‑stage pipeline (reasoning → full‑model safety). This reduces compute cost, memory usage, and training time while preserving the hard‑earned reasoning abilities.
• Guidelines: Use rank = 1, apply LoRA only to MLP up‑projection matrices, and focus on a middle block of transformer layers (≈ 40 % of total depth). This configuration yields the best Pareto frontier between safety and reasoning.
• Theoretical insight: The work provides a clear criterion—low‑rank safety task + high‑rank base task—for when LoRA will succeed. This can inform future research on other alignment problems (e.g., bias mitigation, factuality) where the target behavior may also be low‑dimensional.

Limitations and future work

• The study focuses on refusal‑type safety (i.e., refusing harmful requests). It does not address jailbreak resistance, adversarial prompt engineering, or nuanced ethical dilemmas that may require richer safety representations.
• Experiments are limited to open‑source models up to 14 B parameters; scaling to 70 B+ models may reveal new dynamics.
• The theoretical analysis relies on a linear regression abstraction; extending it to nonlinear transformer dynamics could deepen understanding.

Conclusion
The paper convincingly demonstrates that LoRA, when applied judiciously during safety supervised fine‑tuning, can bypass the notorious “Safety Tax” for reasoning LLMs. By leveraging the low‑rank nature of safety behavior, a rank‑1 LoRA update to MLP up‑projections in middle layers preserves almost all reasoning capability while attaining safety levels on par with full‑model alignment. The work offers both strong empirical evidence across multiple models and benchmarks, and a theoretical framework that explains when and why this approach works. It provides a clear, cost‑effective recipe for practitioners seeking safe, high‑performing reasoning models.