Understanding In-Context Learning Beyond Transformers: An Investigation of State Space and Hybrid Architectures

이 논문은 행동 분석과 개입 기반 방법을 결합하여 다양한 아키텍처의 대규모 언어 모델에서 문맥 학습 (ICL) 의 내부 메커니즘을 탐구하며, 특히 기능 벡터가 매개변수 지식 검색에 중요한 역할을 하지만 Mamba2 는 다른 메커니즘을 사용할 수 있음을 규명했습니다.

핵심

🧠 1. 연구의 배경: "새로운 문제"를 푸는 능력 (ICL)

대부분의 AI 는 미리 학습된 지식만 가지고 있습니다. 하지만 최신 AI 는 학습 없이도 시험 문제 앞에 제시된 몇 가지 예시 (예: "사과→과일, 당근→채소"를 보고 "포도→?"를 맞히는 것) 를 보고 그 패턴을 파악해 답을 맞출 수 있습니다. 이를 **맥락 학습 **(In-Context Learning)이라고 부릅니다.

기존 연구는 이 능력이 주로 **변환기 **(Transformer)라는 특정 뇌 구조에서 어떻게 작동하는지 분석했습니다. 하지만 최근에는 Mamba나 **하이브리드 **(Transformer+Mamba)라는 새로운 뇌 구조의 AI 들도 등장했는데, 이들이 똑같은 일을 할 때 내부적으로 어떤 다른 방식으로 작동하는지는 아직 알려지지 않았습니다.

🔍 2. 연구 방법: 두 가지 다른 시험 유형

연구진은 AI 들에게 두 가지 종류의 시험을 치르게 했습니다.

1. **기억력 테스트 **(Parametric Knowledge Retrieval)
   • 비유: "대한민국의 수도는?" 같은 사실 지식을 묻는 문제.
   • AI 가 이미 머릿속에 저장해 둔 지식을 꺼내와서 답을 찾는 과정입니다.
2. **이해력 테스트 **(Contextual Knowledge Understanding)
   • 비유: "이 글에서 화난 사람은 누구인가?"처럼 지문 내용을 읽고 추론하는 문제.
   • AI 가 지문이라는 '새로운 맥락'을 이해하고 답을 찾아야 합니다.

🧪 3. 주요 발견: 겉은 비슷해도 속은 다르다!

① 겉보기엔 비슷하지만, 내부 엔진은 다름

모든 AI 모델 (변환기, Mamba, 하이브리드) 이 시험 점수는 비슷하게 잘 받았습니다. 하지만 내부 작동 방식을 들여다보니 완전히 달랐습니다.

• 비유: 같은 '자동차'라도 엔진이 가솔린인지, 전기인지, 하이브리드인지에 따라 작동 원리가 다르듯이, AI 도 겉으로는 똑같은 일을 하지만 내부 회로는 다릅니다.

② '기능 벡터 (FV)'라는 특수 부대의 역할

기존 연구에서는 AI 가 새로운 문제를 풀 때 **'기능 벡터 **(Function Vectors)라는 특정 부서가 핵심 역할을 한다고 했습니다. 마치 특수 부대가 작전을 수행하는 것처럼요.

• **기억력 테스트 **(사실 지식)
  • **변환기 **(Transformer)와 Mamba 모델에서는 이 '특수 부대' (자율 주의 Attention 레이어) 가 아주 중요한 역할을 했습니다.
  • 하지만 Mamba2 모델은 이 특수 부대를 거의 쓰지 않았습니다. Mamba2 는 완전히 다른 방식으로 기억력 테스트를 해결하고 있었습니다.
• **이해력 테스트 **(지문 이해)
  • 사실 지식을 묻는 문제와 달리, 지문을 이해하는 문제에서는 '특수 부대'의 중요도가 크게 떨어졌습니다. AI 는 다른 방식으로 맥락을 이해하고 있었습니다.

③ 하이브리드 모델의 비밀: "이중 엔진"

변환기와 Mamba 를 섞은 하이브리드 모델은 흥미로운 결과를 보였습니다.

• 비유: 하이브리드 차처럼 두 가지 엔진 (Attention 과 Mamba) 이 달렸지만, 새로운 문제를 풀 때는 주로 'Attention 엔진'이 주도했습니다.
• 특히 Mamba 엔진 쪽을 건드리면 오히려 성능이 떨어지기도 했습니다. 이는 하이브리드 모델이 새로운 학습을 할 때, 기존에 잘 알려진 'Attention' 방식에 더 의존한다는 뜻입니다.

