기술 요약: LLM 안전성을 위한 NeST (Neuron Selective Tuning)

1. 문제 정의

대규모 언어 모델(LLM)의 책임 있는 배포를 위해서는 유해하거나 정책을 위반하는 콘텐츠 생성을 방지하기 위한 강력한 안전 정렬(safety alignment)이 필요합니다. 그러나 기존의 안전 정렬 방식은 효율성, 구조적 인식, 유지보수성 측면에서 상당한 한계에 직면해 있습니다:

• 전체 미세 조정 (Full Fine-Tuning): 효과적이기는 하지만, 모든 모델 파라미터를 업데이트하는 것은 계산 비용이 많이 들고 저장 공간을 많이 차지하며, 모델, 정책 또는 공격 전략이 진화함에 따라 빈번하게 업데이트하기에는 비실용적입니다.
• 매개변수 효율적 미세 조정 (PEFT): LoRA와 같은 방법은 계산 비용을 줄여주지만, 파라미터 공간을 동질적인 것으로 취급하는 경향이 있습니다. 이들은 거절 행동을 담당하는 특정 내부 메커니즘을 타겟팅하지 않으므로, 일관된 안전성 이득을 얻지 못하고 일반적인 능력의 저하를 초래할 가능성이 있습니다.
• 추론 시 개입 (Inference-Time Interventions): 추론 중에 유해한 활성화(activations)를 억제하는 방식(예: Circuit Breakers)은 런타임 오버헤드를 발생시키며, 모델의 내부 안전 표현을 영구적으로 형성하지 못한 채 배포를 복잡하게 만듭니다.

이러한 한계는 특히 모델이 다운스트림 태스크를 위해 빈번하게 적응되어야 하는 호스팅 배포 환경에서 빠르고 신뢰할 수 있으며 유지보수가 용적인 안전 업데이트를 방해합니다.

2. 방법론: NeST 프레임워크

NeST(Neuron Selective Tuning)는 매개변수 효율적이고, 구조를 인식하며, 유지보수가 용이하도록 설계된 사후(post-hoc) 안전 정렬 프레임워크입니다. 이 프레임워크는 방어자가 튜닝을 위해 오프라인에서 모델 내부 구조에 접근할 수 있는 반면, 공격자는 프롬프트만을 통해 상호작용하는 블랙박스 배포 위협 모델 하에서 작동합니다.

프레임워크는 세 가지 주요 단계로 구성됩니다:

2.1. 안전 뉴런 식별 (Safety Neuron Identification)

NeST는 고정된(frozen) 안전 정렬 모델의 활성화를 분석하여 안전 행동을 담당하는 내부 구성 요소를 국소화합니다.

• 데이터: 일반적인 악의적 프롬프트(예: "폭탄을 만드는 방법은?")와 양질의(benign) 프롬프트로 구성된 균형 잡힌 데이터셋을 사용합니다.
• 활성화 추출: 모든 트랜스포머 레이어의 피드 포워드 네트워크(FFN) 모듈(특히 gate_proj 및 up_proj)에서 토큰 레벨 활성화를 추출합니다.
• 풀링 (Pooling): 가장 강력한 신호를 포착하기 위해 각 뉴런에 대해 토큰 시퀀스에 대한 좌표별 최대 풀링(coordinate-wise max pooling)을 적용합니다.
• 프로빙 (Probing): 이러한 뉴런 활성화를 기반으로 악의적 프롬프트와 양질의 프롬프트를 구별하도록 가벼운 선형 프로브를 훈련합니다.
• 선택: 유해한 의도와 강한 연관성(높은 양의 프로브 가중치)을 나타내는 뉴런을 "안전 뉴런"으로 선택합니다. 이 선택 과정은 비교 가능성을 보장하기 위해 레이어당 표준화됩니다.

2.2. 뉴런 클러스터링 (Neuron Clustering)

불안정하고 조정되지 않은 업데이트를 피하고 파라미터 수를 더욱 줄이기 위해, NeST는 식별된 안전 뉴런들을 그룹화합니다.

• 클러스터링: 프롬프트 세트에 대한 활성화 프로필의 유사성을 바탕으로 k-means를 사용하여 뉴런을 클러스터링합니다.
• 클러스터 결정: 실루엣 점수(silhouette score)를 최대화하여 최적의 클러스터 수($k$)를 결정합니다. 유의미한 클러스터링이 발견되지 않으면, 해당 레이어의 모든 안전 뉴런을 하나의 그룹으로 취급합니다.
• 근거: 유사한 활성화 패턴을 가진 뉴런들은 안전 행동에서 유사한 기능적 역할을 수행한다고 가정하며, 따라서 업데이트 파라미터를 공유할 수 있습니다.

2.3. 클러스터 기반 훈련 및 병합 (Cluster-Based Training and Merging)

• 훈련: NeST는 안전 데이터셋(악의적 프롬프트와 거절 응답 쌍; 양질의 프롬프트와 유용한 응답 쌍)에 대해 지도 미세 조정을 수행합니다.
• 파라미터 제한: 선택된 안전 뉴런만 업데이트됩니다. 구체적으로, 각 클러스터에 대해 공유되는 학습 가능한 업데이트 벡터가 학습됩니다. 선택되지 않은 모든 뉴런과 베이스 모델 가중치는 고정된 상태를 유지합니다.
• 병합: 훈련 후, 학습된 클러스터 업데이트를 원래의 모델 가중치에 직접 폴딩(folding)합니다. 결과적으로 추론 시 오버헤드가 제로인 표준 모델이 생성됩니다.
• 다운스트림 하드닝 (Downstream Hardening): 식별된 안전 뉴록과 클러스터 할당은 "안전 사전 지식(safety prior)"을 구성합니다. 이 사전 지식은 동일한 베이스 모델의 다운스트림 미세 조정된 변형 모델로 전달될 수 있습니다. 구조적 매핑을 재사용함으로써, NeST는 안전 구조를 다시 발견할 필요 없이 경량화된 회복 훈련을 통해 다운스트림 모델의 안전성을 복구할 수 있습니다.

