자체 추측 디코딩을 통한 더 빠른 텍스트 생성

LayerSkip 방법을 통해 구현된 자체 추측 디코딩은 자체 초기 레이어를 사용하여 토큰을 초안 작성하고 더 깊은 레이어를 사용하여 검증함으로써 단일 대형 언어 모델(LLM)의 텍스트 생성 속도를 높입니다. 이 접근 방식은 두 개의 별도 모델이 필요한 전통적인 추측 디코딩과 관련된 메모리 사용량과 계산 지연을 줄입니다.

자체 추측 디코딩 대 전통적 추측 디코딩

전통적인 추측 디코딩은 토큰 시퀀스를 생성하기 위해 작은 초안 모델에 의존하고, 이를 검증하기 위해 더 큰 검증 모델에 의존합니다. 반면, 자체 추측 디코딩은 두 역할에 동일한 모델을 사용합니다. 중간 레이어에서 모델을 종료하여 토큰 초안을 생성함으로써, 시스템은 보조 모델의 가중치가 필요 없게 되어 메모리를 절약하고 대규모 추론에 대한 전체 하드웨어 발자국을 줄입니다.

기술적 구현: 조기 종료 및 언임베딩

자체 추측 디코딩을 가능하게 하려면 모델이 사전에 지정된 중간 레이어에서 생성 프로세스가 중단되는 '조기 종료' 기능을 갖추어야 합니다.

언임베딩 로짓

언어 모델(LM) 헤드는 일반적으로 최종 레이어만을 위해 훈련되므로, 중간 레이어의 로짓을 '언임베딩'하기 위해 특정 훈련이 필요합니다. 이 과정은 중간 레이어의 출력을 LM 헤드에 투영하여 다음 토큰을 예측합니다.

훈련 수정 사항

사전 훈련 또는 파인튜닝 중에 조기 종료를 가능하게 하는 두 가지 주요 수정 사항이 사용됩니다:

  • Layer Dropout: 모델은 특정 레이어를 건너뛰도록 훈련되며, 더 깊은 레이어일수록 드롭아웃 비율이 증가합니다.これにより後続 레이어에 대한 의존도가 줄어들고 일반화가 향상됩니다.
  • Early Exit Loss: 각 종료(중간 레이어)에 정규화된 손실이 적용되어, LM 헤드가 모델 전체의 다양한 깊이에서의 출력을 언임베딩하는 방법을 배우도록 강제합니다.

추론 과정: 초안 작성 및 검증

추론은 두 가지 구분된 단계에서 발생합니다:

  1. Self-Drafting: 중간 레이어에서 모델을 종료하여 토큰을 생성합니다. 추측 토큰의 수와 특정 종료 레이어는 속도와 정확도의 균형을 맞추는 하이퍼파라미터입니다.
  2. Self-Verification: 전체 모델이 초안 토큰을 검증합니다. 이를 최적화하기 위해 시스템은 초안 작성 단계에서 캐시된 초기 레이어의 결과를 재사용하고, 검증에 필요한 나머지 레이어만 계산합니다.

하드웨어 및 성능 최적화

자체 추측 디코딩은 중복 계산을 최소화하기 위해 KVQ 캐시(KV 캐시와 종료 쿼리 캐시의 조합)를 사용합니다.これにより 두 모델 시스템보다 세 가지 주요 장점이 제공됩니다:

  • Shared Weights: 첫 번째 $E$ 레이어는 초안 작성과 검증에 모두 재사용됩니다.
  • Shared KV Cache: 첫 번째 $E$ 레이어의 키-값 쌍이 재사용됩니다.
  • Shared Compute: 종료 쿼리 캐시는 종료 레이어 $E-1$의 쿼리 벡터를 저장하여, 검증 프로세스가 레이어 $0$부터 $E-1$까지의 계산을 건너뛸 수 있도록 합니다.

벤치마크 및 성능 트레이드오프

A100 GPU에서 수행된 벤치마크에 따르면, Llama 3.2 1B, Llama 3 8B, Llama 2 13B, Llama 2 7B 등을 포함한 대부분의 모델 크기에서 조기 종료 자체 추측 디코딩이 전통적인 두 모델 추측 디코딩보다 빠릅니다.

조기 종료의 "Sweet Spot"에 대해

종료 레이어와 성능 사이에는 비선형 관계가 있습니다. 너무 일찍 종료하면 정확도가 낮아지고 초안 토큰의 수용률이 낮아지며, 너무 늦게 종료하면 초안 단계의 계산 오버헤드가 증가합니다. 대부분의 모델에서는 예측 정확도와 생성 오버헤드 사이의 트레이드오프가 초당 최대 토큰 수를 위해 최적화되는 "Sweet Spot"이 존재합니다.

모델별 결과

  • Llama 2 70B: 자기 회귀 디코딩보다 znacz하게 빠르지만, 다른 모델에 비해 속도 향상이 제한적일 수 있습니다. 이는 Llama 2 7B 변형(52B 토큰)에 비해 토큰 수가 적은(328M 토큰) continual pretraining 때문일 수 있습니다.
  • Baseline Models: LayerSkip 레시피로 훈련되지 않은 모델은 초기 레이어가 출력을 예측하도록 훈련되지 않아 토큰 수용률이 매우 낮아 속도 향상을 보지 못합니다.

Transformers와의 통합

자체 추측 디코딩은 Hugging Face transformers 라이브러리에 통합되어 있습니다. 모델 체크포인트가 LayerSkip 레시피로 훈련된 경우, generate() 함수에 assistant_early_exit 인수를 추가하여 활성화할 수 있습니다.

Sources