Multi-token prediction · training heads và speculative decode

1Định nghĩa

Multi-Token Prediction (MTP) là thiết kế mở rộng mô hình Decoder-only: thay vì tại mỗi vị trí chỉ có một đầu ra duy nhất đoán token kế tiếp t+1, kiến trúc gắn thêm một hoặc vài nhánh phụ (MTP heads/modules) nối tiếp nhau để dự đoán luôn các token t+2, t+3 ngay từ trạng thái ẩn hiện tại.

Cơ chế này mang lại lợi ích kép ở cả hai pha: trong lúc huấn luyện, việc buộc mô hình dự đoán trước nhiều bước giúp các lớp ẩn học biểu diễn ngữ cảnh sâu hơn; khi đưa ra phục vụ suy luận, chính các nhánh MTP này có thể được bật lên làm bộ đề xuất token nháp (drafter) cho Speculative Decoding ngay trong cùng một mô hình mà không cần vận hành thêm một mô hình nháp riêng biệt.

Ví dụ đời thường: Multi-Token Prediction (MTP) giống như một kỳ thủ cờ vua hoặc người đánh máy lành nghề: thay vì chỉ nghĩ đúng 1 nước đi tiếp theo rồi dừng lại, ngay tại thời điểm hiện tại họ đã dự tính sẵn chuỗi 2–3 nước đi liên tiếp trong đầu — vừa giúp nước đi hiện tại sâu sắc hơn lúc tập luyện, vừa giúp đi nhanh liền một mạch nhiều nước khi đối thủ đi đúng như dự đoán.

2Luồng

flowchart TD
  step1["Huấn luyện nhiều offset"]
  step2["Đề xuất draft"]
  step3["Verify và cập nhật cache"]
  step1 --> step2
  step2 --> step3

Luồng mô tả cơ chế hoặc quy trình kiểm tra. Các ví dụ số là giả định giáo dục; demo Atlas không chạy kiến trúc production này.

3Giải pháp

Ghi số MTP module, shared weights, teacher-forcing dependencies và policy verification. Hai module không tự bằng tốc độ gấp ba.

Speculative greedy và speculative sampling có quy tắc khác nhau; muốn giữ distribution phải dùng acceptance/correction đúng cho sampling. Benchmark phải dùng cùng decoding settings.

Native MTP trong DeepSeek và Nemotron là lựa chọn theo checkpoint. Dedicated drafters trong Gemma là artifact khác; kiểm tra model card thay vì suy từ tên MTP chung.

4Profiling

• Acceptance & useful tokens

  Đếm draft proposed/accepted và output cuối, không dùng proposed tokens làm throughput.

• Latency và parity

  Đo TTFT, TPOT, verification overhead; so greedy output hoặc sampling distribution.

Mẹo đo lường: ghim số liệu trên workload thật, ghi cả p50/p95/p99 thay vì chỉ trung bình, và luôn ghi lại điều kiện chạy (batch, cache, đồng thời, phần cứng) trước và sau khi tối ưu. Số đo thiếu điều kiện chạy thì không so sánh được.

5Ví dụ

6Thuật ngữ

Auxiliary objective

Loss huấn luyện bổ sung, không đồng nghĩa token được chấp nhận khi serving.

Draft / verify

Đề xuất token rồi kiểm tra bằng target và quy tắc sampling phù hợp.

7Bài liên quan

8Nguồn