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

Fine-tuning large language models (LLMs) via Fine-Tuning-as-a-Service (FTaaS) platforms can erode inherent safety alignment, necessitating post-training recovery mechanisms that restore safety without destroying task utility. The proposed TRACE framework addresses this by shifting the focus from online merging operations to offline patch learning, aiming to create a safety patch that is disentangled from user task updates.

The gist

TRACE proposes a trajectory-based framework that simulates harmful tuning trajectories to generate progressively corrupted model states and then optimizes a universal plug-in patch offline to recover safety while strictly preserving task utility across varying corrupted base states.

Problem Formulation and Bottleneck

Existing merging-based methods suffer from task-safety update entanglement, where the directions of the safety patch and user task updates often overlap, leading to mutual interference. This entanglement makes it difficult to calibrate a single repair strength, as weak repairs fail to eliminate insecure behavior while aggressive repairs degrade utility. The core dilemma is that a fixed repair strength cannot optimally balance safety and utility across varying downstream fine-tuning intensities.

TRACE Framework Components

TRACE operates in an offline training stage followed by a calibration-free online serving stage, achieving calibration-free online service. The framework consists of two primary components:

1. Fine-tuning trajectory simulation: This simulates a degradation trajectory, denoted as Tsim, by iteratively computing model states where each subsequent state is updated based on the loss function combining harmful and benign data: Lsim(θ; Bh, Bt) = LCE(θ; Bh) + LCE(θ; Bt). This captures the variation of different training intensities.

2. Decisive safety patch optimization: TRACE optimizes a universal safety patch, denoted as "ϕ, across the resulting family of corrupted model states using the objective function: Lrec(θ ⊕ ϕ; Bref, Bt) = LCE(θ ⊕ ϕ; Bref) + ωLCE(θ ⊕ ϕ; Bt). This optimization encourages the patch to be disentangled from task-relevant updates and decisive over harmful shifts."

Key Properties and Optimization

The optimization process is executed via an alternating simulate-and-learn paradigm, involving Phase A (simulate safety-eroding fine-tuning) and Phase B (learn the disentangled, decisive safety patch). The resulting learned patch ϕ exhibits two key properties:

1. Disentanglement: The learned safety patch operates along directions largely orthogonal to user task updates, which is confirmed by a reduction in saliency overlap between the learned patch and task vectors to below 0.079.

2. Decisiveness: The patch is designed to overpower the harmful vectors because its disentanglement prevents it from canceling user’s harmful drift directly, allowing safety recovery even when harmful updates coexist with benign ones.

Experimental Validation and Results

Extensive experiments across two models (Llama and Qwen) and three out-of-distribution harmful datasets alongside three utility benchmarks (dialogue summarization, SQL generation, mathematical solving) validate TRACE's superiority. TRACE consistently achieves at least 94% safety rate across all combinations of benchmarks and models while maintaining task utility within ±1.7% of the No Defense baseline. Notably, on the most challenging benchmark, TRACE improves the safety rate from 23% to 100%, over four times higher than the second-best value, while preserving downstream task performance with negligible deviation from the undefended baseline (within ±1.7%).

Conclusion and Impact

TRACE demonstrates that post-training realignment can be formulated as an offline safety transfer problem, decoupling the provider’s safety investment from the per-user deployment pipeline. The framework successfully resolves the safety-utility dilemma by learning a universal low-rank adapter, enabling amortized safety maintenance at scale. Furthermore, TRACE can be seamlessly integrated with existing merging operations, boosting their performance by substituting standard patches with its superior learned patch. TRACE achieves the fastest online deployment among all methods, requiring only 0.41 seconds.

Ablation and Robustness

The necessity of trajectory simulation is confirmed by ablation studies; removing it causes safety to collapse to near-zero levels comparable to an undefended model, demonstrating that the progressive degradation captured by the full trajectory is essential for capturing decisive safety directions. The framework also shows high robustness, maintaining perfect safety across all fine-tuning depths (5–30 epochs) with no degradation, and preserving utility within ±1.6% of No Defense. TRACE breaks existing baseline Pareto frontiers by establishing a superior performance regime that requires no per-user coefficient tuning.

