푸시-풀 분포 정렬을 통한 LLM 미세 조정 안전성 확보

Safeguarding LLM Fine-tuning via Push-Pull Distributional Alignment

대규모 언어 모델(LLM)의 고유한 안전 정렬은 미세 조정 과정에서, 심지어는 겉보기에는 무해한 데이터셋을 사용하는 경우에도 손상되기 쉽습니다. 기존의 안전성 확보 방법들은 데이터 선택을 통해 이러한 문제를 완화하려고 시도하지만, 일반적으로 휴리스틱한, 개별 데이터 수준의 평가에 의존하며, 데이터 분포의 전체적인 구조를 간과하고 유해한 패턴을 명시적으로 제거하지 못합니다. 이러한 문제를 해결하기 위해, 본 연구에서는 안전 최적 수송(Safety Optimal Transport, SOT)이라는 새로운 프레임워크를 제안합니다. SOT는 안전한 미세 조정을 개별 데이터 필터링의 문제로 접근하는 대신, 최적 수송(Optimal Transport, OT)을 기반으로 한 분포 수준의 정렬 문제로 재구성합니다. SOT의 핵심은 이중 참조 ``푸시-풀'' 가중치 학습 메커니즘입니다. SOT는 샘플의 중요도를 최적화하면서, 신뢰할 수 있는 안전한 기준점(safe anchor)을 향해 분포를 적극적으로 끌어당기는 동시에, 일반적인 유해한 기준점(harmful reference)으로부터 분포를 밀어내는 방식으로 작동합니다. 이를 통해 견고한 기하학적 안전 경계를 구축하여 훈련 데이터를 효과적으로 정제합니다. 다양한 모델 아키텍처와 도메인에서의 광범위한 실험 결과, SOT는 모델의 안전성을 크게 향상시키면서도 경쟁력 있는 성능을 유지하며, 기존 방법들에 비해 우수한 안전-유용성 균형을 달성한다는 것을 보여줍니다.

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.