Prompt Lookup Decoding — Speculative decoding không cần draft model

Bạn đã từng nghe về Speculative Decoding — dùng một model nhỏ để "đoán" nhanh các token rồi để model lớn kiểm chứng. Giải pháp đó tuy nhanh nhưng tốn gấp đôi VRAM và đau đầu về việc cân bằng distribution giữa draft và target. Prompt Lookup Decoding (PLD) giải quyết triệt để vấn đề này bằng một insight đầy chiến lược: thay vì thuê thêm "chuyên gia phỏng vấn", hãy để model "mở sách" tra cứu trực tiếp từ chính prompt đầu vào. Không cần training, không cần GPU thêm, chỉ cần một thuật toán tìm kiếm đơn giản.

Vấn đề

Autoregressive generation bị một bottleneck cứng: memory bandwidth. Khi sinh từng token một, GPU phải load toàn bộ weights (từ HBM lên SRAM) cho mỗi token — dù FLOPs rất ít, nhưng băng thông nhớ bị bão hòa khiến GPU ngồi không 90% thời gian.

Giải pháp truyền thống là Speculative Decoding: dùng một model nhỏ (draft) chạy nhanh sinh ra 3-5 token tạm, rồi model lớn (target) verify cả chuỗi trong một forward pass. Nếu đúng, tiết kiệm được nhiều bước; nếu sai, quay lại vị trí lệch. Nhưng cách này có hai vấn đề lớn:

1. Memory overhead: Phải giữ cả draft model (ví dụ 7B) và target model (70B) trong VRAM cùng lúc, cộng thêm hai KV cache riêng biệt.
2. Alignment risk: Draft model thường có distribution khác target, dẫn đến tỷ lệ accept thấp, hoặc cần finetune để "ăn khớp" — tốn thời gian và data.

Có cách nào đạt được tốc độ tương đương mà không trả giá bằng VRAM và complexity không?

Ý tưởng cốt lõi

PLD bỏ hoàn toàn khái niệm draft model. Thay vào đó, nó khai thác một đặc điểm của hàng loạt use case thực tế: RAG, summarization, code editing — những task mà output thường là sự trích xuất và lắp ráp từ input context, hơn là sáng tạo từ thin air.

Cơ chế cực kỳ đơn giản:

1. Lấy n-gram: Tại bước sinh token thứ t, lấy n token vừa sinh ra (thường n=3) làm "search key".
2. Tìm trong prompt: Dùng string matching (sliding window) để tìm key này trong toàn bộ prompt input.
3. Lấy continuation: Nếu tìm thấy, lấy k token tiếp theo ngay sau vị trí match trong prompt làm "candidate sequence" (thường k=10).
4. Verify: Đưa cả chuỗi candidate vào target model, chạy một forward pass để tính probability. Token nào khớp với distribution của target thì giữ lại, token nào lệch thì reject và resample từ điểm đó.

The "Open Book Exam" Analogy:

Hãy tưởng tượng bạn làm bài thi mở sách. Speculative decoding truyền thống giống như đưa cho bạn một "trợ lý nhanh nhưng hay đoán sai" viết trước đáp án, rồi bạn sửa lại. PLD thì đơn giản hơn: khi gặp câu hỏi, bạn lật sách ra tìm đoạn văn liên quan và chép y nguyên cụm từ tiếp theo. Nếu đoạn chép đó hợp lý (verify đúng), bạn tiết kiệm được thời gian suy nghĩ; nếu không, bạn tự viết lại từ đó.

Tại sao n-gram từ prompt lại là candidate tốt?

Trong các tác vụ như tóm tắt văn bản hay trả lời câu hỏi dựa trên tài liệu (RAG), LLM thực chất đang làm "copy-paste có chọn lọc" — nó trích xuất entity names, dates, code snippets từ context và lắp ráp lại. Do đó, chuỗi token tiếp theo thường đã xuất hiện nguyên văn trong prompt. PLD chỉ đơn thuần là nhận ra điều này một cách tường minh và tận dụng nó để parallelize việc sinh token.

