기술 요약: N-vium – 가속화된 정확한 생성을 위한 혼합 탈출 (Mixture-of-Exits) 트랜스포머

1. 문제 제기

표준 자기회귀형 대규모 언어 모델 (LLM) 은 스택 형태의 디코더 레이어를 통해 토큰을 순차적으로 처리합니다. 이 순방향 전파는 (1) 미래 어텐션을 위해 모든 레이어에서 키 - 값 (KV) 표현을 구축하고, (2) 최종 은닉 상태를 투영하여 다음 토큰을 예측하는 두 가지 결합된 기능을 수행합니다. 결과적으로, 예측 분포가 이미 수렴했는지 여부와 상관없이 모든 토큰은 모든 레이어를 통한 완전한 순방향 전파를 필요로 합니다.

조기 탈출 (Early Exit) 방법과 같은 기존 추론 효율성 개선 접근법들은 예측 헤드를 중간 레이어에 부착함으로써 이러한 기능을 분리하려 시도합니다. 그러나 이러한 방법들은 일반적으로 중간 예측을 최종 레이어 분포의 근사치로 취급합니다. 이는 다음과 같은 여러 단점을 초래합니다:

• 근사 오차: 중간 예측은 종종 전체 깊이 예측보다 약하여 훈련 - 테스트 불일치를 만듭니다.
• 휴리스틱 라우팅: 언제 탈출할지 결정하는 것은 종종 보정이 필요한 수작업 신뢰도 임계값에 의존합니다.
• 불완전한 KV 캐시: 연산을 절약하기 위해 조기 탈출 방법들은 우회된 레이어에 대한 KV 캐시 계산을 종종 생략하며, 이는 근사적 복사나 생략을 요구하여 품질을 저하시킵니다.
• 품질 저하: 속도는 종종 퍼플렉시티와 생성 품질의 대가를 치르고 달성됩니다.

본 논문은 "로짓 렌즈 (logit lens)"를 통한 경험적 관찰 (예: 예측 분포가 종종 최종 레이어 이전에 안정화됨) 이 있음에도 불구하고, 현재 방법들은 품질 손실 없이 정확한 생성을 위해 이를 활용하지 못한다고 주장합니다.

2. 방법론: N-vium

저자들은 표준 트랜스포머 아키텍처를 엄격하게 일반화하는 N-vium을 소개합니다. 이는 조기 탈출을 근사치로 취급하는 대신, 모델의 출력을 여러 탈출 지점의 학습된 혼합 분포로 재정의합니다.

아키텍처

• 혼합 탈출 (Mixture of Exits): 모델은 백본의 $L$개 레이어를 $N_V$개의 블록으로 분할합니다. 각 블록 (접합점 $k$) 의 끝에서 라우터는 해당 레이어의 분포 $\pi_k$에서 다음 토큰을 샘플링할 확률을 결정합니다.
• 토큰 적응형 라우팅: 라우터 MLP 는 접합점 $k$에서의 은닉 상태 $h^{(\ell_k)}_t$를 처리하여 확률 $w_k$를 출력합니다. 접합점 $k$에서 탈출할 확률 $p_k$는 다음과 같이 계산됩니다:
  $$p_k(x_t) = w_k(h^{(\ell_k)}_t) \prod_{j=1}^{k-1} (1 - w_j(h^{(\ell_j)}_t))$$
  여기서 관례에 따라 $w_{N_V} = 1$입니다.
• 혼합 분포: 최종 다음 토큰 분포는 모든 탈출 분포의 가중합입니다:
  $$\pi_{mix}(x_{t+1} | x_{1:t}) = \sum_{k=1}^{N_V} p_k(x_t) \pi_k(x_{t+1} | x_{1:t})$$
• 공유 헤더: 모든 탈출 지점은 최종 언어 모델링 (LM) 투영 헤더 $W_{lm}$을 공유하며, 경량 어댑터 MLP 가 중간 은닉 상태를 헤더의 기대 기하학적 구조에 맞게 변환합니다.

훈련

• 목적 함수: 모델은 개별 헤드를 최종 레이어와 일치시키기 위해 고립적으로 최적화하는 대신, 결합 혼합 분포 $\pi_{mix}$의 교차 엔트로피 손실을 최소화하도록 훈련됩니다.
• 연산 페널티: 라우터가 초기 탈출 지점에 확률 질량을 할당하도록 장려하여 모델을 효율성 쪽으로 유도하는 보조 손실 항 $\mathcal{L}_{compute}$가 추가됩니다:
  $$\mathcal{L}_{total} = \mathcal{L}_{mix} + \beta \mathcal{L}_{compute}$$
  여기서 $\mathcal{L}_{compute}$는 기대 정규화 레이어 깊이입니다.
• 워밍업: 초기에 모든 토큰이 최종 레이어로 라우팅되는 "라우터 붕괴"를 방지하기 위해 균일한 탈출 분포에서의 편차를 패널티로 부과하는 라우터 워밍업 단계가 사용됩니다.

추론: 피기백 (Piggybacking)

계산을 생략하지 않고 정확한 KV 캐시 복구를 보장하는 피기백 기술이 핵심 혁신입니다:

