Review

TL; DR

Summary

• 연구진 : 미국 켄터키대학, NYU, UCLA, Meta

• 인용수 : 89

연구 동기

• LLM은 주로 마지막 layer의 출력을 downstream task에 사용
  • “얕은 layer는 단순히 low-level 정보를 담는다”는 일반적인 가정에 기반함

• 저자들은 이 가정에 의문을 제기!
  • “마지막 layer가 항상 최고의 representation을 제공하는가?”

→ 실제로 중간 layer가 더 풍부한 표현력을 가지고 있을 수 있으며, 다양한 task에서 좋은 성능을 보여줄 수 있음을 실증적으로 확인해보자!

Key Findings

1. 중간 layer가 일관되게 마지막 layer보다 더 우수함을 실험적으로 증명

2. 모델 아키텍처(종류, 크기)와 학습 진행도에 따른 표현력 차이를 비교

3. CoT 파인튜닝이 중간 layer 표현을 풍부하게 만듦

평가지표 설계

• Representation의 quality를 어떻게 평가할 수 있는가?

  1. Representation이 얼마나 압축(compressed) 되었는지?

  2. Input이 perturbation / augmentation에 대해 얼마나 robust한지?

  3. 서로 다른 input을 어떻게 기하학적으로 구성하는지?

• Matrix-Based Entropy: 공통된 수학적 관점

  • 고유값 (Eigenvalue) $\lambda_i$를 이용해 entropy 계산

    

    • $\lambda_i$ : 데이터를 표현하는 axis인 eigenvector 방향으로 포함된 정보의 양
    • $\mathbf{Z}\in\mathbb{R}^{n\times d}$ : Representation matrix $(n개의 데이터, d 차원)$
    • $\mathbf{K}=\mathbf{Z}\mathbf{Z}^T$ : Gram matrix (Representation 간 유사도 행렬)
    • $\alpha$ : Smoothing 지수

    → Input 샘플이 몇 개의 고유값에 집중되어 있는지를 측정
    → 직관적으로 eigenvalue가 고르게 퍼져있으면 high entropy

  • Insight 1 (정보 압축 관점)

    • 고유값 중 몇 개만 큰 경우 → 낮은 차원에 정보가 집중되며 일부분의 axis로만 압축됨 → low entropy

    • 고유값이 고르게 분포한 경우 → 정보가 여러 axis에 비슷하게 분산됨 → high entropy

  • Insight 2 (Geometry 관점)

    • 토큰 임베딩이 부드럽게 이어진 경로를 따라가면 → 곡률(curvature)이 낮고 high entropy

    • 갑작스럽게 꺾이면 (즉, 연속된 토큰의 임베딩 방향이 급변하면) → 곡률이 높고 고유값 분포가 한쪽으로 쏠려 low entropy

    → Embedding 곡률도 결국 고유값 분포(=entropy)로 반영 가능

  • Insight 3 (Input perturbation/augmentation에 따른 Robustness 관점)

    • Strong invariance(=robust) → 같은 의미의 샘플이 embedding 공간에서 안정적으로 클러스터링됨 → 엔트로피 유지

• 7가지 Representation 평가 지표

  • 정보이론 기반

    • Prompt Entropy
      • 하나의 프롬프트 안에서 토큰 임베딩이 얼마나 다양하게 퍼져 있는가?
        • 높은 entropy → 표현이 다양함 → 덜 중복되고 풍부한 특징
        • 낮은 entropy → 표현이 비슷함 → 정보가 압축됨

    • Dataset Entropy
      • 여러 프롬프트 임베딩이 데이터셋 전반에서 얼마나 다양하게 퍼져 있는가?
        • 높은 entropy → 서로 다른 프롬프트 간 표현이 잘 구별됨
        • 낮은 entropy → 입력에 상관없이 표현이 유사해짐 (정보 손실 가능성)

    • Effective Rank
      • 표현 공간이 실제로 몇 차원으로 구성되나?
        • 값이 낮을수록 → 대부분 정보가 소수 차원에 압축됨
        • 값이 높을수록 → 정보가 고르게 퍼짐

  • 기하학 기반

    • Curvature

      

      • 곡률을 계산하여 연속된 토큰 벡터들 간의 방향 전환 정도 측정
        • 높은 curvature → 임베딩 벡터 방향이 급격히 전환됨 → local 정보 강조
        • 낮은 curvature → 완만한 궤적 → global 문맥 강조

  • 변형 불변성 기반

    • InfoNCE
      • 같은 의미의 입력 pair가 임베딩 공간에서 얼마나 서로 가까운지 측정
      • InfoNCE loss가 낮을수록 robust함

    • LiDAR
      • Augmentation 전후 임베딩이 얼마나 잘 클러스터링되는지 측정
      • LiDAR 스코어가 높을수록 robust함

    • DiME
      • Augmented pair가 random pair 대비 얼마나 잘 정렬되어있는지 측정
      • 높을수록 정렬 잘됨

