Sparse Attention Post-Training for Mechanistic Interpretability

이 논문은 성능 저하 없이 어텐션 연결성을 0.4% 수준으로 극도로 희소화하여 모델의 계산 중복성을 제거하고, 이를 통해 회로 단순화와 해석 가능성을 획기적으로 향상시키는 새로운 후학습 기법을 제안합니다.

핵심

이 논문은 **"거대한 인공지능 (LLM) 이 어떻게 생각하는지 이해하기 어렵다면, 그 뇌를 더 간결하게 정리해 보자"**는 아이디어를 담고 있습니다.

기존의 거대 언어 모델은 너무 복잡해서, "이 모델이 왜 이런 답을 냈는지"를 설명하는 것이 마치 수만 개의 전선이 뒤죽박죽 섞인 거대한 회로판을 보는 것과 같습니다. 연구자들은 이 복잡한 전선들을 정리해서, 핵심만 남기고 나머지는 잘라내도 성능은 그대로 유지하면서, 내부 작동 원리를 명확하게 볼 수 있게 만드는 방법을 개발했습니다.

이 내용을 일상적인 비유로 설명해 드릴게요.

---

1. 문제: "너무 많은 사람이 참여하는 회의"

지금까지의 AI 모델은 수천 명의 직원이 모여서 회의를 하는 것과 같습니다.

• 상황: "오늘 점심 메뉴를 정하자"라고 질문하면, 수천 명의 직원이 모두 입을 열어 의견을 냅니다.
• 문제: 누가 어떤 말을 했는지, 누가 결정에 영향을 줬는지 파악하기가 불가능합니다. 모든 사람이 다 말하니까 소음만 크고, 진짜 핵심 아이디어가 어디에서 나왔는지 찾기 어렵습니다. (이것이 AI 의 '해석 불가능성' 문제입니다.)

2. 해결책: "필요한 사람만 남기는 '스파르타' 회의"

연구자들은 AI 를 다시 훈련시켜서, 정말 필요한 사람만 회의에 참여하게 만들었습니다.

• 방법: AI 가 학습할 때, "너무 많은 사람이 말하면 안 돼. 중요한 사람 1~2 명만 말하게 해"라고 규칙을 세웠습니다. (논문에서는 이를 '희소성 (Sparsity) 규제'라고 부릅니다.)
• 결과: 놀랍게도, 99% 이상의 직원을 회의에서 내보내도 점심 메뉴를 정하는 능력 (성능) 은 전혀 떨어지지 않았습니다.
• 비유: 마치 수천 명의 군중 대신, 3~4 명의 전문가만 모여서 문제를 해결하는 것과 같습니다.

3. 놀라운 발견: "정리된 뇌의 구조"

이렇게 직원을 줄이자 (주의를 집중시키자) AI 의 내부 구조가 완전히 달라졌습니다.

• 전에는: "A 라는 단어를 이해하려면 B, C, D, E... 수천 개의 연결고리가 필요해!"라고 복잡하게 생각했습니다.
• 지금은: "A 라는 단어를 이해하려면 B 라는 사람 하나만 보면 돼."라고 명확하게 생각하게 되었습니다.
• 효과:
  • 회로 단순화: AI 가 특정 문제 (예: "12+34 는 얼마?") 를 풀 때, 필요한 부품 (연결선) 이 100 배 이상 줄어듭니다.
  • 이해 가능성: 이제 연구자들은 "아, 이 AI 는 '12'와 '34'를 더할 때, 오직 이 두 숫자만 연결하는 회로만 사용했구나!"라고 정확하게 추적할 수 있게 되었습니다.

4. 구체적인 예시: "거꾸로 읽기 vs 직관적인 읽기"

논문 속의 작은 실험 (두 자리 수 덧셈) 을 보면 더 명확합니다.

• 기존 모델: 모든 숫자를 서로 뒤죽박죽 연결해서 답을 냅니다. 마치 모든 숫자를 섞어서 뭉개는 것처럼 복잡합니다.
• 새로운 모델: "일단 1 자리수를 더하고, 10 자리수를 더하고, 필요하면 '올림'만 처리한다"는 명확한 알고리즘을 스스로 찾아냈습니다. 마치 사람이 손으로 덧셈을 할 때처럼 단계별로 깔끔하게 계산하는 것입니다.