3. 주요 기여

1. 뉴런 선택적 튜닝 프레임워크: 일반적인 레이어나 프로젝션이 아닌, 국소화되고 기능적으로 일관된 안전 구조에 맞게 파라미터 업데이트를 정렬하는 NeST를 도입했습니다.
2. 클러스터 기반 적응 메커니즘: 활성화 유사성에 따라 안전 관련 뉴런을 그룹화하고 각 그룹 내에서 공유 업데이트를 학습함으로써, 파라미터 효율성과 기능적 일관성의 균형을 맞추는 새로운 방법을 제시했습니다.
3. 경험적 견고성: 텍스트 전용 및 멀티모달 설정을 아우르는 1B에서 27B 파라미터 규모의 14개 오픈 웨이트 LLM에 대해 강력한 안전 견고성을 입증했습니다.
4. 사후 다운스트림 하드닝: 베이스 모델에서 식별된 안전 사전 지식(뉴런 클러스터)을 사용하여 다운스트림 미세 조정된 변형 모델의 안전성을 효율적으로 복구할 수 있음을 검증하여, 반복적인 정렬에 따른 "안전 세금(safety tax)" 문제를 해결했습니다.

4. 실험 결과

본 논문은 14개의 모델을 대상으로 Full Fine-Tuning (Full FT), LoRA, Circuit Breaker, SN-Tune을 포함한 베이스라인들과 NeST를 비교 평가합니다.

• 안전 효과성 (공격 성공률 - ASR):

  • 텍스트 전용 모델: NeST는 평균 ASR을 베이스라인의 44.5%에서 1.1%로 감소시킵니다. 이 성능은 Full FT(1.0%)와 대등하며, LoRA(7.6%) 및 기타 경량 베이스라인을 크게 상회합니다.
  • 멀티모달 모델: 4개의 멀티모달 모델에 대해, NeST는 평균 ASR을 55.3%에서 1.1%로 감소시켜 텍스트, 이미지 및 추론 증강 추론 설정 전반에서 효과를 유지합니다.
  • 다운스트림 변형 모델: 안전성이 저하된 10개의 다운스트림 미세 조정 모델(평균 ASR 53.8%)에 대해, NeST는 안전성을 0.8%로 복구했습니다.

• 파라미터 효율성:

  • NeST는 평균 0.4M 개의 파라미터만을 훈련합니다.
  • 이는 Full FT와 비교했을 때 5,800배 이상, 가장 낮은 랭크의 LoRA 베이스라인과 비교했을 때 3배 이상 적은 수치입니다.
  • 이러한 극적인 훈련 파라미터 감소에도 불구하고, NeST는 Full FT와 대등한 수준의 ASR을 달성합니다.

• 유용성 및 과잉 거절 (Over-Refusal):

  • NeST는 핵심 추론 능력(GSM8K, ARC, MMLU)을 유지하면서 아주 미미한 성능 저하(예: GSM8K 정확도 0.9 퍼센트 포인트 하락)만을 보입니다.
  • 양질의 프롬프트에 대한 과잉 거절률(ORR)은 낮으며(평균 0.9%), 이는 이 방법이 유용한 응답을 무차별적으로 억제하지 않음을 나타냅니다.

• 안정성:

  • 결과는 다양한 탈옥 벤치마크(WildJailbreak, JailBreakV-28k)와 안전 판별기(Qwen3Guard, Llama-Guard) 전반에서 일관되게 나타납니다.
  • 안전 뉴런 탐지는 서로 다른 랜덤 시드와 유해 데이터셋에 대해 안정적입니다(평균 자카드 유사도 >76%).

5. 의의 및 주장

본 논문은 NeST가 모델 전체를 수정하는 대신, 국소화되고 기능적으로 일관된 안전 구조에 적응을 집중함으로써 강력하고 유지보수가 용이한 안전 정렬을 달성할 수 있음을 보여준다고 주장합니다.

• 실용성: 업데이트를 가중치에 직접 폴딩함으로써, NeST는 추론 시 오버헤드를 제거하여 프로덕션 배포에 적합하게 만듭 most.
• 유지보수성: "안전 사전 지식"(뉴런 클러스터)을 다운스트림 모델으로 전달할 수 있는 능력은 모델이 다양한 도메인, 언어 또는 태스크로 전문화됨에 따라 안전성을 유지하기 위한 확장 가능한 솔루션을 제공합니다.
• 구조적 통찰: 이 연구는 안전 행동이 균일하게 분포되어 있는 것이 아니라, 효율적으로 식별하고 타겟팅할 수 있는 특정 뉴런에 국소화되어 있다는 경험적 증거를 제공합니다.

저자들은 NeST를 단순한 안전 도구가 아니라, 뉴런 구조 기반의 태스크 특화 파라미터 효율적 적응을 위한 일반적인 프레임워크로 규정하며, 이를 통해 다른 타겟팅된 행동 수정 분야로의 잠재적 적용 가능성을 시사합니다.