💡 4. 결론: 왜 이 연구가 중요한가요?

이 연구는 **"모든 AI 가 똑같은 방식으로 학습하는 게 아니다"**라는 사실을 밝혀냈습니다.

• 과거의 생각: AI 는 모두 같은 원리 (특수 부대) 로 학습한다.
• 새로운 발견:
  • Mamba2는 아예 다른 원리로 학습한다.
  • 기억력과 이해력은 AI 내부에서 서로 다른 부서가 담당한다.
  • 하이브리드 모델은 두 엔진 중 하나에 더 의존한다.

🌟 한 줄 요약

"똑똑한 AI 들이 시험을 잘 치는 건 비슷해 보이지만, 그 속을 들여다보면 기억력과 이해력을 처리하는 방식이 모델마다 천차만별입니다. 특히 최신 모델인 Mamba2는 기존 AI 들과 완전히 다른 '비밀 무기'를 쓰고 있었죠!"

이 연구를 통해 우리는 AI 의 내부 작동 원리를 더 정교하게 이해하게 되었고, 앞으로 더 효율적이고 똑똑한 AI 를 설계하는 데 중요한 단서를 얻었습니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• 컨텍스트 학습 (ICL) 의 확장: 대규모 언어 모델 (LLM) 의 핵심 능력인 컨텍스트 학습 (ICL) 은 기존 연구에서 주로 트랜스포머 (Transformer) 아키텍처에 집중되어 있었습니다. 특히, ICL 을 수행하는 메커니즘으로 '함수 벡터 (Function Vectors, FVs)'와 '인덕션 헤드 (Induction Heads)'가 트랜스포머의 셀프 어텐션 헤드에서 발견되었습니다.
• 새로운 아키텍처의 등장: 최근 **상태 공간 모델 (State Space Models, SSMs, 예: Mamba, Mamba2)**과 **하이브리드 아키텍처 (트랜스포머 + SSM 결합, 예: Hymba, Zamba2)**가 등장하며 효율성을 높이고 있습니다.
• 연구 격차: 이러한 새로운 아키텍처들이 ICL 을 수행할 수 있다는 것은 알려져 있으나, 내부 메커니즘이 어떻게 작동하는지, 특히 트랜스포머와 동일한 FV 메커니즘을 사용하는지 여부는 명확히 규명되지 않았습니다. 또한, ICL 작업 유형에 따라 (지식 기반 vs 문맥 이해) 내부 작동 방식이 어떻게 다른지도 연구가 부족했습니다.

2. 연구 방법론 (Methodology)

저자는 행동적 분석 (Behavioral Analysis) 과 기계적 해석 가능성 (Mechanistic Interpretability) 분석을 결합하여 다양한 아키텍처를 평가했습니다.

• 평가 대상 모델:
  • SSM: Mamba-1.4B, Mamba2-1.3B
  • 트랜스포머: Gemma-3-1B, Llama-3.2-1B, Qwen2.5-1.5B
  • 하이브리드: Hymba-1.5B (병렬 구조), Zamba2-1.2B (직렬 구조)
  • 모든 모델은 약 1B 파라미터로 제한하여 공정한 비교를 진행했습니다.
• 작업 유형 분류 (Jin et al., 2025 기준):
  1. 파라미터적 지식 검색 (Parametric Knowledge Retrieval): 모델 내부 지식 (국가 - 수도, 통화 등) 을 기반으로 답을 찾는 작업.
  2. 문맥적 지식 이해 (Contextual Knowledge Understanding): 프롬프트 내 문맥을 이해하고 추론해야 하는 작업 (감정 분석, 혐오 발언 분류 등).
• 실험 기법:
  • 행동 실험: 라벨 무작위화 (Label Randomization), 라벨 뒤집기 (Label Flipping) 실험을 통해 모델이 노이즈에 얼마나 강건한지, 그리고 새로운 관계를 학습할 수 있는지 확인.
  • 기계적 해석 (Mechanistic Interpretability):
    • FV 헤드 식별: Todd et al. (2024) 의 방법을 확장하여 SSM 과 하이브리드 모델의 헤드 (Attention 및 Mamba Stream) 에서 ICL 에 기여하는 함수 벡터 (FV) 를 식별.
    • 개입 실험 (Intervention):
      • Steering (조작): 식별된 FV 헤드의 출력을 특정 값으로 조작하여 성능 변화 측정.
      • Ablation (제거): FV 헤드를 제거 (제로 또는 평균값 대체) 하여 성능 저하 정도 측정.