5. 왜 이것이 중요한가요?

이 기술은 AI 를 **블랙박스 (검은 상자)**에서 투명한 유리상자로 바꿔줍니다.

• 안전성: AI 가 왜 잘못된 답을 냈는지, 혹은 편견을 가졌는지 정확한 원인을 찾아낼 수 있습니다.
• 신뢰: "이 AI 는 복잡한 계산 없이, 이 핵심 논리만 따라서 답을 냈어"라고 설명할 수 있게 되어, 우리가 AI 를 더 믿고 사용할 수 있습니다.

요약

이 논문은 **"AI 를 더 똑똑하게 만드는 게 아니라, AI 를 더 '간결하게' 만들어서 우리가 그 두뇌를 이해할 수 있게 만들자"**는 제안입니다.

마치 복잡한 도시의 교통 체증을 해결하기 위해, 불필요한 도로를 막고 핵심 도로만 남긴 뒤, 그 핵심 도로만 따라가면 목적지에 쉽게 도착할 수 있게 만든 것과 같습니다. AI 의 성능은 그대로 유지하면서, 그 안을 들여다보는 우리가 훨씬 더 명확하게 볼 수 있게 된 것입니다.

기술 요약

1. 문제 제기 (Problem)

• 복잡성과 불투명성: 최근 LLM 의 성능 향상은 모델 크기와 복잡성의 증가와 동반되지만, 이로 인해 모델 내부의 작동 원리는 여전히 불투명합니다.
• 해석의 병목 현상: 기존 기계적 해석 가능성 연구에서는 특정 작업을 수행하는 '회로 (Circuit)'를 발견하려 하지만, 밀집된 (Dense) 어텐션 구조 때문에 회로가 수백 개의 어텐션 헤드와 MLP(다층 퍼셉트론) 로 복잡하게 얽혀 있어 해석이 어렵습니다.
• 기존 방법의 한계: 계산 효율성을 위한 희소 어텐션 (Sparse Attention) 방법들은 고정된 패턴 (예: 슬라이딩 윈도우) 을 사용하거나, Top-k 선택 시 $k$값을 고정해야 하는 등 모델이 학습한 내부 계산 구조를 반영하지 못하거나 유연성이 부족합니다.

2. 방법론 (Methodology)

저자들은 사전 훈련된 모델을 다시 처음부터 훈련하지 않고, 후훈련 (Post-Training) 단계에서 어텐션 패턴을 희소하게 만드는 방법을 제안합니다.

• 희소 어텐션 레이어 (Sparse Attention Layer):

  • 표준 소프트맥스 (Softmax) 어텐션 대신 SPARTAN 프레임워크의 희소 어텐션을 도입합니다.
  • 학습 가능한 분포 (Bernoulli) 를 통해 어텐션 가중치를 이진 (0 또는 1) 으로 샘플링하여, 불필요한 어텐션 엣지를 '0'으로 만듭니다.
  • Gumbel-Softmax 트릭을 사용하여 이 이진 샘플링 과정을 미분 가능하게 만듭니다.
  • 하드 샘플링 후의 함수 형태는 표준 어텐션과 동일하므로, 사전 훈련된 가중치를 그대로 초기화하여 사용할 수 있습니다.

• 제약된 손실 최적화 (Constrained Optimisation):

  • 목표: 어텐션 엣지의 수를 최소화 (희소성 극대화) 하되, 다음 토큰 예측 손실 (Cross-Entropy) 은 사전 훈련된 베이스 모델 수준으로 유지합니다.
  • GECO 알고리즘: Lagrange 승수 ($\lambda$) 를 사용하여 손실 제약 조건을 자동으로 조절합니다.
    • 목적 함수: $\min_{\theta} \sum E[|A_l|]$ (희소성) subject to $CE \le \tau$ (손실 제약).
    • 훈련 중 $\lambda$를 업데이트하여 모델 성능이 저하되지 않는 선에서 희소성 정규화 강도를 적응적으로 증가시킵니다.

