Localizing Safety Alignment: MLP Layers and Mid-Network Blocks Encode Refusal Behavior in Large Language Models

기술 요약: 대규모 언어 모델(LLM)의 안전 정렬 국소화 (Localizing Safety Alignment in Large Language Models)

문제 제기
대규모 언어 모델(LLM)의 안전 정렬(Safety alignment)은 일반적으로 전체 네트워크의 분산된 속성으로 취급됩니다. 그러나 정렬된 거절 행동(refusal behavior)이 보여주는 실질적인 취약성—즉, 작은 업데이트, 탈옥(jailbreak), 또는 프롬프트 변형만으로도 안전 장치를 우회하거나 무해한 입력에 대해 과잉 거절(over-refusal)을 유발하는 현상—은 거절 메커니즘이 특정된 더 작은 파라미터 집합에 집중되어 있을 가능성을 시사합니다. 기존의 기계론적 연구들이 가지치기(pruning)를 통해 저차원 활성화 방향이나 "안전 레이어"를 식별해 왔으나, 직접적인 가중치 공간 국소화 문제, 즉 *어떤 특정 파라미터가, 어느 깊이에서, 정렬된 모델로부터 미정렬된 베이스 모델로 안전 거절 특성을 전이시키는가?*라는 질문은 여전히 미해결 상태로 남아 있습니다.

방법론
저자들은 동일한 아키텍처를 공유하지만 안전 정렬 여부만 다른 두 쌍의 매칭된 오픈 웨이트 모델을 통해 이 문제를 조사합니다:

1. RealSafe-R1-7B (정렬됨) vs. DeepSeek-R1-Distill-Qwen-7B (베이스)
2. saferlhf ultra sft (정렬됨) vs. Llama-3.1-8B (베이스)

본 연구는 베이스 모델의 가중치 일부를 대응하는 정렬 모델의 가중치로 교체하는 **선택적 가중치 이식(selective weight transplantation)**을 통해 하이브리드 모델을 구축합니다. 실험은 다음과 같은 다양한 이식 정밀도를 적용합니다:

• 컴포넌트 수준(Component-level): 모든 어텐션(Attention) 가중치 교체 vs. 모든 MLP 가중치 교체.
• 연속 영역(Contiguous regions): 처음/마지막/중간 5개 레이어 교체 또는 네트워크의 첫 번째/두 번째 절반 교체.
• 블록 수준(Block-level): MLP 스택을 연속적인 4개 레이어 블록(예: 블록 1 = 레이어 0–3)으로 분할하고, 개별 블록 또는 조합을 이식.

평가는 네 가지 벤치마크인 TwinPrompt 및 SGXSTest (쌍을 이루는 유해/무해 프롬프트), AdvBench (유해 지시문), 그리고 OR-Bench (과잉 거절을 테스트하기 위해 설계된 무해 프롬프트)를 통해 수행됩니다. 주요 지표는 **악성 거절(Malicious Refusal, MR)**과 **무해 과잉 거절(Benign Over-Refusal, BOR)**입니다. 최적의 조합을 식별하기 위해, 저자들은 가장 높은 거절 이득을 주는 MLP 블록을 반복적으로 추가하는 탐욕적 전방 선택(greedy forward selection) 알고리즘을 사용하여 비가산적(non-additive) 효과를 분석합니다.

주요 결과
실험은 두 모델 군 모두에서 세 가지 일관된 경험적 패턴을 도출했습니다:

1. MLP 지배성(MLP Dominance): 전이 가능한 거절 행동은 어텐션 가중치보다는 주로 MLP 가중치에 의해 매개됩니다. 어텐션 파라미터를 교체할 때보다 MLP 파라미터를 교체할 때 악성 프롬프트에 대한 거절 능력이 훨씬 더 많이 회복되었으며, 모든 악성 벤치마크에서 최소 2.7배의 이득을 보였습니다. 그러나 이 경로 또한 과잉 거절 신호를 함께 전달합니다. RealSafe-R1-7B 계열의 경우, 전체 MLP 이식은 MR과 BOR을 모두 유의미하게 증가시켰으며, 이는 바람직한 거절과 오탐(false-positive) 거절이 MLP 경로에 공존함을 나타냅니다.
2. 중간 네트워크 집중(Mid-Network Concentration): 거절 관련 파라미터는 MLP 스택 내에 균등하게 분포되어 있지 않습니다. **블록 3 (레이어 8–11)**은 두 모델과 세 가지 악성 데이터셋에 걸친 6개의 탐욕적 검색 과정에서 일관되게 가장 먼저 선택되었으며, 이는 거절을 매개하는 파라미터가 특정 중간 깊이에 집중되어 있음을 식별해 줍니다. RealSafe-R1-7B 계열에 대한 추가 분석 결과, 레이어 8이 특히 결정적이며, 단일 레이어 이식만으로도 전체 블록의 절반에 가까운 거절률을 달려 성과를 낼 수 있음을 확인했습니다.
3. 비가산적 상호작용(Non-Additive Interactions): 안전 관련 블록의 구성은 **비가산적(non-additive)**입니다. 6개의 탐욕적 궤적 중 5개에서, 특정 지점(종종 6개 블록 부근)을 넘어 더 많은 정렬된 블록을 추가하면 선택적 부분 집합을 사용할 때보다 거절 성능이 오히려 감소했습니다. 이러한 "부의 한계 효용(negative marginal effect)"은 특정 블록의 효용이 이식된 주변 가중치 집합에 의존함을 의미합니다. 결과적으로, 선택된 블록의 부분 집합은 악성 거절과 무해 과-거절 모두에서 전체 MLP 이식보다 더 나은 성능을 보이는 경우가 많습니다.

의의 및 주장
본 논문은 현재 LLM의 안전 정렬이 국소적이며 상호작용에 민감하다고 주장합니다.

• 취약성 설명: 거절 신호가 소수의 MLP 파라미터(특히 중간 네트워크 블록)에 집중되어 있다는 점은 왜 정렬된 행동이 취약한지에 대한 기계론적 설명을 제공합니다. 즉, 이 작은 부분 집합에 대한 표적 업데이트가 안전 행동을 크게 변화시킬 수 있습니다.
• 정밀도 대 커버리지(Precision vs. Coverage): 블록 선택 순서를 도출하는 데 사용된 소스 벤치마크가 다운스트림 행동에 영향을 미친다는 트레이드오프를 강조합니다. 서로 다른 벤치마크에서 도출된 순서는 공격적인 안전 전이(높은 MR)와 과잉 일반화(높은 BOR) 사이의 서로 다른 균형을 만들어냅니다.
• 표적 개입(Targeted Interventions): 본 연구 결과는 효과적인 안전 개입을 위해 전체 정렬 모델을 이식할 필요가 없음을 시사합니다. 대신, 특정하고 비가산적인 MLP 블록의 부분 집합을 식별하고 선택함으로써 더 유리한 운영 지점을 달성할 수 있으며, 이는 더욱 정밀하고 견고한 안전 튜닝을 향한 경로를 제공할 수 있습니다.

저자들은 자신들의 국소화 패턴이 연구된 7B–8B 모델 범위 내에서는 견고하지만, 이러한 특정 깊이 범위(레이어 8–11)가 다양한 모델 규모, 아키텍처 및 정렬 절차에 걸쳐 일반화될 수 있는지 확인하기 위해 향후 연구가 필요하다고 결론짓습니다.