Experiments

• Downstream Task Performance

  “마지막 layer가 항상 최선인가?”에 대한 실증적 검증

  

  • 비교모델
    • Pythia : Decoder-only Transformer
    • Mamba : State Space Model
    • BERT-base : Encoder-only Transformer

  • 벤치마크 : MTEB (Massive Text Embedding Benchmark)
    • 32개 Task: Span classification, Semantic textual similarity, clustering, reranking…

  → 거의 모든 태스크에서 중간 layer가 마지막 layer보다 더 높은 성능을 기록

• Metrics vs. Performance Correlation

  앞서 정의한 representation 품질에 대한 평가 metric들이 실제로 성능과 관련 있는지 검증 (dCor metric 활용)

  • dCor (distance Correlation) $\in[0,1]$ : Linear & Non-linear 모두에 대해 두 행렬 (representation 집합) 사이의 거리를 기반으로 의존성을 측정하는 지표

  

  → Label 없이도 "가장 좋은 표현을 가진 layer"를 찾을 수 있다!

• Architectural and Scale Differences

  모델 아키텍처와 크기에 따라 representation 품질 변화 패턴이 다른지 파악

  • 아키텍쳐에 따른 품질 변화

    

    • BERT : 전체적으로 고른 정보량 유지→ bottleneck 적음

    • Pythia : 중간 layer에서 bottleneck 발생

    • Mamba : 부드럽게 변화하는 형태 (BERT와 Pythia의 중간 정도)

  • 모델 크기에 따른 품질 변화

    

    • 모델 크기가 커질수록 (14M → 1B) entropy valley가 더 뚜렷해지며, curvature 값이 낮아지고, invariant해진다!

    → Feature 표현이 더 부드럽게 정제되고, robust해짐

  • 파인튜닝 효과

    

    • Unsupervised 방식 → 다양한 토큰 representation 생성을 통해 entropy 증가

    • 방식에 상관없이 파인튜닝 자체가 invariance 개선!

• Impact of Training Progression

  학습이 진행되면서 각 layer에서의 representation이 어떻게 변화하는지 파악

  

  • 초기 layer : 학습 진행률과 관계없이 representation이 안정됨

  • 중간 layer : 학습이 진행될수록 entropy 감소 & LiDAR score 최저점
    • 곡률은 smooth해짐 → 학습하면서 전역적인 의미 구조를 반영하기 시작함!

• Impact of CoT Finetuning

  CoT 파인튜닝 적용 시 entropy는 어떻게 변화하는 지 파악

  

  • 파인튜닝된 모델이 더 높은 entropy를 보임
    • 모델이 문제를 여러 단계로 나누어 표현하려고 시도하고, 이에 따라 표현 자체가 넓고 다양해지기 때문

• Extreme Input Conditions

  비정상적인 프롬프트가 주어질 때 모델 반응 확인

  (Random 텍스트 & 텍스트 반복)

  

  • Repetition → 의미 없는 반복을 즉시 감지하고, 중간 layer에서 표현을 단순화/무시함

  • Random 텍스트 → 초기 layer가 노이즈에 민감하여 무작위 input에 표현 다양성 증가시킴 → 중간 layer에서는 압축됨

  • Prompt 길이가 길수록 표현 압축 효과는 미미해짐

• Comparison to Vision Transformers

  비전 모델도 경향성이 유사한가?

  

  

  • Autoregressive Image Model (AIM) (위 figure에서 파랑, 주황) : GPT처럼 순차 예측 진행하는 모델 → Pythia와 같은 decoder-only Transformer 처럼 중간 layer에서 entropy가 가장 낮음

  • 반면 ViT (핑크)와 같이 non-autoregressive 방식의 학습을 활용하는 모델은 steady한 곡선을 보임 (BEiT 제외)
    • Non-autoregressive 모델은 중간에서 정보를 압축할 필요가 적다!

  → 결론 : 학습 방식 (autoregressive 여부)의 차이가 표현력 차이를 이끌어낸다!