Medusa Heads

Medusa 헤드는 대규모 언어 모델의 마지막 은닉 레이어에 추가하는 경량 디코딩 헤드로, 각 헤드는 모델 자체 출력 레이어보다 더 앞선 토큰을 예측하도록 학습됩니다. 한 번의 순전파에서 모델은 평소처럼 다음 토큰을 예측하는 동시에 그 이후 몇 개 위치의 토큰도 추측합니다. 다음 순전파에서 트리 어텐션 단계가 그 추측을 확인하고 모델이 동의하는 가장 긴 연속 구간을 유지합니다. Medusa는 별도의 초안 모델이 필요 없는 투기적 디코딩 방식이므로 학습하거나 호스팅하거나 동기화해야 하는 두 번째 네트워크가 없습니다. Tianle Cai, Tri Dao와 동료들은 2024년 1월 Medusa: 여러 디코딩 헤드를 사용한 간단한 LLM 추론 가속 프레임워크에서 이를 소개했습니다.

헤드가 앞선 토큰을 예측하는 방식#

일반적인 텍스트 생성 단계에서는 위치 t의 마지막 은닉 상태를 언어 모델 헤드에 입력해 t + 1의 토큰 점수를 계산합니다. Medusa 헤드 k는 같은 은닉 상태를 읽고 t + k + 1의 토큰 점수를 계산합니다. 각 헤드는 SiLU 활성화와 잔차 연결을 포함하는 단일 피드포워드 레이어 뒤에 어휘 투영을 적용합니다. 투영 레이어는 언어 모델 헤드의 복사본으로 초기화하고 피드포워드 레이어는 0으로 초기화하므로 학습 전에는 모든 헤드가 기본 모델의 다음 토큰 예측을 그대로 재현합니다. 논문에서는 헤드 5개면 충분하다고 설명합니다. 이 PyTorch 예시는 헤드 3개를 만들고 하나의 은닉 상태에 대한 출력을 기본 모델 헤드의 출력과 비교합니다:

import torch
import torch.nn.functional as F

hidden_size, vocab_size, num_heads = 256, 1000, 3
lm_head = torch.nn.Linear(hidden_size, vocab_size, bias=False)  # 고정된 모델의 출력 레이어

# 각 Medusa 헤드: 0으로 초기화한 잔차 SiLU 레이어 하나와 LM 헤드의 복사본
w1 = [torch.nn.Linear(hidden_size, hidden_size) for _ in range(num_heads)]
w2 = [torch.nn.Linear(hidden_size, vocab_size, bias=False) for _ in range(num_heads)]
for layer, out in zip(w1, w2):
    torch.nn.init.zeros_(layer.weight)
    torch.nn.init.zeros_(layer.bias)
    out.weight.data.copy_(lm_head.weight.data)

# 모든 헤드는 같은 마지막 은닉 상태를 읽으며, 헤드 k는 위치 t + k + 1을 대상으로 합니다.
h = torch.randn(1, hidden_size)
head_logits = [out(h + F.silu(layer(h))) for layer, out in zip(w1, w2)]
print([torch.equal(logits, lm_head(h)) for logits in head_logits])

출력 [True, True, True]은 0으로 초기화한 잔차 레이어 덕분에 각 새 헤드가 기본 모델의 예측을 정확히 복사한 상태로 시작함을 보여줍니다. 이후 학습에서는 헤드 k가 자체적으로 더 앞선 위치를 예측하도록 조정되며, lm_head과 백본은 변경되지 않습니다.

트리 어텐션과 수락#

각 헤드의 최상위 추측을 조합해 후보 연속 토큰을 만듭니다. 헤드 1에서 상위 2개 토큰을, 헤드 2에서 상위 3개 토큰을 선택하면 데카르트 곱으로 접두부를 공유하는 후보 6개가 생성되어 트리를 이룹니다:

다음 순전파에서는 모든 후보 토큰을 입력에 한 번에 추가합니다. 어텐션 마스크는 각 토큰이 트리에서 자신의 조상만 참조하도록 하므로 모델이 모든 분기를 병렬로 점수화합니다. 모델이 수락하는 가장 긴 분기를 유지하고, 그 분기의 마지막 은닉 상태를 다음 단계의 헤드 입력으로 사용합니다. 모든 헤드가 같은 은닉 상태를 읽으므로 헤드 2는 각 헤드 1 분기 아래에서 동일한 토큰을 제안합니다. 저자들은 가장 정확한 조합으로 만든 64개 노드의 희소 트리가 256개 노드의 조밀한 트리보다 더 높은 가속률을 보인다고 밝혔습니다.

탐욕적 디코딩에서는 후보 토큰이 모델의 최상위 선택과 일치할 때 이를 수락합니다. 샘플링의 경우 Medusa는 전형적 수락 방식을 제공합니다. 이 방식은 기본 모델의 확률이 분포의 엔트로피에 따라 정해지는 임계값을 넘는 후보를 모두 유지하므로 정확한 분포 일치의 대가로 더 많은 토큰을 수락합니다.

Medusa-1 및 Medusa-2 학습#

Medusa-1은 고정된 백본을 그대로 둔 채 새 헤드만 학습하므로 모델 출력 품질을 유지하며, 논문은 생성 속도가 2.2배 이상 빨라졌다고 보고합니다. Vicuna-7B 헤드 학습에는 단일 A100 GPU와 ShareGPT 대화 60,000개를 사용해 약 5시간이 걸렸습니다. 이는 freeze=23이 Ultralytics YOLO26 파인튜닝에서 설정하는 분할 방식과 같습니다. 레이어 0~22는 고정하고 레이어 23의 탐지 헤드만 새 데이터로 학습합니다.

Medusa-2는 헤드와 백본을 함께 파인튜닝해 헤드 정확도를 높이고 속도 향상을 2.3~3.6배로 끌어올립니다. 이때 백본의 기존 능력을 보호하는 학습 레시피를 사용합니다. 기존 학습 데이터를 사용할 수 없는 경우에는 자기 증류를 통해 모델이 자체적으로 생성한 텍스트로 헤드를 학습합니다.

Medusa는 대상 모델 자체에서 토큰을 초안으로 생성하는 여러 방법 중 하나입니다. EAGLE은 작은 자기회귀 레이어를 사용해 모델의 내부 특징으로부터 초안을 만들며, DeepSeek-V3에서 사용된 다중 토큰 예측은 사전 학습 중 추가 예측 모듈을 학습하고 서빙 엔진이 이를 초안 생성기로 재사용합니다. 세 방법은 모두 상호작용형 챗봇의 작은 배치 크기에서 추론 지연 시간을 가장 크게 줄입니다. 각 순전파 중 GPU가 대부분 메모리 대기 상태에 있기 때문입니다.