1. 연산 연기: 토큰이 접합점 $k$에서 탈출할 때, $k+1$부터 $L$까지의 레이어에 대한 계산은 연기됩니다.
2. 병렬 배치: 미래 토큰이 처리됨에 따라, 이들의 순방향 전파는 연기된 레이어에 도달합니다. 이러한 깊이에서 연기된 토큰들의 은닉 상태는 현재 토큰의 상태와 연결되어 병렬로 처리됩니다.
3. 결과: 이는 모든 연기된 토큰에 대한 누락된 KV 상태를 단일 병렬 패스로 채워, KV 캐시가 완전하고 정확하며 샘플링된 토큰과 일관되도록 보장합니다. 계산이 생략되는 것이 아니라 순차적 실행에서 병렬 배치 실행으로 재배열될 뿐입니다.

3. 주요 기여

1. 정확한 샘플링 혼합: N-vium 은 표준 조기 탈출 방법의 고유한 근사 오차를 피하면서 학습된 분포에서 정확한 샘플링을 가능하게 하는 토큰 적응형 깊이 라우팅을 지원하는 첫 번째 방법입니다.
2. 완전한 KV 복구: 상부 레이어 계산을 연기하고 배치함으로써 N-vium 은 다른 적응형 계산 방법에서 종종 손실되는 완전하고 정확한 KV 캐시를 복구합니다.
3. 훈련 프레임워크: 저자들은 학습 가능한 연산 페널티를 통해 모델을 효율성 쪽으로 유도하는 전략을 포함하여 혼합 목적 함수를 직접 최적화하는 실용적인 훈련 파이프라인을 제공합니다.
4. 확장성: 이 방법은 최대 15 억 파라미터 규모까지 검증되어, 이 접근 방식이 소형 모델에 국한되지 않음을 입증했습니다.

4. 실험 결과

저자들은 LLaMA 2 아키텍처를 사용하여 C4 데이터셋에서 N-vium 모델 (특히 4 개의 탈출 지점을 가진 "Quadrivium") 을 사전 훈련하여 파라미터 및 데이터가 일치하는 조밀한 베이스라인과 비교했습니다.

• 속도 향상 vs 퍼플렉시티:
  • 가장 큰 모델 (15 억 파라미터, 48 레이어, 1536 은닉 폭) 은 퍼플렉시티 저하 없이 (0% 차이) 조밀한 베이스라인 대비 57.9% 의 월클럭 속도 향상을 달성했습니다.
  • 더 작은 모델들은 23% 에서 41% 사이의 일관된 속도 향상을 보였으며, 약간의 퍼플렉시티 개선을 동반했습니다 (예: 36 레이어 모델은 3.7% 퍼플렉시티 개선과 함께 36.4% 속도 향상을 보임).
• 파레토 프론티어: 품질 - 속도 플롯에서 N-vium 은 "상단 - 우측" 사분면 (더 빠르고 동등하거나 더 나은 품질) 을 차지하는 반면, CALM(재구현된 CALM*) 과 같은 방법들은 "하단 - 우측"(더 빠르지만 품질이 떨어짐) 에 위치하고, LayerSkip 은 베이스라인 아래에 있는 일정한 퍼플렉시티 선 위에 있습니다.
• 하류 성능: Tulu 3 혼합 데이터셋에 대한 지도 미세 조정 (SFT) 후, N-vium 모델들은 하류 벤치마크 (MMLU, ARC, HellaSwag 등) 에서 조밀한 베이스라인과 일치하거나 초과하여, 모델이 여전히 제어 가능하며 작업 성능을 희생하여 교차 엔트로피 목적 함수에 과적합되지 않음을 확인했습니다.
• 하드웨어 견고성: A100, RTX A6000, RTX 3090 GPU 에서 측정한 속도 향상 수치는 일관된 결과 (~57%) 를 보였으며, 네이티브 bfloat16 지원이 없는 구형 Turing 아키텍처 GPU 는 더 낮은 이득을 보였습니다.

5. 중요성과 주장

본 논문은 N-vium 이 추론 효율성에서 패러다임 전환을 나타낸다고 주장합니다. 대부분의 방법이 토큰당 FLOP 를 감소시킴으로써 (종종 근사를 통해) 초당 토큰 수를 개선하는 반면, N-vium 은 초당 FLOP 를 증가시킴으로써 효율성을 개선합니다. 이는 순차적 계산을 병렬 배치 패스로 재구성하여 어떤 계산도 폐기하지 않고 하드웨어 병렬성을 더 잘 활용함으로써 이를 달성합니다.

저자들은 N-vium 이 표준 트랜스포머의 엄격한 일반화임을 강조합니다. 중간 탈출 지점에 대한 라우팅 가중치가 0 일 때, 모델은 표준 트랜스포머와 정확히 동일하게 동작합니다. 이 방법은 저자들이 기존 다른 어떤 방법도 동시에 충족하지 않는다고 주장하는 네 가지 중요한 기준을 충족합니다:

1. 정확한 추론: 근사 없이 학습된 분포에서 샘플링.
2. 완전한 KV 캐시 복구: 생략된 레이어나 근사적 캐시 채움 없음.
3. 학습된 라우팅: 토큰별 계산은 휴리스틱이 아닌 훈련된 라우터에 의해 할당됨.
4. 디코드 시간 속도 향상: 자기회귀 생성 중 월클럭 지연 시간 감소.

논문은 N-vium 훈련이 여러 접합점에서 헤드를 평가하기 때문에 토큰당 더 높은 FLOP 비용을 수반하지만, 추론 비용이 지배적인 생산 배포에서는 이 오버헤드가 정당화되며 모델 크기가 커질수록 상대적 오버헤드는 감소한다고 결론지었습니다.