Secure LLM Fine-Tuning via Safety-Aware Probing

TL;DR

This paper introduces a safety-aware probing framework that improves the safety-utility balance during LLM fine-tuning by steering updates away from unsafe regions while maintaining task performance.

Contribution

It proposes a novel contrastive safety signal-based optimization method that enhances LLM safety during fine-tuning without sacrificing utility.

Findings

SAP reduces harmful scores significantly compared to standard fine-tuning.

SAP outperforms strong baselines in safety-utility tradeoff.

SAP demonstrates robustness against harmful data poisoning and adversarial attacks.

Abstract

Large language models (LLMs) have achieved remarkable success across many applications, but their ability to generate harmful content raises serious safety concerns. Although safety alignment techniques are often applied during pre-training or post-training, recent studies show that subsequent fine-tuning on adversarial or even benign data can still compromise model safety. In this paper, we revisit the fundamental question of why fine-tuning on non-harmful data may nevertheless degrade safety. We show that the safety and task-performance loss landscapes are partially decoupled, so updates that improve task-specific performance may still move the model toward unsafe regions. Based on this insight, we propose a safety-aware probing (SAP) optimization framework for mitigating safety risks during fine-tuning. Concretely, SAP uses contrastive safety signals to locate safety-correlated directions, and optimizes a lightweight probe that perturbs hidden-state propagation during fine-tuning, thereby steering parameter updates away from harmful trajectories while preserving task-specific learning. Extensive experiments show that SAP consistently improves the safety--utility tradeoff across multiple models and tasks. Averaged over multiple LLMs, SAP reduces the harmful score significantly relative to standard fine-tuning, outperforming strong baselines while maintaining competitive task-specific performance. SAP also demonstrates stronger robustness under harmful data poisoning, adversarial fine-tuning, and a dedicated post-fine-tuning adaptive attack, validating that SAP is an effective and scalable framework for preserving LLM safety during fine-tuning. Our code is available at https://github.com/ChengcanWu/SAP.