The Generalization Ridge: Information Flow in Natural Language Generation

논문 요약: 자연어 생성에서의 일반화 능선 (Generalization Ridge) 과 정보 흐름

이 논문은 트랜스포머 기반 언어 모델 (LLM) 이 자연어 생성 (NLG) 태스크를 수행할 때, 내부 계층을 통해 정보가 어떻게 흐르고 변환되는지 정보 이론적 관점에서 분석합니다. 저자들은 모델의 깊이에 따라 예측 정보가 비단조적 (non-monotonic) 으로 변화하며, 중간 계층에서 '일반화 능선 (Generalization Ridge)'이 형성된다는 새로운 현상을 발견하고 이를 정량화하는 프레임워크 InfoRidge를 제안합니다.

1. 문제 정의 (Problem)

• 배경: 트랜스포머 기반 언어 모델은 NLG 분야에서 뛰어난 성능을 보이지만, 학습 과정에서 태스크 관련 정보를 어떻게 통합하고 일반화하는지에 대한 내부 메커니즘은 여전히 불명확합니다.
• 기존 연구의 한계: 이전 연구들은 중간 계층이 최종 계층보다 더 일반화 가능한 표현을 가진다고 제안했으나, 이 일반화 능력이 학습 과정에서 어떻게 발생하고 계층 간에 전파되는지, 그리고 일반화 (generalization) 와 암기 (memorization) 가 계층별로 어떻게 조직화되는지에 대한 명확한 설명이 부족했습니다.
• 핵심 질문: NLG 태스크에서 학습 중 계층을 거치며 정보가 어떻게 진화하며, 네트워크의 어떤 계층이 일반화를 지원하는 기능적 역할을 수행하는가?

2. 방법론 (Methodology)

저자는 InfoRidge라는 정보 이론적 프레임워크를 제안하여 모델의 계층별 정보 흐름을 정량화합니다.

• 핵심 지표:
  1. 예측 정보 (Predictive Information, $I(Z_\ell; Y)$): $\ell$번째 계층의 은닉 표현 $Z_\ell$과 다음 토큰 레이블 $Y$ 사이의 상호 정보 (Mutual Information) 를 측정합니다. 이는 해당 계층이 태스크 관련 정보를 얼마나 잘 보존하고 있는지를 나타냅니다.
  2. 점증적 정보 획득 (Incremental Information Gain, $I(\Delta Z_\ell; Y)$): 인접한 두 계층 간의 잔차 변화 $\Delta Z_\ell = Z_\ell - Z_{\ell-1}$와 레이블 사이의 상호 정보를 측정합니다. 이는 각 트랜스포머 블록이 추가적으로 제공하는 예측 신호의 양을 나타냅니다.
• 추정 기법: 행렬 기반 상호 정보 추정 (Matrix-based Mutual Information Estimation, Giraldo et al., 2014) 을 사용하여 고차원 벡터 간의 상호 정보를 계산합니다. 가우시안 커널을 기반으로 한 Gram 행렬을 생성하고, 이를 통해 엔트로피와 상호 정보를 추정합니다.
• 추가 분석:
  • Attention 패턴 분석: 어떤 토큰에 주의를 기울이는지 분석하여 의미론적으로 중요한 정보 (신호) 가 집중되는 계층을 파악합니다.
  • 잔차 스케일링 (Residual Scaling): 각 잔차 블록에 학습 가능한 스칼라 계수 $\beta_\ell$를 도입하여, 분포 변화 (OOD) 상황에서 각 계층의 기여도가 어떻게 변하는지 인과적 증거를 확보합니다.

3. 주요 결과 (Key Results)

다양한 모델 (GPT-2, Qwen, LLaMA) 과 데이터셋 (CLUTRR, ECQA, Synthetic Arithmetic 등) 에 대한 실험을 통해 다음과 같은 일관된 패턴을 발견했습니다.

