Title: Shape it Up! Restoring LLM Safety during Finetuning
Venue: NeurIPS 2025
Accepted: yes
Keywords: LLM

LLM finetuning

LLM finetuning risks

harmful finetuning

guardrail

safe guard

safety shaping
Abstract: Finetuning large language models (LLMs) enables user-specific customization but introduces important 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 — an atomic treatment we term static safety shaping. In contrast, we propose dynamic safety shaping (DSS), a dynamic shaping 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 ★DSS, a DSS method 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


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces dynamic safety shaping (DSS), a fine-tuning framework for large language models that leverages fine-grained safety signals to reinforce safe content and suppress unsafe segments within model responses. By repurposing guardrail models to evaluate safety at the segment level, the authors propose Safety Trajectory Assessment of Response (STAR), a token-level signal enabling dynamic safety control during training. Experiments show that the proposed method achieves substantial safety gains across models and datasets without sacrificing performance on intended tasks.\n\nSTRENGTHS AND WEAKNESSES:\nPros:\n\n- 1 This paper studies an important problem in LLM safety. It formalises the limitations of “static” sample-level shaping (e.g., RS) and reframes safety mitigation as a sequence-level weighting problem.\n- 2  The method is evaluated across multiple guardrails and benchmark datasets,  outperforming a few baseline methods in Safety metrics.\n- 3  The theoretical bound (Theorem 1) is provided. It ties post-finetuning harmfulness to the expected sequence-level KL divergence. \n\nCons:\n- 1 The key idea of this paper is similar to Deep token, making the contribution of this paper a bit of limited. \n- 2 The proposed method is heavily dependent on guardrail quality. Safety still degrades almost linearly with the guard-rail FN rate, even though the decline is gentler than RS, DSS cannot handle entirely novel jailbreak styles unseen by the guard-rail.\n- 3 Code and data release are “upon acceptance”; check-points and guardrail configs are not yet public, which may slow reproducibility.\n- 4 Only short-answer benchmarks (MMLU, ARC-C, GSM-8K) are used; no tests on generation-heavy tasks (long-form reasoning, coding) where the KL term might over-regularise helpful content.\n\nQUESTIONS:\nSee cons.\n\nLIMITATIONS:\nNo. The paper needs to discuss the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes Dynamic Safety Shaping (DSS), a novel framework for maintaining the safety of large language models (LLMs) during fine-tuning. Traditional static methods like rejection sampling treat entire training examples as safe or unsafe, overlooking harmful content embedded within otherwise benign data. To address this, the authors introduce the STAR (Safety Trajectory Assessment of Response) score—a token-level signal derived from guardrail models that estimates the evolving safety risk throughout a response. Leveraging this signal, they design STAR-DSS, a dynamic loss function that adaptively balances imitation learning and safety regularization at the token level. The paper provides both theoretical guarantees and extensive empirical evaluations, demonstrating that STAR-DSS successfully maintains model safety across various threats and models without sacrificing task performance.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper proposes a dynamic loss function that **adjusts the loss across different tokens/chunks using STAR scores**. This fine-grained shaping significantly reduces the risk of safety degradation during finetuning compared to static methods.\n2. The experiments span multiple LLMs, datasets, harm levels, and guardrails. This comprehensive setup strengthens the generalizability and credibility of the results.\n3. STAR-DSS consistently outperforms baselines in both worst-case and ideal-case settings, achieving high safety without sacrificing task performance.\n\nWeaknesses:\n1. The interplay between STAR granularity (chunk size M) and computational cost vs. safety-performance trade-off is acknowledged but not deeply optimized or ablated.\n2. More discussion on limitations is needed, such as time cost (**total** time of fine-tuning using STAR-DSS vs. baseline), requirements on the Guardrail model (intuitively, the Guardrail model should usually be of similar size to the fine-tuned model to ensure adequate safety).\n\nQUESTIONS:\n1. The proposed method relies heavily on guardrail accuracy. Can the authors provide a more detailed analysis (e.g., ablation or sensitivity study) on how STAR-DSS behaves when guardrail false negatives increase or shift under adversarial adaptation?\n2. As shown in Figure 4(b), could you provide some analysis on Llama Guard-3 1B is better than 8B?\n\nLIMITATIONS:\nDiscussed in weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper is about a novel loss for fine-tuning models via SFT, that take into account the safety of the completions at the token level and thus limit the degradation of model safety performance. The authors place their work specifically in the context of fine-tuning platforms which are available online. The authors introduce the STAR score: given a chunk of text and a safety classifier, the STAR score measures the probability that it’s safe (via the ratio of the logit for “safe” and that for “safe” + that for “unsafe”). The novel loss - DSS for dynamic safety shaping - is a weighted cross-entropy loss with the STAR score,  balanced with the KL divergence to the original model, with a weight of 1- STAR. The authors proceed to demonstrate a theoretical result that bounds the probability for the fine-tuned model to produce unsafe completions with that of the original model, a term linked to the KL divergence of the two models, and the false negative rate of the classifier. Finally, the authors provide experimental results in the following settings\n* “Vanilla” SFT or “malicious” user SFT - meaning the input dataset is 100% unsafe;\n* “Benign” user SFT - the user has a 5% unsafe dataset;\n* “Sophisticated malicious” user SFT - the user has appended misleading prefixes or suffixes to all the completions which are otherwise 100% unsafe, or is conducting a prompt poisoning attack, or trying to do harmful prefilling.\n\nSTRENGTHS AND WEAKNESSES:\n**Major comments**\n\nIt’s great that the authors are showing results w/ a few safety classifiers, input checkpoints, fine-tuning datasets. However, experimental results are not convincing for a few reasonsM\n* There are no standard deviations - e.g. over eval re-runs (the easier). Safety results are noisy, and in the case at hand, results from the proposed method are close to that of rejection sampling, which is likely to be much faster and more compute-efficient.\n* In relation to this, no data is presented regarding time and cost of a run w/ DSS. It would be interesting to see how that changes w/ M (chunk size). What's the tradeoff exactly if someone decided to use the DSS rather than rejection sampling?\n* I don’t find the “vanilla” SFT scenario very credible. Providers have access to safety classifiers too, and it’s much simpler to plan for pausing someone’s account once a certain threshold of harmful content is passed, to investigate, rather than force them to use another loss, that’d additionally incur the user or them more cost.\n* There should be more references to other literature mentioning prefix/suffix appending type of attacks, as otherwise they seem a little anecdotal (“response adaptation attack” l. 355). More generally, there should be more coverage of the background. It seems from most of the paper that benign fine-tuning needs to include unsafe examples to deteriorate the safety performance (benign user SFT scenario), but it’s not the case. It would be interesting to see how the DSS fares in that case (see minor comment below for a suggestion of where to add it).\n\nOverall, I find that the most convincing application of the authors’ method is to defend against data poisoning attacks and harmful prefilling attacks, which are mentioned in section 6.4. In that case, yes the model fine-tuner themselves are looking for solutions to mitigate against someone else’s attacks. It’s likely that depending on the user's application they’d be willing to incur more cost / longer training time to make sure that even if their data was poisoned, backdoors could not be activated. I think it could worth it to focus experimental results more on these aspects in a future version of the paper.\n\n**Minor comments**\n* Example in figure 1 is not super convincing. “Here is” is not essentially safe, and “creating scripts for exploiting vulnerabilities” isn’t essentially unsafe either. I’d suggest to find a few examples that are more illustrative and display them. \n* The baseline is missing in table 2 - what are the results if the user fine-tunes on their corrupted dataset? What about if they fine-tune on GSM8K alone? The baseline would also be a great addition to figure 4.c, keeping the same colour for it over all subplots in Figure 4 and explaining it in the figure caption (its color changing from blue in 4.a to grey in 4.d).\n* It’s surprising that the authors are presenting results that involve fine-tuning on one dataset (eg PureBad), and safety results on other ones (HexPhi and AdvBench). Often, one wants to fine-tune on the same data distribution that they want to apply the model to.\n* Fig. 5b is hard to read.\n\nQUESTIONS:\nSee strengths and weaknesses above\n\nLIMITATIONS:\nThe paper is missing a study of when DSS fails (besides the safety classifier performing poorly). In which conditions does the method break? Are there attacks that it cannot help fight against?", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses safety risks in large language models (LLMs), where even a few harmful examples can undermine safety alignment. Traditional mitigation, termed static safety shaping, updates models uniformly on safe and unsafe parts of responses, which is suboptimal due to shifting safety contexts within a single example. The authors propose dynamic safety shaping (DSS), a framework that uses fine-grained safety signals to reinforce learning from safe response segments while suppressing unsafe content. They introduce the Safety Trajectory Assessment of Response (STAR), a token-level safety signal derived by repurposing guardrail models to evaluate partial responses and track safety risks dynamically. Their method, DSS, guided by STAR scores, robustly mitigates finetuning risks, significantly improving safety across diverse threats, datasets, and model families without compromising task performance. The paper advocates for future research to build on dynamic shaping principles to address evolving safety challenges in LLM finetuning.\n\nSTRENGTHS AND WEAKNESSES:\n## Strength\n- **Novelty in Token-Level Safety Optimization:** The paper introduces dynamic safety shaping (DSS) and Safety Trajectory Assessment of Response (STAR), offering a novel way to address safety risks in LLM finetuning by dynamically adjusting updates based on fine-grained safety signals, a significant improvement over static methods.\n- **Technically Sound:** The proposed DSS framework is grounded in a robust technical foundation, repurposing guardrail models to evaluate partial responses and generate token-level STAR scores. This approach provides a clear and practical mechanism for dynamically tracking and mitigating safety risks during finetuning, ensuring both theoretical rigor and implementability.\n- **Experimentally Effective:** The DSS method demonstrates strong empirical performance, achieving significant safety improvements across diverse threats, datasets, and model families without compromising task performance. The broad applicability and consistent results underscore the method’s effectiveness in real-world scenarios.\n- **Informative Insight Analyze:** The analysis is very informative. For instance, authors find that conventional SFT strategies are sensitive to guardrail errors, and vulnerable to context entanglement. This part is quite informative.\n\n## Weakness\nI think this paper is well motivated and technically sound, my question is mainly lies in the details of the implementation.\n- **Lack of Detailed Explanation of STAR Chunk Accuracy:** One of the risks is how to cut the sequence into different types of chunks for STAR training efficiently and effectively. How is the robustness of the method towards the FP and FN chunks?\n\nQUESTIONS:\nSee weankess\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-05"}