Fine-tuning language models to avoid rare safety failures
Beyond Average Safety: Chance-Constrained LLM Fine-tuning
Machine LearningArtificial Intelligence
Summary
When improving large language models, sometimes changes make the model less safe on rare but important examples. The authors propose a new way to fine-tune models that controls the chance of these safety issues happening rather than just measuring average safety. They develop a mathematical method that focuses on avoiding cases where safety is significantly worse, and show with extensive tests that their approach works better than existing methods. This means tuning models can be safer by watching out for rare but serious problems.
What this means in practice
- •For machine learning engineers: Safely fine-tune language models to minimize rare but serious safety regressions on specific safety-critical inputs.
- •For ai service operators: Improve reliability of deployed AI assistants by limiting degradation on harmful prompt responses during model updates.
Authors
Taha Entesari, Mahyar Fazlyab
Abstract
Fine-tuning large language models on new objectives can improve helpfulness, instruction following, or domain-specific performance, but it can also induce regressions on safety-critical prompts. Existing safety-preserving fine-tuning methods typically control average safety loss or use weighted auxiliary penalties, which can obscure rare but severe failures. We propose a chance-constrained formulation for safety-preserving fine-tuning that limits the fraction of safety examples whose degradation relative to a reference model exceeds a prescribed threshold. Because the resulting empirical chance constraint contains a discontinuous indicator, we introduce a differentiable majorization of the violation rate, yielding a tractable conservative constraint. We then develop a constraint-aware gradient descent method that treats the majorized constraint as a safe set in parameter space and minimally modifies the fine-tuning direction to preserve feasibility. The resulting update admits a closed form and produces a tail-aware safety correction that emphasizes examples near or above the degradation threshold. We conduct an extensive set of experiments on harmful fine-tuning across three different tasks and three models and show that our approach consistently outperforms the baselines that exist in the literature. These results suggest that safety preservation in LLM fine-tuning is better viewed as a reliability-constrained optimization problem than as average-risk regularization.