NLSR: Neuron-Level Safety Realignment of Large Language Models Against Harmful Fine-Tuning

(AAAI 2025, Accept)

Introduction

사용자가 업로드한 소량의 악성 데이터만으로도 LLM 모델의 미세 조정을 미묘하게 조작하여 모델의 안전 정렬 상태를 손상시킬 수 있음
→ 미세 조정 전후의 safety-critical 가중치의 유사도 차이를 기반으로 LLM의 안전 정렬을 복원하는 train-free 프레임워크인 Neuron-Level Safety Realignment (NLSR) 제안

Methonology

1. Construction of a Safety Reference Model: Super-aligned 모델 $F_{W_e}$를 생성. $W_{weak}$는 초기 미세조정된 모델 처럼 안전성이 약한 LoRA 가중치이며, $W_{strong}$은 안전성이 강하게 정렬된 안전성 특화 모델의 LoRA 가중치임

   $W_{medium} = \alpha W_{strong} + (1 - \alpha) W_{weak}, \quad \alpha \in (0, 1] \quad$

   모델 $F_{W_e}$의 LoRA 가중치인 We를 구하는 수식에서 $W_a$는 $W{strong}$의 유무에 따라 값이 달라짐

   $W_e = \frac{1}{\alpha} W_a - \frac{1 - \alpha}{\alpha} W_0 = (1 + \beta)W_a - \beta W_0 \quad$

   $W_{strong}$이 존재하는 경우 $W_a$는 $W_{medium}$이 되며, 존재하지 않는경우 DPO, RLHF 등의 최적화 기법을 바탕으로 preference-alinged된 모델의 LoRA 가중치를 $W_a$로 함. 이때 $W_{strong}$의 존재 여부 판단은 그 모델의 안전성 정렬 정도를 바탕으로 주관적으로 판단함. $W_{strong}$이 없으면 최소한 preference-alinged된 모델이라도 있어야 됨
   $\beta$만큼 $W_0$의 가중치는 상쇄되고 $W_a$의 가중치가 증폭되면서 $W_e$는 안전성 관련 가중치가 강조되는 행렬을 가짐

2. Recognition of Safety-Critical Neurons: 안전-중요 뉴런을 식별하기 위해 지정된 희소성 비율(sparsity rate) PSR에 따라 LoRA 가중치에 rank reduction을 적용. 원본 출력과 근사 출력 간의 Frobenius norm 차이를 최소화하는 low-rank matrix $\hat{W_j}$를 찾음
   즉, $W_j \hat{X^i_j}$의 특이값 분해(Truncated SVD) $USV^T \approx W_j \hat{X^i_j}$ 를 통해 안전성 관련 주요 가중치 행렬을 추출하고, $\hat{W_j} = UU^TW_j$로써 원본 가중치 행렬 $W_j$를 이 안전성 관련 주요 가중치에 투영함으로써, $\hat{W_j} = \arg \min_{\text{rank}(\hat{W_j}) \le r^*} \|W_j \hat{X^i_j} - \hat{W_j} \hat{X^i_j}\|_F^2 \quad$를 만족하는 $\hat{W_j}$를 구함
   이 때, $r^* = r \times (1 - \text{PSR})$
   이 $\hat{W_j}$를 바탕으로 전체 $N^*$뉴런 중 가중치 절대값 $\text{Top-k} = N^* \times (1 - \text{PSR})$의 뉴런을 안전-중요 가중치라고 함

   $\text{indices} = \text{argsort}(-\vert \hat{W_j} \vert)[:, : \text{Top-k}] \quad$

3. Restoration for Safety-Broken Neurons:

   1. Probability-based Layer Pruning: indices를 바탕으로 안전-중요 가중치의 위치만 남기고 마스킹하는 $M_{j,i^*} = \begin{cases} 1, & \text{if } i^* \in \text{indices} \\ 0, & \text{otherwise} \end{cases}$를 구함
      사용자의 미세조정에 의해 안전성이 훼손된 LoRA 가중치를 $W_{t,j} = B_{t,j} A_{t,j}$라고 하고 $A \in \mathbb{R}^{r \times k}, B \in \mathbb{R}^{d \times r}$일 때, Super-aligned 모델의 LoRA 가중치 $W_e$ 및 $F_{W_t}$ 중 안전-중요 가중치만 남긴 행렬을 각각 $W'{e,j}$와 $W'{t,j}$로 함

      $W^{\prime}_{e,j} = (M^B_j \odot B_j)(M^A_j \odot A_{e,j})$
      $W^{\prime}_{t,j} = (M^B_j \odot B_j)(M^A_j \odot A_{t,j}) \quad$

      이 각각의 안전-중요 가중치만 남은 LoRA 가중치 간의 유사도를 비교하여 $S_j = \frac{\langle W'{e,j}, W'{t,j} \rangle_F}{\|W'_{e,j}\|F \|W'{t,j}\|_F}$ 유사도가 높은 순으로 순위를 매김 $\{r_1, r_2, \dots, r_N\} = \text{rank}(S_1, S_2, \dots, S_N)$
      순위를 바탕으로 Pruning probability 를 $P_j = P_L + \delta \frac{r_j}{N} \quad$를 구성하고 Probability-based Layer Pruning $\gamma_j \sim \text{Bernoulli}(P_j) \quad$를 수행함. 이 때 $P_L$은 기본 레이어 가지치기 확률, $\delta$는 증가 계수, $N$은 총 레이어 수
      

   2. Neuron-Level Correction: 모든 레이어의 가지치기 상태 $\Gamma = \{\gamma_1, \gamma_2, ..., \gamma_N\}$가 주어졌을 때, $F_{W_t}$의 LoRA 가중치는 $W_{t,j}^{\prime\prime}$로 변경됨으로써 훼손된 안전성을 복원함

      $W^{\prime\prime}_{t,j} = \begin{cases} W_{e,j}^{\prime} + \hat{W}_{t,j}^{\prime} & \text{if } \gamma_j = 0 \\ W_{t,j}^{\prime} & \text{otherwise} \end{cases}$

   이 때 $\hat{W}'{t,j} = ((1 - M^B_j) \odot B{t,j})((1 - M^A_j) \odot A_{t,j})$

Result

Table 1: Fine-tuning performance on SST2 with Llama3-8B, varying harmful instruction ratios from 0.01 to 0.3. Methods with require no extra training
Table 2: Fine-tuning performance under different alignment methods, including SFT, DPO, ORPO, KTO, and SimPO. Methods marked with indicates that no additional training