SimLens for Early Exit in Large Language Models: Eliciting Accurate Latent Predictions with One More Token

대형 언어 모델 (LLM) 은 깊이가 깊어질수록 점진적으로 예측을 형성합니다. 중간 계층 (Intermediate Layer) 의 은닉 상태 (Hidden State) 에서 모델의 잠재적 예측을 정확하게 읽어낼 수 있다면, 추론 비용을 줄이고 (Early Exit) 모델 내부의 의미론적 형성 과정을 이해하는 데 도움이 됩니다.

그러나 기존 연구들은 다음과 같은 한계가 있었습니다:

• 기존 Lens 방법의 한계: Logit Lens, Tuned Lens 등 기존 방법들은 중간 계층의 상태를 직접 선형 변환 (Linear Readout) 을 통해 해독합니다. 이는 간단하지만, 특히 초기 계층에서는 모델의 최종 예측과 크게 벗어나는 (Drift) 경향이 있어 신뢰도가 낮습니다.
• 정보 부재가 아닌 해독의 어려움: 초기 계층에도 정답에 필요한 정보가 이미 존재함에도 불구하고, 선형 변환만으로는 이를 정확하게 읽어내지 못하는 것이 주요 병목 현상임을 발견했습니다.

2. 방법론 (Methodology)

SimLens (Simple Training-free Decoder)

저자들은 더 강력한 선형 프로젝터를 설계하는 대신, **"단 하나의 추가 토큰"**을 통해 더 정확한 예측을 이끌어낼 수 있는지 제안합니다.

• 핵심 아이디어: 중간 계층에서 전체 시퀀스를 계속 진행하는 대신, **시작 토큰 <s>**와 후보 정답 토큰 <a> 두 개만 남긴 상태에서 모델의 나머지 상위 계층 (Upper Layers) 을 가볍게 통과시킵니다.
• 동작 원리:
  1. 중간 계층 $l$에서 <s>와 <a>로 구성된 2 토큰 상태를 형성합니다.
  2. 이 상태를 모델의 마지막 계층 $L$까지 순전파 (Forward pass) 시킵니다.
  3. 최종 계층의 출력을 통해 <a>의 확률을 계산하고, 가장 높은 점수를 받은 후보를 선택합니다.
• 역할:
  • <s>: 전역 조건 (Global Conditioning) 과 위치/컨텍스트 앵커 역할을 수행 (Attention Sink 현상 활용).
  • <a>: 평가 대상인 후보의 의미론적 앵커 (Semantic Anchor) 역할을 수행.

Linear SimLens (경량화 근사)

Early Exit 시 빈번한 신뢰도 확인을 위해 SimLens 를 선형적으로 근사화한 방법입니다.

• SimLens 의 출력을 타겟으로 하여, 중간 계층의 표현을 최종 계층의 표현으로 매핑하는 선형 변환 행렬을 학습합니다.
• 전체 시퀀스 전파 없이 엔트로피 기반의 신뢰도 점수 (Confidence Score) 를 빠르게 계산하는 데 사용됩니다.

SimExit (하이브리드 Early Exit 프레임워크)

SimLens 와 Linear SimLens 를 결합한 조기 종료 메커니즘입니다.

• 신뢰도 평가: 각 계층에서 Linear SimLens를 사용하여 엔트로피 (Entropy) 기반의 신뢰도 점수를 계산합니다.
• 종료 및 해독: 신뢰도 임계값을 충족하면 전체 시퀀스 전파를 중단하고, SimLens를 사용하여 축소된 시퀀스 (2 토큰) 로만 최종 정답을 도출합니다.
• 알고리즘: 평가 간격 ($N$) 을 설정하여 모니터링 오버헤드와 종료 시점의 균형을 맞춥니다.

3. 주요 기여 (Key Contributions)

1. SimLens 도입: 학습이 필요 없는 (Training-free) 간단한 디코더로, 2 토큰 연속 전파를 통해 중간 계층의 정확한 잠재 예측을 유도합니다.
2. 토큰 역할 규명: <s>(전역 조건) 와 <a>(의미론적 앵커) 가 각각 고유한 역할을 하며, 이를 유지하는 것이 초기 계층 읽기 정확도를 획기적으로 높인다는 것을 증명했습니다.
3. SimExit 프레임워크: 저비용 엔트로피 모니터링 (Linear SimLens) 과 정확한 축소 시퀀스 디코딩 (SimLens) 을 결합한 적응형 조기 종료 메커니즘을 제안했습니다.
4. 공정한 비교 분석: 추가 2 토큰 전파 오버헤드를 포함한 '공정 계산 (Fair Compute)' 기준에서 기존 방법보다 우수한 성능을 입증했습니다.

4. 실험 결과 (Results)

LLaMA-7B 와 Vicuna-7B 모델을 ARC, BoolQ, HeadQA 데이터셋에서 평가했습니다.

• 정확도 향상 (Iso-Compute):
  • 모든 6 가지 설정 (모델 x 데이터셋) 에서 SimLens 는 기존 선형 기반 방법 (Logit Lens, Tuned Lens 등) 보다 높은 정확도를 보였습니다.
  • 평균적으로 **+0.43%**의 정확도 향상을 기록했습니다 (공정한 계산 비용 포함).
  • 특히 BoolQ 와 HeadQA 에서 초기 계층의 정확도 개선 폭이 매우 컸습니다.
• 속도 향상 (SimExit):
  • 최적의 정확도 운영점에서 평균 1.15 배의 속도 향상을 달성했습니다.
  • 정확도를 1%p 까지 허용할 경우, 평균 1.40 배의 속도 향상을 얻었습니다.
• 일반화 능력:
  • ARC 에서 학습된 Linear SimLens 를 BoolQ, HeadQA 등 다른 데이터셋에 재학습 없이 적용했을 때도 우수한 성능을 유지했습니다.
• Ablation Study:
  • <s>를 제거하면 평균 정확도가 10.68%p 하락했습니다.
  • <a>를 잘못된 토큰으로 교체하면 정확도가 13.24%p 하락했습니다. 이는 두 토큰이 모두 필수적임을 보여줍니다.

5. 의의 및 결론 (Significance)

• 선형적 한계를 넘어선 통찰: 중간 계층의 정보를 읽어내는 데 있어 단순한 선형 변환보다는, 모델의 비선형성을 활용하는 **매우 적은 양의 추가 전파 (One More Token)**가 훨씬 효과적임을 증명했습니다.
• 실용적인 효율성: SimExit 은 추론 비용을 줄이면서도 높은 정확도를 유지할 수 있는 실용적인 Early Exit 솔루션을 제공합니다.
• 분석적 관점: 이 현상을 '공간 정렬 (Space Alignment)' 또는 '추가 비선형성 (Extra Non-linearity)'의 관점에서 해석하며, 단순한 메커니즘 주장보다는 분석적 통찰을 제공합니다.
• 한계 및 향후 과제: 현재는 단일 토큰 결정 작업 (객관식 등) 에 최적화되어 있으며, 긴 생성 (Long-form generation) 으로 확장 시 KV 캐시 백필 (Backfilling) 비용이 병목이 될 수 있음을 지적했습니다.

이 논문은 LLM 의 중간 표현을 활용하여 효율성을 극대화하는 새로운 패러다임을 제시하며, 단순한 구조적 변경으로도 모델의 잠재 능력을 크게 끌어올릴 수 있음을 보여줍니다.