Shape it Up! Restoring LLM Safety during Finetuning

TL;DR

This work tackles safety degradation during LLM finetuning by moving from static, sample-level safety filtering to dynamic, token-level shaping. It introduces STAR, a guardrail-derived trajectory signal that tracks safety across each generated token, and STAR-DSS, a loss that adaptively emphasizes safe segments while constraining unsafe ones. The authors provide a theoretical bound on the harmfulness of the finetuned model and demonstrate substantial safety gains across diverse models, guardrails, datasets, and attack scenarios without sacrificing capability. The results advocate adopting dynamic shaping principles in finetuning safety to guard against evolving risks in real-world settings.

Abstract

Finetuning large language models (LLMs) enables user-specific customization but introduces critical safety risks: even a few harmful examples can compromise safety alignment. A common mitigation strategy is to update the model more strongly on examples deemed safe, while downweighting or excluding those flagged as unsafe. However, because safety context can shift within a single example, updating the model equally on both harmful and harmless parts of a response is suboptimal-a coarse treatment we term static safety shaping. In contrast, we propose dynamic safety shaping (DSS), a framework that uses fine-grained safety signals to reinforce learning from safe segments of a response while suppressing unsafe content. To enable such fine-grained control during finetuning, we introduce a key insight: guardrail models, traditionally used for filtering, can be repurposed to evaluate partial responses, tracking how safety risk evolves throughout the response, segment by segment. This leads to the Safety Trajectory Assessment of Response (STAR), a token-level signal that enables shaping to operate dynamically over the training sequence. Building on this, we present STAR-DSS, guided by STAR scores, that robustly mitigates finetuning risks and delivers substantial safety improvements across diverse threats, datasets, and model families-all without compromising capability on intended tasks. We encourage future safety research to build on dynamic shaping principles for stronger mitigation against evolving finetuning risks. Our code is publicly available at https://github.com/poloclub/star-dss.

Figures (7)

• Figure 1: Dynamic safety signals reveal evolving risks within each training sample, motivating finer-grained mitigation of llm finetuning risks. In the finetuning-as-a-service setting, where users upload data and providers return finetuned llm, the safety context within a single training example can shift across tokens, mixing safe and unsafe content. Treating such examples as atomic and updating on the entire sequence is suboptimal. We propose the star score, a token-level safety signal computed using a guardrail model, that tracks evolving risk across each response, and introduce $\bigstar$DSS, which uses it to suppress unsafe patterns while preserving model capability. The star score shown in the figure is computed using Llama Guard-3-8B.
• Figure 2: Static safety shaping via rs delivers notable safety gains, but its atomic view of each training example creates blind spots. High guardrail fn and complex, mixed-safety content leak into finetuning and degrade the safety of the resulting model.
• Figure 3: The star score enables fine-grained safety assessment within training samples, addressing key limitations of static safety shaping. We plot average star scores as a function of response progression and show that it reliably captures evolving safety risks across different datasets.
• Figure 4: Our $\bigstar$DSS generalizes across (a) llm, (b) guardrails, (c) harm levels, and (d) finetuning datasets, achieving robust safety gains without compromising capability. Orange consistently denotes our $\bigstar$DSS across all subplots. Gray represents the vanilla sft baseline. Blue (in subplot A only) highlights safety degradation from harmful finetuning. Triangle markers denote safety scores, and circle markers denote capability scores.
• Figure 5: Qualitative comparisons of model responses to broader threats in finetuning-as-a-service. We present how different finetuned llm behave under (a) response adaptation, (b) prompt poisoning, and (c) harmful prefilling attacks, demonstrating that $\bigstar$DSS consistently produces safer generations across all cases.

Theorems & Definitions (14)

• Theorem 1
• Definition 1: Harmfulness
• Definition 2: False–negative rates
• Definition 3: Sequence‑level KL divergence
• Lemma 1: Chain--rule equivalence
• Lemma 2: Optimizer control of $\varepsilon_{\mathrm{KL}}$
• proof
• Lemma 3: Pinsker inequality for input prompt
• proof
• Lemma 4: Expected harmfulness gap