On the Predictive Power of Representation Dispersion in Language Models

1. 문제 정의 (Problem)

대형 언어 모델 (LLM) 은 텍스트 생성 및 코드 생성 등 다양한 작업에서 뛰어난 성능을 보이지만, 그 내부 임베딩 공간의 기하학적 구조에는 이방성 (anisotropy) 또는 랭크 붕괴 (rank collapse) 현상이 자주 관찰됩니다. 이는 은닉 상태 (hidden states) 가 좁은 원뿔 (cone) 안에 밀집하거나 저차원 부분 공간에 국한되는 현상을 의미합니다.

기존 연구들은 이러한 기하학적 구조가 모델의 표현 능력을 제한할 수 있다고 가정했으나, 임베딩 공간의 '넓이' (breadth) 가 모델의 자동 회귀적 텍스트 생성 능력 (예: 퍼플렉시티) 과 어떻게 구체적으로 연결되는지는 명확히 규명되지 않았습니다. 즉, 모델이 얼마나 잘 예측하는지와 그 내부 표현이 얼마나 넓게 퍼져 있는지 사이의 인과적, 상관적 관계가 부족했습니다.

2. 방법론 (Methodology)

이 논문은 **표현 분산 (Representation Dispersion)**이라는 새로운 지표를 도입하여 언어 모델의 예측 능력을 분석하고 활용하는 방법을 제시합니다.

• 측정 지표 (Representation Dispersion):

  • 특정 레이어 (기본적으로 최종 레이어) 의 은닉 벡터들 간의 **평균 쌍별 코사인 거리 (average pairwise cosine distance)**를 계산합니다.
  • 수식: $D = \frac{1}{\binom{N}{2}} \sum_{1 \le i < j \le N} \left[ 1 - \frac{E_i \cdot E_j}{\|E_i\|\|E_j\|} \right]$
  • 값이 클수록 임베딩 공간이 더 넓게 퍼져 있음을 의미하며, 값이 작을수록 표현이 압축되어 있음을 의미합니다.

• 실험 설계:

  • 모델 및 도메인: LLaMA, Qwen, Gemma, Mistral, Phi 등 다양한 모델 패밀리와 위키텍스트, 뉴스, 과학 초록 등 다양한 도메인에서 실험 수행.
  • 분석:
    1. 전역 분석: 퍼플렉시티 (Perplexity) 와 표현 분산의 상관관계, 레이어별 패턴, 파인튜닝 (Full FT, LoRA) 의 영향 분석.
    2. 세밀 분석: 의미적 클러스터 (동일한 10-gram 이어붙임) 내부 및 클러스터 간 거리 변화를 훈련 과정에서 추적.
    3. 응용 실험:
       • 레이블 없는 난이도 순위: 분산을 통해 데이터의 난이도 (예측 오차 가능성) 를 순위 매기기.
       • 모델 선택: 라벨 없이 사전 훈련된 여러 체크포인트 중 성능이 좋은 모델을 선별.
       • kNN-LM 레이어 선택: 검색 증강을 위해 어떤 내부 레이어의 표현을 키 (key) 로 사용할지 결정.
       • 훈련 신호: 분산을 증가시키는 보조 손실 함수 ("push-away" loss) 를 도입하여 모델 성능 향상.

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

A. 핵심 발견: 분산과 퍼플렉시티의 강한 음의 상관관계

• 관찰: 모델이 텍스트를 더 잘 예측할수록 (퍼플렉시티가 낮을수록), 그 표현 공간의 분산 (평균 코사인 거리) 은 더 큽니다.
• 데이터: LLaMA, Qwen 등 다양한 모델과 도메인에서 일관되게 관찰됨 (피어슨 상관계수 -0.6 ~ -0.92).
• 해석: 약한 모델은 유사한 문맥을 좁은 클러스터로 압축하는 반면, 강한 모델은 의미적으로 유사한 문맥조차도 임베딩 공간에서 더 넓게 분리하여 더 명확한 차이를 만듭니다.
• 레이어별 패턴: 초기 레이어보다는 깊은 레이어에서 이 상관관계가 더 뚜렷하게 나타납니다.