• 실용적 기법:

  • 대규모 모델 (7B 파라미터) 에서는 LoRA (Low-Rank Adaptation) 파인튜닝을 사용하여 계산 비용을 절감했습니다.
  • FlashAttention과 지식 증류 (Distillation) 기법을 적용하여 훈련 안정성과 효율성을 높였습니다.

3. 주요 기여 및 결과 (Key Contributions & Results)

A. 성능 유지와 극도의 희소성 달성

• 모델: GPT-2 (1.24 억 파라미터) 및 OLMo-7B (70 억 파라미터) 모델에서 실험 수행.
• 결과:
  • 모델의 성능 (Cross-Entropy) 은 베이스 모델과 거의 동일하게 유지됨.
  • 어텐션 엣지의 활성 비율이 약 0.22% (GPT-2) 및 0.44% (OLMo) 수준으로 급격히 감소 (약 99.5% 이상 희소화).
  • 이는 기존 밀집 모델 대비 약 0.4% 수준의 연결성만으로도 동일한 성능을 낼 수 있음을 의미합니다.

B. 단순화된 회로 발견 (Circuit Discovery)

• 활성화 패치 (Activation Patching) 실험: 특정 작업 (예: 복사, IOI, Greater-Than) 을 수행하는 데 필요한 최소한의 구성 요소를 찾음.
• 결과:
  • 희소 모델은 밀집 모델보다 4 배 적은 어텐션 헤드와 최대 100 배 적은 엣지로 동일한 작업의 90% 이상을 설명 가능.
  • 예: '복사 (Copy)' 작업에서 희소 모델은 명확한 '인덕션 헤드 (Induction Head)' 패턴을 보이며, 밀집 모델은 61 개의 헤드가 필요한 반면 희소 모델은 9 개의 헤드로 해결.

C. 어텐션 귀속 (Attribution) 의 용이성

• 크로스레이어 트랜스코더 (Cross-Layer Transcoders) 활용: 특징 (Feature) 간 상호작용을 분석하는 어트리뷰션 그래프를 생성.
• 문제 해결: 기존 밀집 모델에서는 특징 간 연결을 매개하는 어텐션 구성 요소가 너무 많아 해석이 불가능했으나, 희소 모델에서는 이를 매개하는 구성 요소가 극도로 줄어듦.
• 결과:
  • 특징 간 연결을 설명하는 데 필요한 키 - 쿼리 쌍이 16.1 배, 어텐션 헤드가 3.4 배 감소.
  • 계산 비용이 수 시간에서 수 분으로 단축되었으며, 특징 기반 관점과 회로 기반 관점을 통합된 시각으로 이해할 수 있게 됨.

4. 의의 및 결론 (Significance)

• 계산의 중복성 확인: 트랜스포머의 어텐션 메커니즘은 현재 사용 중인 것보다 훨씬 더 희소할 수 있으며, 대부분의 계산이 중복됨을 시사합니다.
• 해석 가능성의 새로운 패러다임: 모델 설계 단계나 후훈련 단계에서 **희소성을 유도적 편향 (Inductive Bias)**으로 활용하면, 성능을 희생하지 않으면서도 인간이 이해하기 쉬운 구조화된 회로를 가진 모델을 만들 수 있습니다.
• 미래 방향: 이 방법은 어텐션뿐만 아니라 MLP, 모델 가중치 희소화, 또는 강화학습 (RL) 기반 후훈련 등 다른 영역으로 확장 가능하며, 더 투명하고 신뢰할 수 있는 AI 모델 개발의 핵심 원칙이 될 수 있습니다.

요약

이 논문은 **"성능을 유지하면서 어텐션 연결을 99% 이상 줄이면, 모델 내부의 복잡한 계산이 단순하고 해석 가능한 회로로 재구성된다"**는 가설을 실험적으로 입증했습니다. 이는 기계적 해석 가능성 연구의 핵심 병목 현상인 '복잡성'을 해결하고, 더 투명하고 안전한 AI 시스템을 구축하는 데 중요한 통찰을 제공합니다.