Compressed Sensing for Capability Localization in Large Language Models

이 논문은 트랜스포머 기반 대규모 언어 모델에서 다양한 능력이 소수의 어텐션 헤드에 국한되어 존재한다는 것을 발견하고, 압축 센싱 기법을 활용해 이러한 기능적 구성 요소를 효율적으로 식별하는 방법을 제시합니다.

핵심

거대 언어 모델의 '비밀 병기' 찾기: 압축 센싱을 활용한 능력 국소화 연구

이 논문은 최근 인공지능, 특히 거대 언어 모델 (LLM) 이 어떻게 작동하는지에 대한 흥미로운 비밀을 밝혀냈습니다. 마치 거대한 도시의 전신주에서 특정 전선 하나만 끊었을 때, 그 전선과 관련된 특정 구역만 정전되는 현상을 발견한 것과 비슷합니다.

이 연구의 핵심 내용을 일상적인 비유로 쉽게 설명해 드리겠습니다.

---

1. 핵심 발견: "작은 부품 하나가 거대한 능력을 좌우한다"

거대 언어 모델 (예: Llama, Qwen 등) 은 수학 문제 풀이, 코드 작성, 시 짓기 등 다양한 능력을 가지고 있습니다. 연구진은 이 모델들이 이 모든 능력을 두루뭉술하게 처리하는 것이 아니라, 매우 작고 구체적인 부품들 (주의 헤드, Attention Heads) 에 특화되어 있다는 것을 발견했습니다.

• 비유: 거대 언어 모델을 거대한 오케스트라라고 상상해 보세요.
  • 보통은 모든 악기가 함께 연주해서 아름다운 음악을 만들어낸다고 생각합니다.
  • 하지만 이 연구에 따르면, 수학 문제를 풀 때는 오케스트라 중 '바이올린 5 대'만 집중적으로 연주하고, 코드를 작성할 때는 '트럼펫 3 대'만 집중적으로 연주하는 것입니다.
  • 나머지 수백 개의 악기들은 그 순간에는 거의 침묵하고 있거나 다른 일을 합니다.

2. 실험: "특정 부품만 꺼내면 어떻게 될까?"

연구진은 이 가설을 검증하기 위해 모델의 특정 '부품 (주의 헤드)'들을 끄고 (Zeroing out) 성능을 측정했습니다.

• 결과: 수학 능력을 담당하는 '바이올린 5 대'만 끄자, 모델의 수학 실력은 65% 이상 급격히 떨어졌습니다.
• 신기한 점: 하지만 수학 헤드만 꺼낸 것이지, 다른 악기들은 그대로 두었기 때문에 코드 작성이나 일반적인 대화 능력은 거의 변하지 않았습니다. 마치 오케스트라에서 수학 연주를 담당하던 바이올린만 잠잠하게 만들자, 그날의 코딩 연주는 여전히 완벽하게 이루어진 것과 같습니다.

3. 방법론: "수천 개의 전선을 일일이 확인하지 않는 지혜"

이렇게 특정 부품을 찾으려면, 수천 개의 헤드 중 하나씩 끄며 테스트하는 '일일이 확인하는 방법 (Greedy Search)'을 쓸 수 있습니다. 하지만 이 방법은 시간과 비용이 너무 많이 듭니다. (수천 번의 모델 실행이 필요함).

연구진은 **'압축 센싱 (Compressed Sensing)'**이라는 수학적 기술을 도입했습니다.

• 비유: 어두운 방에 숨겨진 5 개의 보석 (특수 헤드) 을 찾는 상황입니다.
  • 기존 방법: 방 구석구석을 하나씩 비추며 찾는 것. (매우 느림)
  • 이 연구의 방법: 방 전체를 한 번에 비추는 특수 카메라를 사용합니다. 보석은 아주 희귀하고 (희소성), 빛을 반사하는 방식이 단순하다 (선형성) 는 가정을 합니다.
  • 결과: 아주 적은 횟수의 촬영 (모델 평가) 만으로도, 어디에 보석이 숨겨져 있는지 정확히 찾아냅니다. 기존 방법보다 50 배 이상 빠르고 효율적입니다.

4. 흥미로운 추가 발견: "만능 부품"과 "모델 크기의 차이"

연구진은 두 가지 더 재미있는 현상을 발견했습니다.

A. 만능 부품 (Universal Heads)

어떤 부품은 수학, 코드, 언어 이해 등 모든 작업에 동시에 관여합니다.

• 비유: 오케스트라의 지휘자나 메인 스테이지 조명과 같습니다.
• 이 부품만 끄면, 수학도 못하고 코드도 못 쓰며, 심지어 말이 안 되는 소리를 내거나 반복되는 이상한 행동을 합니다. 이는 특정 능력보다는 모델이 말을 잘하기 위한 기본 토대를 담당하는 부품들입니다.

