엔트로피와 CKA로 추측 디코딩에 쓸 LLM 층을 고른다
Surfing the Activation Wave: Analyzing the Information and Geometry of LLM Hidden States for Speculative Decoding
무엇인가
추측 디코딩(speculative decoding)은 작은 드래프트 모델이 여러 토큰을 미리 생성하고 타깃 모델이 이를 한 번에 검증하는 방식으로 생성 지연시간을 줄인다. EAGLE3 계열 구현에서는 드래프트 모델이 타깃 모델의 은닉 상태(hidden state)를 입력으로 받아 쓰는데, 여기서 실무적으로 골치 아픈 질문이 하나 남는다. 깊이가 N인 타깃 모델에서 어느 층의 은닉 상태를 뽑아 쓸 것인가. 기존 방식은 [2, N//2, N-3]처럼 위치 기반의 고정 규칙으로 몇 개 층을 집는 것이 사실상 표준이었다. 그런데 층의 부분집합을 고르는 일은 본질적으로 조합 최적화 문제이고, 적당히 깊은 모델만 되어도 모든 조합을 평가하는 것은 불가능하다. 이 논문은 그 선택을 감이나 관성 대신 측정 가능한 두 가지 축으로 정당화하려 한다.
어떻게 동작하나
제안하는 프레임워크의 두 축은 정보성(informativeness)과 다양성(diversity)이다. 정보성은 행렬 엔트로피(matrix-entropy) 분석으로 측정한다. 어떤 층의 은닉 상태가 담고 있는 정보량이 얼마나 큰지를 엔트로피 관점에서 정량화하는 것이다. 다양성은 CKA(centered kernel alignment)로 잰다. CKA는 두 표현 사이의 유사도를 재는 지표이므로, 두 층의 은닉 상태가 서로 얼마나 겹치는지, 즉 중복(redundancy)이 얼마나 큰지를 알려준다. 결국 목표는 정보량은 많으면서 서로 겹치지 않는 층 집합, 다시 말해 '많이 담고 있으면서 최소한으로 중복되는' 조합을 찾는 것이다.
무엇과 다른가
이 기준을 세우기에 앞서 저자들은 디코더 표현의 구조 자체를 먼저 관찰한다. 층별 섀넌 엔트로피(Shannon entropy), 은닉 상태를 이어 붙였을 때(concatenation)의 총합 엔트로피, 그리고 깊이에 따른 쌍별 CKA의 변화 양상을 함께 측정한다. 그 결과 디코더 표현 안에 기하학적으로나 정보적으로 구분되는 뚜렷한 구획(compartment)들이 존재한다는 점을 보고한다. 이 관찰이 이후 층 선택 알고리즘의 근거가 된다.
어떻게 쓰나
이 관찰 위에서 저자들은 엔트로피와 CKA를 함께 쓰는 층 선택 알고리즘을 제안한다. 알고리즘은 앞서 정의한 두 기준을 만족하는 은닉 상태 집합을 선호하도록 설계된다. 여기에 더해 토큰 엔트로피 필터링(token-entropy filtering)이라는 구성 요소를 도입한다. 예측 엔트로피가 낮은 토큰, 즉 모델이 거의 확신하고 있어 정보가 별로 없는 토큰에 대응하는 은닉 상태를 통계에서 제거하는 것이다. 이 필터링을 거치면 층 선택에 쓰이는 통계 자체의 품질이 올라간다는 것이 저자들의 주장이다.
전제와 한계
실험은 EAGLE3 스타일의 추측 디코딩 설정에서 이뤄진다. 저자들이 선택한 층 집합은 고정 베이스라인 선택들, 그중에서도 기본 규칙인 [2, N//2, N-3]을 능가한다고 보고된다. 동시에 이 프레임워크는 왜 어떤 기본값들이 여전히 경쟁력을 유지하는지에 대한 설명도 제공한다. 다만 제공된 원문에는 실험 표와 구체적인 수치가 포함되어 있지 않아, 사용된 데이터셋 이름이나 베이스라인별 정확한 수치, 개선폭은 원문에 제시되지 않았다. 이 부분은 본문의 실험 섹션을 직접 확인해야 한다.
개발자 입장에서 이 논문의 실용적 가치는 두 가지다. 첫째, 드래프트 모델에 어떤 층을 넣을지 더 이상 관성적인 기본값에 의존하지 않고 엔트로피와 CKA라는 측정 지표로 결정할 수 있다. 둘째, 토큰 엔트로피 필터링은 층 선택뿐 아니라 은닉 상태 통계를 수집하는 어떤 파이프라인에도 이식하기 쉬운 아이디어다. 다만 실제로 적용하려면 모델 깊이나 아키텍처가 바뀔 때마다 선택 집합을 다시 계산해야 하고, CKA와 엔트로피 계산에 드는 비용을 감수할 수 있는지 따져야 한다.
한계도 분명하다. 층 부분집합 선택이 조합 폭발 문제라는 점 자체는 해결되지 않았고, 제안된 알고리즘은 그 문제를 우회하는 휴리스틱이다. 실험은 EAGLE3 스타일의 특정 추측 디코딩 설정에 국한되어 있으며, 저자들 스스로 이 분석을 다른 응용으로 확장하는 일은 향후 과제로 남겨 둔다. 또한 이 논문은 NeurIPS 워크숍(Axiom, GDDL, PriGM)에 채택된 워크숍 논문으로, 대규모 벤치마크 전반에서의 일반화는 별도로 검증되어야 한다.