• 일반화 능선 (Generalization Ridge) 의 발견:
  • 예측 정보 ($I(Z_\ell; Y)$) 는 초기 계층에서 서서히 증가하다가 **중간~상위 중간 계층에서 정점 (Peak)**에 도달한 후, 최종 계층으로 갈수록 감소하는 비단조적 (non-monotonic) 인 3 단계 패턴을 보입니다.
  • 이 정점을 이루는 계층을 **"일반화 능선"**이라고 명명했습니다.
• 일반화 vs. 암기의 트레이드오프:
  • 중간 계층 (능선): 추상적이고 일반화 가능한 태스크 관련 특징을 가장 강력하게 인코딩합니다. 이 계층에서 분포 외 (OOD) 성능이 가장 높게 나타납니다.
  • 최종 계층: 예측 정보는 감소하지만 분포 내 (In-Distribution) 정확도는 100% 에 수렴합니다. 이는 최종 계층이 훈련 데이터의 특정 패턴을 **암기 (Memorization)**하여 일반화 능력을 희생하고 있음을 시사합니다.
• 점증적 정보 획득: 중간 계층에서 예측 정보의 증가량이 가장 큽니다. 이는 일반화 가능한 특징이 주로 중간 블록에서 생성됨을 의미합니다.
• Attention 및 잔차 스케일링 분석:
  • Attention 분석은 중간 계층이 태스크에 중요한 신호 토큰에 집중하는 반면, 최종 계층은 마지막 토큰이나 특정 데이터 포인트에 집중함을 보여줍니다.
  • 잔차 스케일링 실험에서, OOD 데이터에 적응시키려면 최종 계층의 가중치 ($\beta_\ell$) 를 낮추고 중간 계층의 가중치를 높여야 성능이 향상됨을 확인했습니다. 이는 일반화 능선이 분포 변화에 강건함을 입증합니다.
• 멀티-토큰 생성으로의 확장: CNN/DailyMail 요약 태스크와 같이 긴 시퀀스를 생성하는 경우에도, 생성 단계가 진행됨에 따라 전체 정보량은 증가하지만, **계층별 구조 (중간 계층에서 피크)**는 일정하게 유지됨을 확인했습니다.

4. 주요 기여 (Contributions)

1. 정보 흐름과 일반화의 연결: 학습 전 과정을 추적하여 예측 정보의 흐름과 일반화 행동 간의 명확한 연결을 확립했습니다. 특히 중간 계층의 비단조적 피크 (일반화 능선) 를 발견했습니다.
2. InfoRidge 프레임워크 제안: 자기회귀 언어 모델에 행렬 기반 상호 정보 추정을 적용하여 정보 흐름을 분석하는 새로운 정보 이론적 도구를 제시했습니다.
3. 인과적 증거 제공: Attention 패턴 분석과 학습 가능한 잔차 스케일링 계수를 통해 일반화 능선 현상에 대한 인과적 근거를 제시했습니다.

5. 의의 및 중요성 (Significance)

• 이론적 통찰: 트랜스포머 내부 메커니즘에 대한 이해를 심화시켜, 모델이 어떻게 일반화와 암기 사이에서 균형을 잡는지 구조적으로 설명합니다.
• 실용적 함의:
  • 모델 최적화: 일반화 성능을 극대화하기 위해 중간 계층의 표현을 활용하거나, 최종 계층의 과도한 암기를 줄이는 전략 (예: 조기 종료, 계층 프루닝) 에 대한 지침을 제공합니다.
  • 분포 외 (OOD) 일반화: OOD 환경에서 모델의 강건성을 높이기 위해 중간 계층에 더 의존하도록 모델을 설계하거나 튜닝할 수 있음을 시사합니다.
  • NLG 태스크 이해: 기존 분류 태스크 중심의 연구에서 벗어나, 자연어 생성 (Generative) 태스크에서의 정보 역학을 규명했습니다.

이 논문은 언어 모델의 "블랙박스"를 해명하는 중요한 단계로, 모델의 깊이에 따른 기능적 분업 (Functional Specialization) 을 정보 이론적으로 규명했다는 점에서 의의가 큽니다.