Tại sao nó hoạt động

Memory Parallelism vs Sequential Choking

Bản chất tại sao PLD nhanh nằm ở việc biến đổi bài toán từ "sinh tuần tự" sang "verify song song". Khi target model verify một chuỗi dài, nó tính attention và logits cho tất cả vị trí trong một lần forward duy nhất — tận dụng được matrix multiplication parallelism của GPU. Điều này hiệu quả hơn nhiều so với việc chạy nhiều bước autoregressive riêng lẻ, nơi mỗi bước đều phải load weights từ HBM.

Lossless Guarantee

Một hiểu lầm phổ biến là "lookup từ prompt sẽ làm output cứng nhắc". Thực tế, PLD là lossless — nó bảo toàn distribution của target model. Tại sao? Vì nếu candidate token có probability bằng 0 (hoặc cực thấp) theo target model, nó sẽ bị reject ngay lập tức. Quá trình verification thực chất là rejection sampling: nó chấp nhận token chỉ khi target model cũng sẵn sàng sinh ra token đó. Nếu chuỗi candidate hoàn toàn sai, PLD fallback về autoregressive thuần túy, không mất gì về chất lượng.

Efficiency qua Match Rate

Hiệu quả của PLD phụ thuộc vào "acceptance rate" — tỷ lệ token được giữ lại sau verification. Trên các tác vụ như summarization (CNN/Dailymail) hoặc long-context QA, match rate thường đạt 60-70%, dẫn đến speedup 2-3x. Tuy nhiên, với creative writing từ zero (không có context để copy), match rate gần 0 và PLD không mang lại lợi ích — điều này hoàn toàn hợp lý vì đó không phải use case của nó.

Ý nghĩa thực tế

So sánh với các phương pháp khác

PLD chiến thắng ở mục "zero overhead" — bạn có thể bật/tắt nó trong vLLM hoặc HuggingFace Transformers mà không cần chuẩn bị model phụ hay sửa kiến trúc.

Benchmarks thực tế (Mistral-7B-Instruct trên A100):

• Summarization: 2.4x speedup
• Long-context QA: 2.4x speedup
• Multi-turn chat: Cải thiện nhẹ (phụ thuộc overlap)

Limitations

PLD không phải silver bullet. Nó thất bại thảm hại với:

• Creative writing từ zero: Không có prompt nào để lookup.
• Prompt ngắn: Không đủ n-gram để match.
• Output diverge sớm: Nếu generation đi sang hướng khác prompt ngay từ đầu, PLD fallback về autoregressive thường.

Nhưng trong hệ sinh thái RAG và agentic workflows — nơi context dài và output phụ thuộc nặng vào input — PLD là optimization "free lunch".

Đào sâu hơn

Paper gốc:

• "PLD+: Accelerating LLM inference by leveraging Language Model" (arXiv:2412.01447) — cải tiến cho phép multiple candidate sequences song song.
• Implementation gốc: apoorvumang/prompt-lookup-decoding

Cùng cụm (Inference Frontier):

• EAGLE — Draft head gắn vào internal layer, 3x speedup — So sánh với PLD khi bạn có thể train một draft head nhẹ thay vì lookup tĩnh.
• TurboQuant — Random rotation trước khi quantize, 3-bit KV cache không mất accuracy — Kết hợp PLD với KV cache nén để tối ưu cả compute và memory cho long-context.
• KV Cache Quantization — Nén cache INT4/INT3, mở khóa ultra-long context — PLD hiệu quả nhất khi prompt dài; nén KV cache giúp bạn chứa được context triệu token để lookup.

Đọc tiếp:

• Speculative Decoding — Model nhỏ draft, model lớn verify song song — Hiểu nền tảng speculative decoding để thấy tại sao PLD là bước tiến hóa logic.
• RAG — Tìm trước, generate sau — Use case điển hình cho PLD, nơi output thường là trích xuất từ context.
• KV Cache — Đừng tính lại cái đã tính — Cơ chế caching cơ bản mà PLD dựa vào để verify nhanh.