B. 모델 크기에 따른 차이 (Scale Dependence)

모델의 크기가 커질수록 능력의 국소화 (특정 부품에 집중) 가 더 뚜렷해집니다.

• 작은 모델: 수학 문제도, 코드도, 일반 대화도 유사한 몇 개의 부품이 함께 처리합니다. (모두가 다재다능하지만 전문성은 떨어짐)
• 큰 모델: 각 능력이 완전히 다른 전용 부품으로 분리되어 있습니다. (전문성이 매우 뚜렷함)
• 비유: 작은 식당에서는 요리사 한 명이 모든 요리를 하지만, 대형 호텔에서는 수석 셰프, 소스 셰프, 디저트 셰프가 각각 완벽하게 분업합니다.

5. 왜 이 연구가 중요한가요?

이 발견은 인공지능의 '블랙박스'를 조금씩 열어젖히는 중요한 열쇠가 됩니다.

1. 해석 가능성 (Interpretability): 모델이 왜 그런 답을 냈는지, 어떤 부품이 그 일을 했는지 이해할 수 있게 됩니다.
2. 모델 편집 (Model Editing): 원하지 않는 능력 (예: 혐오 발언 생성) 만을 정확히 제거하고, 다른 능력은 그대로 유지할 수 있습니다. 마치 특정 악기 소리만 줄이는 EQ 조절처럼요.
3. AI 안전성: 위험한 지식이나 능력을 가진 특정 부품을 찾아내어 차단함으로써, 더 안전한 AI 를 만들 수 있습니다.

요약

이 논문은 **"거대 언어 모델은 모든 것을 다 잘하는 만능 기계가 아니라, 수많은 작은 전문가 (부품) 들이 모여 일하는 팀"**임을 증명했습니다. 그리고 우리는 매우 적은 노력으로 이 '전문가 팀' 중 누가 어떤 일을 하는지 찾아낼 수 있는 방법을 개발했습니다. 이는 AI 를 더 투명하고 안전하며 통제 가능하게 만드는 큰 진전입니다.

기술 요약

1. 문제 정의 (Problem Statement)

대형 언어 모델 (LLM) 은 수학적 추론, 코드 생성, 언어적 행동 등 다양한 고수준 능력을 보유하고 있습니다. 기존 연구는 사실적 지식 (factual knowledge) 이 특정 뉴런이나 레이어에 국소화 (localization) 될 수 있음을 보여주었으나, 복잡한 행동 능력 (capabilities) 이 Transformer 아키텍처 내의 어떤 구성 요소에 집중되어 있는지는 명확하지 않았습니다.

이 연구는 다음과 같은 핵심 문제를 다룹니다:

• 국소화 가설: 특정 작업 (Task) 에 필요한 계산이 모델 전체에 분산되어 있는 것이 아니라, Attention Head 의 작은 부분집합 (Sparse Subset) 에 집중되어 있는가?
• 효율적인 식별: 수천 개의 Attention Head 가 있는 모델에서 특정 작업을 수행하는 Head 를 식별하기 위해, 모든 Head 를 하나씩 제거 (Ablation) 하는 비효율적인 탐색 (Greedy Search) 없이도 압축 센싱 (Compressed Sensing) 기법을 통해 효율적으로 찾을 수 있는가?

2. 방법론 (Methodology)

저자들은 압축 센싱 (Compressed Sensing) 프레임워크를 활용하여 희소성 (Sparsity) 을 가진 작업별 Attention Head 를 식별하는 새로운 알고리즘을 제안합니다.

가. 핵심 가정

1. 희소성 (Sparsity): 임의의 특정 작업에 대해 전체 Head(N) 중 극히 일부 (k, $k \ll N$) 만이 성능에 결정적인 기여를 합니다.
2. 가법성 (Additivity): 여러 Head 를 동시에 제거했을 때의 성능 저하는 각 Head 를 개별적으로 제거했을 때의 영향의 합으로 근사할 수 있습니다 (1 차 선형 상호작용 우세).

나. 압축 센싱 기반 식별 알고리즘

기존의 Greedy Search 는 $O(N \times k)$ 번의 모델 평가가 필요하지만, 제안된 방법은 이를 획기적으로 줄입니다.

1. 측정 행렬 구성 ($\Phi$): $M$개의 평가 (Measurement) 를 수행합니다. 각 평가에서 $N$개의 Head 중 일부는 무작위 또는 계층적 (Stratified) 으로 제거 (Zeroing out) 됩니다.
   • Bernoulli Sampling: 각 Head 를 독립적인 확률로 제거.
   • Stratified Sampling: 모든 Head 가 균등하게 제거되도록 제약을 주어 추정의 안정성을 높임 (실험상 더 우수함).
