Work

Safeguarding LLM Fine-tuning via Push-Pull Distributional Alignment

Year: 2026

Type: conference-paper

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 dualreference "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.Our code is available at https://github.com/kasaer/ SOT.

Source: Proceedings of the conference - Association for Computational Linguistics. Meeting

DOI: 10.18653/v1/2026.acl-long.1083

Authors: Haozhong Wang, Zhuo Li, Yibo Yang, He Zhao, Hongyuan Zha, Dandan Guo