B. 실용적 응용 (Label-Free Applications)

1. 레이블 없는 데이터 난이도 순위 (Ranking Example Hardness):
   • 정답 레이블 없이도 분산을 측정하여 모델이 틀릴 가능성이 높은 "어려운 슬라이스 (hard slices)"를 식별할 수 있습니다.
   • 정확도가 높은 샘플일수록 분산이 크고, 낮은 샘플일수록 분산이 작다는 단조로운 관계가 확인됨.
2. 모델 선택 (Model Selection):
   • 동일한 모델 패밀리 내의 다양한 파인튜닝 버전 (예: 수학 전문, 코드 전문) 중 성능이 좋은 모델을 선별할 때, 라벨 없이 **분산 격자 (Dispersion Gap)**를 측정하여 효율적으로 스크리닝할 수 있습니다.
   • 도메인 관련 토큰 간의 분산과 일반 토큰과의 거리를 합산한 지표가 하류 작업 정확도와 0.95 이상의 높은 상관관계를 보임.
3. kNN-LM 을 위한 레이어 선택:
   • kNN-LM 에서 데이터스토어 키로 사용할 내부 레이어를 선택할 때, Attention 서브레이어 출력이 Feed-Forward 서브레이어보다 분산이 더 크고, 이는 더 낮은 퍼플렉시티로 이어짐을 발견했습니다.
   • 이를 통해 전체 파이프라인을 실행하지 않고도 분산 측정만으로 최적 레이어를 선별 가능.
4. 훈련 신호로서의 활용 (Training Signal):
   • 표준 크로스 엔트로피 손실에 보조 "밀어내기 (push-away)" 손실 (은닉 벡터 간 평균 코사인 거리를 최대화) 을 추가했습니다.
   • 단일 도메인 및 크로스 도메인 (위키텍스트 + 코드) 훈련 모두에서 분산을 증가시키는 것이 퍼플렉시티를 직접적으로 개선시켰습니다. 특히 이질적인 데이터 소스를 다룰 때 효과가 두드러졌습니다.

C. 인과성 검증

• 의도적 압축 실험: 분산을 줄이는 ("Squeeze") 보조 손실을 추가하여 모델을 훈련시켰을 때, 퍼플렉시티가 유의미하게 악화됨을 확인했습니다. 이는 분산이 단순한 상관관계가 아니라 모델 성능의 기능적 동인 (driver) 일 가능성을 시사합니다.

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

• 개념적 통찰: 언어 모델의 성능은 단순히 토큰 예측 확률뿐만 아니라, **임베딩 공간의 기하학적 구조 (넓이)**와 밀접하게 연관되어 있음을 증명했습니다. 더 넓은 기하학적 구조는 더 낮은 엔트로피와 더 정확한 다음 토큰 예측을 가능하게 합니다.
• 실용적 가치:
  • 비용 효율성: 라벨링 없이도 모델의 성능을 진단하고, 데이터를 선별하며, 최적의 모델을 선택할 수 있는 저비용 도구를 제공합니다.
  • 훈련 최적화: 분산을 유도하는 간단한 정규화 기법을 통해 모델의 일반화 능력과 도메인 적응력을 향상시킬 수 있습니다.
• 미래 방향: 이 연구는 임베딩 기하학을 제어하여 모델의 강건성 (robustness) 과 해석 가능성 (interpretability) 을 높이는 새로운 연구 방향을 제시합니다.

요약하자면, 이 논문은 **"모델이 얼마나 잘 예측하는지는 그 내부 표현이 얼마나 넓게 퍼져 있는지로 예측할 수 있으며, 이를 통해 모델 평가, 선택, 훈련 전반에 걸쳐 효율적인 개선 전략을 수립할 수 있다"**는 핵심 주장을 데이터와 실험을 통해 입증했습니다.