Language Models are Homer Simpson! Safety Re-Alignment of Fine-tuned Language Models through Task Arithmetic

TL;DR

This work tackles the problem that fine-tuning LLMs often degrades safety. It introduces RESTA, a simple arithmetic realignment that adds a task-safe vector to post-finetuned weights, and augments it with DARE to prune redundant deltas, restoring safety with minimal loss in task performance. The approach is validated across PEFT and Full-FT on multilingual and code/math tasks, using CatQA and three existing safety benchmarks, with GPT-4 acting as the judge. Key findings show substantial reductions in unsafety scores (e.g., from 18.6% to 5.1% under PEFT and 9.2% to 1.5% under Full-FT) while preserving most task capabilities, and results generalize across languages and domains. The work contributes a practical, scalable safety re-alignment technique and a new multilingual safety dataset, enabling safer deployment of fine-tuned LLMs in diverse settings.

Abstract

Aligned language models face a significant limitation as their fine-tuning often results in compromised safety. To tackle this, we propose a simple method RESTA that performs LLM safety realignment. RESTA stands for REstoring Safety through Task Arithmetic. At its core, it involves a simple arithmetic addition of a safety vector to the weights of the compromised model. We demonstrate the effectiveness of RESTA in both parameter-efficient and full fine-tuning, covering a wide range of downstream tasks, including instruction following in Chinese, English, and Hindi, as well as problem-solving capabilities in Code and Math. We also showcase the generalizability of RESTA on three existing safety evaluation benchmarks and a multilingual benchmark dataset proposed as a part of this work, consisting of 550 harmful questions covering 11 categories, each with 5 sub-categories of harm. Overall, RESTA decreases the harmfulness of the compromised model from 18.6% to 5.1% and from 9.2% to 1.5% in parameter-efficient and full fine-tuning, respectively, while maintaining most of the model's performance on the task. We release the source codes at: https://github.com/declare-lab/resta.

Figures (4)

• Figure 1: Illustration of restored safety in (PEFT) fine-tuned models using Resta---At little to no performance trade-off and without any further alignment cost, Resta brings back the safety guardrails by simple addition of the safety vector to the SFT model.
• Figure 2: (Methodology) LLMs can be likened to the Homer Simpson character, known for overlooking consequences in decision-making. LLMs, in pursuit of task performance, may compromise safety, and thus deviate from an ideal optimization path. The inherent safety of LLM is symbolized by a safety cap, often lost during training (SFT). Resta employs straightforward arithmetic operations to effectively restore the cap on the Simpson character without incurring additional alignment costs.
• Figure 3: Categorical harmfulness (unsafety) score of the Uncompromised (base) model, SFT (PEFT/Full-FT), and Resta variants. As judged by GPT-4, the harmfulness of the model increases after SFT (green) which is significantly reduced by Resta (red and yellow). Base model (dark blue) and SFT refer to the Llama-2-7B model and its domain-specific tuned versions, respectively. Resta and Resta$_d$ denote SFT model after the added safety vector with and without DARE, respectively.
• Figure 4: Impact on the safety and performance of the model with different fractions of added safety vector. SFT is on Chinese-Alpaca, evaluated on CatQA.