Large language models fine-tuned to keep safety and perform better
Rethinking LLM Fine-Tuning via Weight Space Reparameterization: Preserving Safety during Downstream Adaptation
Machine LearningArtificial Intelligence
Summary
Fine-tuning large language models often makes them worse at staying safe while trying to do new tasks. The authors suggest a new way called WSR-Tune that changes how the model's settings are organized so the safety parts don't get messed up when the model learns new things. Their method finds the important safety parts and keeps them fixed, while allowing the rest to change for better performance. This helps keep the model safe and improved at the same time.
What this means in practice
- •For machine learning engineers: Improve fine-tuning techniques to maintain model safety while adapting to new tasks using weight space reparameterization.
- •For ai safety teams: Preserve safety behaviors in large language models during updates for downstream applications to reduce risk of unsafe outputs.
Authors
Minseong Kim, Jongbok Won, Jeesup Park, Yonghee Choi, Dongjun Han
Abstract
Fine-tuning large language models (LLMs) on downstream tasks often weakens their previously aligned safety behavior, creating an inherent trade-off between task adaptation and safety preservation. Existing approaches attempt to mitigate this issue by identifying safety-relevant neurons, layers, or update directions in the original parameter space. However, safety-related information is often distributed across many directions in this space, causing downstream updates to overwrite parameters that encode safety behavior. We propose WSR-Tune, a Weight Space Reparameterization-based fine-tuning framework that preserves safety during downstream adaptation. Our approach constructs a safety-conditioned basis and reparameterizes the weight matrices such that safety-relevant information becomes concentrated in a small subset of directions. We then identify safety-critical directions and freeze them in the reparameterized space, while updating the remaining complementary directions for downstream adaptation. By explicitly structuring the parameter space to disentangle safety and task information, our method reduces interference between objectives and enables simultaneous preservation of safety and improvement in downstream performance. Experimental results show that our method achieves a more favorable trade-off between downstream performance and safety retention, demonstrating its effectiveness for reliable LLM fine-tuning.