Connection to Meta-Learning

TRACE’s alternating optimization structure is compared to meta-learning methods like MAML, sharing the motivation of learning parameters that generalize across varying adapted model states.

Here are the specific improvements that can be made to AI systems based on the TRACE framework, and what those improved systems will be able to do:

) 1. Implementation of a Universal, Zero-Shot Safety Patch Transfer Mechanism:

The system will no longer require per-user calibration or model-specific safety fine-tuning during deployment. Instead, the provider pre-trains a universal safety patch (the TRACE adapter, optimized offline across simulated corruption trajectories).

2. Decisive Safety Recovery Across Unseen Fine-Tuning Intensities:

The improved system will maintain near 100% safety performance regardless of how heavily or aggressively a malicious user fine-tunes the model (simulated by varying training epochs/learning rates). It can neutralize harmful drift even when the downstream task updates are highly divergent from the base model's original alignment, overcoming the task-safety update entanglement bottleneck.

3. Guaranteed Utility Preservation During Malicious Exploits:

The system will achieve a safety rate of at least 94% (and up to 100% in specific settings) while maintaining task accuracy within a negligible deviation (±1.7%) of the undefended baseline, even when encountering complex, mixed-profile training data (harmful content alongside benign task data). This means the model can be used for customized tasks without sacrificing its core functional capability or becoming unsafe when exposed to adversarial input.

4. Efficient and Scalable Realignment:

The system shifts the computational burden from expensive online per-user calibration to a one-time offline training phase for the safety patch. Online deployment latency is reduced dramatically (e.g., from 6 seconds for RESTA down to 0.41 seconds for TRACE), enabling high-throughput, scalable Fine-Tuning-as-a-Service (FTaaS) platforms without incurring prohibitive recurring computational overhead.

5. Robustness Against Evolving Adversarial Attacks:

By learning a patch that is disentangled from task updates, the system becomes resilient to adaptive attacks. The safety mechanism operates on directions largely orthogonal to user task modifications, meaning the safety patch cannot be easily neutralized or overwritten by malicious weight shifts, ensuring decisive control over harmful behaviors.

6. Enhanced Cross-Model Generalization:

The framework is designed to learn a universal patch (using LoRA adapters) that generalizes across different base models (e.g., Llama and Qwen). This allows a single safety patch to provide superior defense mechanisms across an entire family of LLMs, simplifying deployment and maintenance for service providers.

This improved AI system will be:

• A provider-grade LLM deployment pipeline capable of securely hosting user-customized models.

• Able to process data from malicious users while strictly enforcing safety guardrails, even when those users apply aggressive fine-tuning techniques.

• Capable of reliably performing a wide variety of downstream tasks (summarization, SQL generation, math reasoning) with consistent accuracy and safety across all fine-tuning scenarios.

• Highly efficient in operation, as the necessary safety realignment is done once offline rather than repeatedly online.

Fine-Tuning-as-a-Service (FTaaS) platforms let users perform supervised fine-tuning (SFT) on customized data, but this pipeline can erode model safety alignment. To recover safety without re-running full alignment, existing realignment methods focus on calibrating the integration of safety patches into fine-tuned models. These methods exhibit a persistent safety-utility trade-off: weak repair leaves harmful behavior intact, while stronger repair increasingly damages the benign task. This paper shifts the focus from online calibration to offline patch learning and aims to learn a safety patch that restores safety while preserving task-specific capabilities. To this end, we propose TRACE, which simulates harmful SFT trajectories to produce progressively corrupted model states, and optimizes a safety patch simultaneously across these states. TRACE trains the safety patch during the offline stage, and reuses it across all user checkpoints without per-user calibration. We evaluate two representative models using three harmful SFT datasets, together with three utility benchmarks. Across six benchmarks and two models, TRACE consistently dominates the safety-utility frontier. TRACE improves the safety rate by up to 77 percentage points over the best baselines, while maintaining comparable utility to the undefended fine-tuned model.