3. 주요 결과 및 발견 (Key Results & Findings)

A. 행동적 분석 결과

• 성능 유사성: 모든 아키텍처는 ICL 수행 능력과 노이즈 라벨에 대한 강건성 측면에서 정성적으로 유사한 성능을 보였습니다.
• 학습 곡선의 차이: 문맥적 지식 이해 작업에서 Mamba2 와 Zamba2 는 'Gold(정답)' 예시를 제공받았을 때 성능 향상이 미미했으나, 'Label Flipping(거짓 관계 학습)' 실험에서는 새로운 관계를 학습할 수 있는 능력을 보여주었습니다. 이는 모델이 단순히 암기하는 것이 아니라 문맥을 학습하고 있음을 시사합니다.

B. 기계적 분석 결과 (핵심 발견)

1. FV 헤드의 위치:
   • 트랜스포머 및 Mamba: ICL 에 중요한 FV 헤드는 주로 셀프 어텐션 (Self-Attention) 및 Mamba 레이어에 집중되어 있습니다.
   • Mamba2: 흥미롭게도 Mamba2 는 다른 모델들과 달리 FV 헤드를 통해 ICL 을 수행하지 않는 것으로 보입니다. FV 조작 (Steering) 에 대한 반응이 미미했으나, 제거 (Ablation) 시에는 성능 저하가 발생하여 다른 메커니즘을 사용하거나 헤드가 다른 목적으로 사용됨을 시사합니다.
2. 하이브리드 모델의 메커니즘:
   • 하이브리드 모델 (Hymba, Zamba2) 에서 ICL 능력은 주로 셀프 어텐션 레이어에 위치한 FV 헤드에 의해 주도됩니다.
   • 병렬 (Hymba) 이나 직렬 (Zamba2) 구조 여부와 상관없이 어텐션 스트림이 핵심 역할을 합니다.
3. 작업 유형에 따른 FV 의 역할 차이:
   • 파라미터적 지식 검색: FV 헤드가 ICL 성능에 결정적인 역할을 합니다.
   • 문맥적 지식 이해: FV 헤드의 영향력이 상대적으로 적으며, 다른 메커니즘이 더 중요하게 작용합니다.
   • 헤드 중첩도: 두 가지 작업 유형을 위해 식별된 상위 FV 헤드는 대부분 겹치지 않습니다 (중첩률 낮음). 이는 모델이 작업 유형에 따라 서로 다른 내부 경로를 사용함을 의미합니다.

4. 주요 기여 (Key Contributions)

1. 아키텍처 확장: 기존 트랜스포머 중심이었던 FV 기반 ICL 분석을 SSM 과 하이브리드 아키텍처로 확장하여, 아키텍처 간 내부 메커니즘의 유사점과 차이점을 규명했습니다.
2. Mamba2 의 독특한 메커니즘 제안: Mamba2 가 기존 Mamba 나 트랜스포머와 다른 메커니즘으로 ICL 을 수행할 가능성을 최초로 제시했습니다.
3. 작업 유형별 메커니즘 분리: ICL 이 단일 메커니즘이 아니라, '지식 검색'과 '문맥 이해'라는 두 가지 다른 작업 유형에 따라 서로 다른 FV 세트와 내부 경로를 활용함을 증명했습니다.
4. 방법론적 제안: 행동적 분석과 기계적 해석을 결합한 접근법의 중요성을 강조하여, LLM 의 능력을 더 정교하게 이해하는 틀을 제공했습니다.

5. 의의 및 결론 (Significance)

이 연구는 LLM 의 컨텍스트 학습 능력이 아키텍처 (트랜스포머 vs SSM) 와 작업 유형 (지식 검색 vs 문맥 이해) 에 따라 어떻게 다른 내부 메커니즘을 통해 구현되는지를 심층적으로 규명했습니다. 특히, 하이브리드 모델에서 어텐션 레이어가 ICL 의 핵심 동력이라는 점과 Mamba2 의 비전통적 메커니즘을 발견함으로써, 차세대 효율적 LLM 설계와 해석 가능성 연구에 중요한 통찰을 제공합니다. 이는 단순히 성능만 비교하는 것을 넘어, 모델이 '어떻게' 학습하고 추론하는지에 대한 근본적인 이해를 넓혔다는 점에서 의미가 큽니다.