Safeguarding LLM Fine-tuning via Push-Pull Distributional Alignment

TL;DR

Safety Optimal Transport (SOT) reframes LLM fine-tuning safety as a distribution-level problem solvable with entropy-regularized OT. It learns a per-sample weight distribution P(\mathbf{w}) by a dual-reference push-pull objective that pulls toward a task-specific safe anchor $Q$ and pushes away from a broad harmful reference $M$, optimizing $\min_{\mathbf{w}} (1-\lambda) \text{OT}(P(\mathbf{w}),Q) - \lambda \text{OT}(P(\mathbf{w}),M)$. The downstream fine-tuning uses a Top-$K$ hard filtering plus soft reweighting loss, $\mathcal{L}_{\text{custom}}$ with reweighted samples, and LoRA for efficiency. Empirically, SOT consistently improves the safety-utility trade-off across diverse model families and domains, achieving state-of-the-art Avg scores and lower Harmfulness Scores, including strong performance in the legal-domain Generalization task. The approach provides a principled, data-driven path to safer, personalized LLM deployment by shaping the training distribution rather than relying on pointwise data filtering alone.

Abstract

The inherent safety alignment of Large Language Models (LLMs) is prone to erosion during fine-tuning, even when using seemingly innocuous datasets. While existing defenses attempt to mitigate this via data selection, they typically rely on heuristic, instance-level assessments that neglect the global geometry of the data distribution and fail to explicitly repel harmful patterns. To address this, we introduce Safety Optimal Transport (SOT), 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). At its core is a dual-reference ``push-pull'' weight-learning mechanism: SOT optimizes sample importance by actively pulling the downstream distribution towards a trusted safe anchor while simultaneously pushing it away from a general harmful reference. This establishes a robust geometric safety boundary that effectively purifies the training data. Extensive experiments across diverse model families and domains demonstrate that SOT significantly enhances model safety while maintaining competitive downstream performance, achieving a superior safety-utility trade-off compared to baselines.

Figures (6)

• Figure 1: Overview of the SOT framework. Stage 1 (Dual-Reference Distributional Weight Learning) utilizes a dual-reference "push-pull" Optimal Transport strategy to learn sample weights by aligning with safe anchors and repelling harmful ones. Stage 2 (Safety-Aware Fine-Tuning) applies these weights for Top-k selection and weighted SFT to produce a final safe and helpful LLM.
• Figure 2: The left figure shows results for different data selection rates, the middle figure presents results for fine-tuned datasets containing varying proportions of harmful samples, the right figure shows the reference data efficiency of SOT using Meta-Llama-3.1-8B-Instruct on the SLIMORCA dataset.
• Figure 3: Distribution of Learned Weights using Meta-Llama-3.1-8B-Instruct on the SLIMORCA dataset.
• Figure 4: Comparative Effectiveness of SOT Method and Baseline Method in Harmful Data Filtering using Meta-Llama-3.1-8B-Instruct on the SLIMORCA dataset.
• Figure 5: The left figure displays experimental results for different $\lambda$ values in Eq.\ref{['eq:sot_loss']}. The right figure shows experimental results for constructing distributions using different layer representations using Meta-Llama-3.1-8B-Instruct on the SLIMORCA dataset.