2. 선형 시스템 모델링: 성능 변화 ($y$) 를 Head 제거 여부 ($\Phi$) 와 각 Head 의 영향력 ($x$) 의 선형 관계로 모델링합니다 ($y = \Phi x + \epsilon$).
3. 희소 회귀 (Sparse Regression): Lasso($L_1$ 정규화) 최적화 문제를 풀어 각 Head 의 영향력 벡터 $\hat{x}$를 추정합니다.
   • 식: $\hat{x} = \arg \min_x \frac{1}{2M} \|y - (\beta_0 + \Phi x)\|_2^2 + \lambda \|x\|_1$
   • 가장 큰 음수 계수 (가장 큰 성능 저하를 유발) 를 가진 Head 들을 식별합니다.

3. 주요 기여 (Key Contributions)

1. Capabilty Localization 발견: LLM 의 다양한 능력 (수학, 코드, 욕설, 운율 등) 이 특정 Attention Head 의 매우 작은 집합에 국소화되어 있음을 실증했습니다.
2. 효율적인 식별 알고리즘: 수천 번의 모델 평가가 필요한 기존 방법 대비 최대 50 배 이상 적은 평가 횟수로 작업별 Head 를 정확히 식별하는 압축 센싱 기반 방법을 개발했습니다.
3. 보편적 Head (Universal Heads) 발견: 특정 작업뿐만 아니라 여러 작업 전반에 걸쳐 핵심적인 역할을 수행하는 Head 들을 발견했습니다. 이를 제거하면 모델이 일관된 언어 생성 능력을 상실하고 병리적 행동 (반복, 퇴화) 을 보입니다.
4. 규모 의존적 국소화 (Scale-dependent Localization): 모델의 규모 (Parameter Size) 에 따라 국소화 패턴이 달라짐을 발견했습니다. 특히 작은 모델에서는 작업별 Head 보다는 '지식 기반 객관식' 형태의 공유 Head 가 두드러지는 반면, 큰 모델에서는 더 명확한 작업별 국소화가 나타납니다.

4. 실험 결과 (Results)

연구진은 Llama (1B8B) 와 Qwen (3B7B) 계열의 5 가지 모델과 수학 (GSM8K, Arithmetic), 코드 (MBPP, HumanEval), 언어 (욕설, 운율) 등 다양한 능력을 대상으로 실험을 수행했습니다.

• 성능 저하: 식별된 상위 5 개의 작업별 Head 를 제거했을 때, 해당 작업의 성능은 최대 65% 까지 급격히 감소했습니다 (예: Qwen 2.5-7B 의 수학 능력 -65.4%).
• 선택성 (Specificity): 작업별 Head 를 제거하더라도 관련 없는 일반 언어 능력 (HellaSwag, MMLU 등) 은 거의 영향을 받지 않았습니다. 이는 모델이 모듈식 (Modular) 으로 조직화되어 있음을 시사합니다.
• 효율성 비교:
  • Greedy Search: 약 5,120 회 평가 필요.
  • Stratified Compressed Sensing: 약 100~200 회 평가로 동등하거나 더 나은 성능 저하 달성 (Table 4 참조).
• 일반화: 한 데이터셋 (예: GSM8K) 으로 식별된 Head 는 다른 데이터셋 (예: Arithmetic) 에서도 동일하게 성능을 저하시켰습니다. 이는 Head 가 특정 데이터가 아닌 '능력' 자체를 담당함을 의미합니다.

5. 의의 및 결론 (Significance & Conclusion)

이 연구는 Transformer 기반 LLM 의 기능적 조직화에 대한 중요한 통찰을 제공합니다.

• 해석 가능성 (Interpretability): 모델이 어떻게 복잡한 능력을 구현하는지 이해하는 데 있어, 뉴런 수준이 아닌 Attention Head 수준에서의 국소화가 핵심임을 보여줍니다.
• 모델 편집 (Model Editing): 특정 능력 (예: 해로운 지식, 특정 언어 습관) 만을 선택적으로 제거하거나 수정하기 위해, 해당 능력을 담당하는 소수의 Head 만을 타겟팅하여 편집할 수 있는 가능성을 열었습니다.
• AI 안전성 (AI Safety): 보편적 Head 와 작업별 Head 를 구분함으로써, 모델의 핵심 언어 능력을 해치지 않으면서도 특정 위험 능력을 제어하는 정밀한 안전 장치 개발에 기여할 수 있습니다.

결론적으로, 대형 언어 모델은 무작위로 분산된 계산이 아니라, 희소하고 기능적으로 구별되는 컴포넌트 (Attention Head) 에 의해 모듈식으로 조직화되어 있으며, 이를 효율적으로 식별하는 압축 센싱 기법은 모델 분석 및 제어의 새로운 표준이 될 수 있음을 시사합니다.