Medusa Heads

Các head Medusa là những head giải mã nhẹ được bổ sung trên lớp ẩn cuối cùng của model ngôn ngữ lớn, mỗi head được huấn luyện để dự đoán một token xa hơn so với lớp đầu ra của chính model. Trong một lượt forward, model tạo ra dự đoán token tiếp theo như thường lệ cùng các dự đoán cho vài vị trí kế tiếp; ở lượt tiếp theo, một bước attention dạng cây kiểm tra các dự đoán đó và giữ lại chuỗi dài nhất mà model đồng ý. Medusa là một dạng giải mã suy đoán không cần model nháp riêng, nên không phải huấn luyện, lưu trữ hay đồng bộ thêm mạng thứ hai. Tianle Cai, Tri Dao và các cộng sự giới thiệu phương pháp này trong Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads vào tháng 1 năm 2024.

Cách các head dự đoán trước#

Một bước tạo văn bản tiêu chuẩn đưa trạng thái ẩn cuối cùng tại vị trí t vào head của model ngôn ngữ, nơi tính điểm token ở t + 1. Head Medusa k đọc cùng trạng thái ẩn đó và tính điểm token ở t + k + 1. Mỗi head là một lớp truyền thẳng duy nhất với activation SiLU và kết nối dư, sau đó là phép chiếu lên từ vựng. Phép chiếu khởi tạo bằng một bản sao của head model ngôn ngữ, còn lớp truyền thẳng khởi tạo bằng 0, nên trước khi huấn luyện, mọi head đều tái tạo dự đoán token tiếp theo của model cơ sở. Bài báo cho thấy năm head là đủ. Ví dụ phác thảo bằng PyTorch này tạo ba head và so sánh đầu ra của chúng với head của model cơ sở trên một trạng thái ẩn:

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)  # lớp đầu ra của model đã đóng băng

# Mỗi head Medusa: một lớp dư SiLU (khởi tạo bằng 0) và một bản sao của head 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)

# Mọi head đều đọc cùng trạng thái ẩn cuối cùng; head k nhắm đến vị trí 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])

Đầu ra [True, True, True] cho thấy lớp dư khởi tạo bằng 0 giúp mỗi head mới bắt đầu giống hệt dự đoán của model cơ sở. Sau đó, quá trình huấn luyện điều chỉnh head k để hướng đến vị trí tương lai tương ứng, còn lm_head và backbone không thay đổi.

Attention dạng cây và cơ chế chấp nhận#

Các dự đoán đứng đầu của các head được kết hợp thành những chuỗi tiếp nối tiềm năng. Lấy hai token đứng đầu từ head 1 và ba token đứng đầu từ head 2, tích Descartes tạo ra sáu ứng viên có chung tiền tố, hình thành một cây:

Trong lượt forward tiếp theo, mọi token ứng viên được nối vào đầu vào cùng lúc; một mask attention cho phép mỗi token chỉ nhìn thấy các nút tổ tiên của nó trong cây, nhờ đó model chấm điểm song song mọi nhánh. Nhánh dài nhất được model chấp nhận sẽ được giữ lại, và trạng thái ẩn cuối cùng của nhánh đó được đưa vào các head cho vòng tiếp theo. Vì mọi head đều đọc cùng trạng thái ẩn, head 2 đề xuất cùng một token trên mọi nhánh của head 1. Các tác giả nhận thấy cây thưa 64 nút, được tạo từ những tổ hợp chính xác nhất, cho tốc độ tăng tốc cao hơn cây dày đặc 256 nút.

Khi giải mã tham lam, token ứng viên được chấp nhận nếu trùng với lựa chọn hàng đầu của model. Với lấy mẫu, Medusa cung cấp cơ chế chấp nhận điển hình, giữ lại mọi ứng viên có xác suất theo model cơ sở vượt ngưỡng gắn với entropy của phân phối; phương pháp này chấp nhận nhiều token hơn nhưng đánh đổi độ khớp chính xác với phân phối.

Huấn luyện Medusa-1 và Medusa-2#

Medusa-1 chỉ huấn luyện các head mới trên backbone đã đóng băng, nhờ đó chất lượng đầu ra của model được giữ nguyên; bài báo báo cáo tốc độ tạo sinh nhanh hơn hơn 2.2 lần. Việc huấn luyện các head cho Vicuna-7B mất khoảng năm giờ trên một GPU A100 duy nhất với 60,000 cuộc hội thoại ShareGPT. Đây cũng là cách chia lớp mà freeze=23 thiết lập khi fine-tuning Ultralytics YOLO26: các lớp 0-22 được giữ cố định, chỉ head phát hiện lớp 23 học từ dữ liệu mới.

Medusa-2 fine-tune đồng thời các head và backbone, giúp các head chính xác hơn và tăng tốc độ lên 2.3-3.6 lần nhờ công thức huấn luyện bảo vệ các khả năng ban đầu của backbone. Khi không có dữ liệu huấn luyện ban đầu, tự chưng cất huấn luyện các head trên văn bản do chính model tạo ra.

Medusa nằm trong số nhiều phương pháp tạo token nháp từ chính model đích. EAGLE tạo bản nháp từ đặc trưng nội bộ của model bằng một lớp tự hồi quy nhỏ; còn dự đoán nhiều token, được sử dụng trong DeepSeek-V3, huấn luyện thêm các module dự đoán trong quá trình tiền huấn luyện để các engine phục vụ sau này dùng làm bộ tạo bản nháp. Cả ba phương pháp giảm độ trễ suy luận nhiều nhất ở kích thước batch nhỏ của các chatbot tương tác, nơi mỗi lượt forward khiến GPU chủ yếu phải chờ bộ nhớ.