To See the Unseen: on the Generalization Ability of Transformers in Symbolic Reasoning

이 논문은 Decoder-only Transformer 모델이 훈련 데이터에서 관찰되지 않은 토큰 (unseen tokens) 을 포함하는 추상적 기호 논리 추론 (abstract symbolic reasoning) 을 수행할 때 발생하는 일반화 실패의 원인을 규명하고, 이를 해결하기 위한 방법론을 제안합니다.

주요 내용은 다음과 같습니다.

1. 연구 배경 및 문제 제기

• 문제: 최근 대규모 언어 모델 (LLM) 은 수학 및 과학과 같은 형식적 추론 작업에서 뛰어난 성과를 보이지만, 입력 토큰의 이름이나 값을 변경하거나 관련 없는 정보를 삽입하는 등의 간단한 변화에도 취약합니다. 이는 모델이 진정한 추론이 아닌 표면적인 패턴 (token bias) 에 의존하고 있음을 시사합니다.
• 기존 연구의 한계:
  • Boix-Adserà et al. (2024) 은 토큰 복사 (copying) 능력이 부족하여 unseen 토큰 처리에 실패한다고 주장했으나, 여러 개의 unseen 변수가 존재할 때 구조적 변경만으로는 해결되지 않음을 발견했습니다.
  • Anand et al. (2025) 은 훈련 중 임베딩을 주기적으로 초기화하는 '능동적 망각 (active forgetting)' 전략이 효과적임을 보였으나, 그 기저 메커니즘을 명확히 설명하지 못했습니다.
• 핵심 질문: 왜 모델은 훈련 중 보지 못한 새로운 기호 (토큰) 를 구별하거나 복사하는 데 실패할까요?

2. 핵심 발견: (Un)Embedding Collapse

논문은 unseen 토큰의 일반화 실패에 대한 새로운 기작을 발견했습니다.

• 현상: 훈련 과정에서 관찰되지 않은 토큰 (unseen tokens) 의 (un)embedding 벡터가 거의 동일한 벡터로 수렴 (collapse) 하는 현상이 발생합니다.
• 원인:
  • 이론적 증명: L2 정규화 (weight decay) 와 LayerNorm 을 사용하는 SGD(또는 AdamW) 기반 훈련에서, 모델이 데이터를 잘 적합 (fit) 시키게 되면 관찰되지 않은 클래스 (레이블) 에 해당하는 가중치 벡터들 사이의 거리가 수렴 단계마다 감소함을 수학적으로 증명했습니다.
  • 결과: 이 수렴 현상으로 인해 모델은 서로 다른 unseen 변수들을 구별하기 어려워지고, 특히 임베딩과 언임베딩 파라미터가 공유 (tied) 되는 아키텍처에서 이 문제가 심화됩니다.
• 실험적 증거: 합성 논리 문제 실험과 Gemma 3 모델 분석을 통해, unseen 토큰들의 코사인 유사도가 훈련 중 급격히 증가하여 거의 동일한 방향을 향하게 됨을 확인했습니다.

3. 제안된 방법론 및 해결책

저자들은 이 '임베딩 붕괴' 문제를 해결하기 위해 세 가지 기술의 조합을 제안합니다.

1. 복사 가능 아키텍처 (Copy-enabled Architecture):
   • Boix-Adserà et al. (2024) 의 제안을 확장하여, 입력 토큰의 임베딩을 직접 출력에 연결하는 **'Copy Attention Head'**를 추가했습니다. 이는 모델이 컨텍스트 내의 토큰을 복사하여 생성하는 능력을 향상시킵니다.
2. 데이터 다양성 (Data Diversity):
   • 훈련 데이터에 포함된 기호 (predicates) 의 종류를 대폭 늘려 (예: 150 개에서 9600 개까지), 모델이 다양한 기호 패턴을 학습하도록 유도합니다.
3. 임베딩 고정 또는 주기적 재초기화:
   • Frozing: 훈련 중 (un)embedding 파라미터를 고정하여 초기화 시의 직교성 (orthogonality) 을 유지하게 합니다.
   • Active Forgetting: 훈련 중 주기적으로 (예: 100 스텝마다) 임베딩을 재초기화하여 붕괴를 방지합니다.

4. 실험 결과

• 합성 논리 추론 실험:
  • 표준 Transformer 는 unseen 토큰이 포함된 쿼리에서 성능이 급격히 떨어졌습니다.
  • Copy Attention + 데이터 다양성: 단일 unseen 토큰이 있는 경우 성능이 크게 개선되었으나, 모든 변수가 unseen 인 경우에는 여전히 실패했습니다.
  • Embedding Freeze/Reset + Copy Attention: 임베딩을 고정하거나 주기적으로 재초기화하고 Copy Attention 을 적용한 모델은 모든 변수가 unseen 인 경우에도 높은 정확도를 달성했습니다. 이는 임베딩 붕괴가 일반화 실패의 주요 원인임을 강력히 지지합니다.
• Gemma 3 모델 분석:
  • 오픈 가중치 모델인 Gemma 3 에서도 99 개의 의도적으로 사용되지 않는 토큰 (unused tokens) 이 높은 코사인 유사도를 보이며 붕괴 현상이 관찰되었습니다.
  • 이러한 collapsed 임베딩을 가진 토큰으로 파인튜닝 (finetuning) 을 수행할 경우, 일반 토큰에 비해 **수렴 속도가 매우 느린 것 (약 10 배 이상)**이 확인되었습니다.

5. 기여 및 의의

1. 메커니즘적 설명: Transformer 의 기호 추론 일반화 실패가 단순한 복사 능력 부족이 아니라, **훈련 중 발생하는 임베딩 벡터의 수렴 (collapse)**에 기인함을 이론적, 실험적으로 증명했습니다.
2. 기존 연구 해석: Boix-Adserà et al. (2024) 의 아키텍처 변경이 여러 unseen 변수 앞에서 왜 실패했는지, 그리고 Anand et al. (2025) 의 '능동적 망각' 전략이 왜 효과적인지에 대한 명확한 설명을 제공합니다.
3. 실용적 제안: 추론 능력을 갖춘 모델을 구축하기 위해 **Copy Attention, 데이터 다양성, 임베딩 관리 (고정/재초기화)**를 결합한 구체적인 아키텍처 및 훈련 전략을 제시했습니다.
4. 대규모 모델 시사점: Gemma 3 와 같은 실제 모델에서도 유사한 붕괴 현상이 발생하며, 이는 다운스트림 파인튜닝 성능에 부정적인 영향을 미칠 수 있음을 경고했습니다.

결론

이 논문은 Transformer 가 추상적 기호를 다룰 때 겪는 일반화 한계의 근본 원인을 '임베딩 붕괴'로 규명하고, 이를 해결하기 위한 구조적 및 훈련적 개입 방법을 제시함으로써, LLM 의 추론 능력 향상을 위한 중요한 통찰을 제공합니다.