TRACE: Trajectory-Based Safety Patch Learning for LLM Post-Training Realignment

• URL: http://arxiv.org/abs/2607.16242v1
• Date: Fri, 26 Jun 2026 03:29:59 GMT
• Title: TRACE: Trajectory-Based Safety Patch Learning for LLM Post-Training Realignment
• Abstract summary: Fine-Tuning-as-a-Service platforms let users train large language models on customized tasks.<n>In practice, service providers need to recover models' safety without re-running full alignment.<n>We propose TRACE, a trajectory-based safety patch learning framework.
• Score: 27.81709722999727
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Fine-Tuning-as-a-Service (FTaaS) platforms let users train large language models (LLMs) on customized tasks, but this pipeline could erode models' safety alignment. In practice, service providers need to recover models' safety without re-running full alignment, or destroying the utility gained from customized tasks. A line of existing work refers to model parameter merging, which adds a safety patch on the fine-tuned model parameters to shift the model away from unsafe tendencies. However, this merging-based paradigm is fundamentally bottlenecked by task-safety update entanglement: downstream task updates and the safety patch often overlap in their dominant directions, so the merge strength is intrinsically hard to calibrate. If the safety vector is scaled too weakly, harmful components could still dominate, preventing the model from returning to a safe region; if it is scaled too aggressively, it suppresses task-relevant directions and degrades utility. To solve this problem, we shift the focus of merging-based methods from designing online merging operators to offline patch learning, and seek a safety patch that minimally interferes with task-relevant directions while retaining decisive control over unsafe behaviors. We propose TRACE, a trajectory-based safety patch learning framework that (i) simulates harmful tuning trajectories to generate progressively corrupted states, and (ii) optimizes a plug-in patch to recover safety while maintaining utility across varying corrupted base states. Across six benchmarks and two models, TRACE consistently dominates the safety-utility frontier. TRACE reaches nearly 100% safety on all settings, while maintaining comparable utility to the undefended fine-tuned model.

Related papers

• Do Models Share Safety Representations? Cross-Model Steering for Safe Visual Generation [52.122731171289665]
  We introduce the first framework for cross-model safety steering.<n>Our pipeline never accesses unsafe data on the target side.<n>We evaluate our approach in text-to-image and text-to-video generation across diverse source-target model pairs.
  arXiv  Detail & Related papers  (2026-06-03T18:00:04Z)
• SafeGene: Reusable Adapters for Transferable Safety Alignment [31.247158830484604]
  We propose SafeGene, a reusable safety-adapter module for cross-task reuse.<n>We show that SafeGene-enhanced models reduce harmful response rates while maintaining downstream performance.
  arXiv  Detail & Related papers  (2026-06-02T14:51:14Z)
• RefusalGuard: Geometry-Preserving Fine-Tuning for Safety in LLMs [5.100622189286672]
  Fine-tuning safety-aligned language models for downstream tasks can lead to substantial degradation of refusal behavior.<n>We introduce REFUSALGUARD, a representation-level fine-tuning framework that preserves safety-relevant structure during model adaptation.
  arXiv  Detail & Related papers  (2026-05-03T14:48:18Z)
• Safety Alignment as Continual Learning: Mitigating the Alignment Tax via Orthogonal Gradient Projection [52.551864761088574]
  Large Language Models (LLMs) often incur an alignment tax: safety post-training can reduce general utility.<n>We argue that this tax primarily arises from continual-learning-style forgetting in sequential alignment.<n>We propose Orthogonal Gradient Projection for Safety Alignment (OGPSA) to balance plasticity and stability.
  arXiv  Detail & Related papers  (2026-02-08T09:53:46Z)
• Q-realign: Piggybacking Realignment on Quantization for Safe and Efficient LLM Deployment [55.14890249389052]
  Existing defenses either embed safety recovery into fine-tuning or rely on fine-tuning-derived priors for post-hoc correction.<n>We propose textttQ-realign, a post-hoc defense method based on post-training quantization.<n>Our work provides a practical, turnkey solution for safety-aware deployment.
  arXiv  Detail & Related papers  (2026-01-13T00:07:24Z)
• Safeguarding LLM Fine-tuning via Push-Pull Distributional Alignment [45.772620696660034]
  Safety Optimal Transport (SOT) is a novel framework that reframes safe fine-tuning from an instance-level filtering challenge to a distribution-level alignment task grounded in Optimal Transport (OT)<n>SOT prioritizes sample importance by actively pulling the downstream distribution towards a trusted safe anchor while simultaneously pushing it away from a general harmful reference.<n>Experiments across diverse model families and domains demonstrate that SOT significantly enhances model safety while maintaining competitive downstream performance.
  arXiv  Detail & Related papers  (2026-01-12T04:48:02Z)
• UpSafe$^\circ$C: Upcycling for Controllable Safety in Large Language Models [67.91151588917396]
  Large Language Models (LLMs) have achieved remarkable progress across a wide range of tasks, but remain vulnerable to safety risks such as harmful content generation and jailbreak attacks.<n>We propose UpSafe$circ$C, a unified framework for enhancing LLM safety through safety-aware upcycling.<n>Our results highlight a new direction for LLM safety: moving from static alignment toward dynamic, modular, and inference-aware control.
  arXiv  Detail & Related papers  (2025-10-02T16:43:33Z)
• Shape it Up! Restoring LLM Safety during Finetuning [65.75757313781104]
  Finetuning large language models (LLMs) enables user-specific customization but introduces critical safety risks.<n>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.<n>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.
  arXiv  Detail & Related papers  (2025-05-22T18:05:16Z)
• What Makes and Breaks Safety Fine-tuning? A Mechanistic Study [64.9691741899956]
  Safety fine-tuning helps align Large Language Models (LLMs) with human preferences for their safe deployment. We design a synthetic data generation framework that captures salient aspects of an unsafe input. Using this, we investigate three well-known safety fine-tuning methods.
  arXiv  Detail & Related papers  (2024-07-14T16:12:57Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.