Beam Search — Giữ N candidate thay vì greedy, tìm output tốt hơn

Khi ChatGPT tạo câu trả lời, nó không phải lúc nào cũng chọn từ "an toàn" nhất ở mỗi bước. Nếu làm vậy, nó sẽ rơi vào bẫy lặp lại vô nghĩa ("the the the..."). Beam Search là giải pháp: thay vì cam kết một đường đi duy nhất, nó giữ $N$ kịch bản "what-if" song song, như người chơi cờ tính toán nhiều nước đi phía trước, đánh đổi bộ nhớ lấy chất lượng toàn cục.

Vấn đề

Greedy decoding—luôn chọn token có xác suất cao nhất tại mỗi bước—là chiến lược "cận thị". Nó tạo ra hai vấn đề nghiêm trọng:

1. Tối ưu cục bộ (Local Optima) Một từ có vẻ tốt ngay bây giờ có thể ép model vào ngõ cụt ngữ pháp sau đó. Ví dụ: chọn "the" sau một noun phrase có vẻ an toàn, nhưng nó khóa bạn vào cấu trúc mạo từ xác định, có thể gây lủng củng khi câu tiếp diễn.

2. Bẫy lặp lại Các mô hình ngôn ngữ gán xác suất cao cho các mẫu generic, an toàn. Greedy decoding liên tục chọn các token này dẫn đến vòng lặp vô nghĩa ("the company said the company said..."). Đây là "lời nguyền xác suất cao"—những gì dễ xảy ra nhất thường là những gì nhàm chán nhất.

Ý tưởng cốt lõi

Thay vì theo đuổi một chuỗi duy nhất, Beam Search duy trì đồng thời $N$ chuỗi giả thuyết (beams). Tại mỗi bước sinh token:

Mở rộng (Expand) Mỗi beam hiện tại nhân ra $V$ khả năng tiếp theo (với $V$ là kích thước từ vựng), tạo ra $N \times V$ ứng viên.

Chấm điểm (Score) Tính tổng log-probability tích lũy cho toàn bộ chuỗi: $\sum_{t=1}^T \log P(y_t | y_{<t})$. Điều này cho phép so sánh công bằng các chuỗi có độ dài khác nhau.

Cắt tỉa (Prune) Chỉ giữ lại $N$ chuỗi có điểm cao nhất, loại bỏ phần còn lại. Các beam bị loại không còn tốn tài nguyên.

Lặp lại (Iterate) Tiếp tục cho đến khi các beam sinh ra token <end>, sau đó trả về chuỗi hoàn chỉnh có điểm cao nhất.

Đây là "aha moment": Bạn không nhìn trước tương lai (điều đó đòi hỏi tìm kiếm luỹ thừa), nhưng bạn cũng không cam kết sớm. Bằng cách giữ $N$ "file save" hoạt động song song, bạn để ngữ cảnh tương lai quyết định lựa chọn hiện tại nào thực sự tốt. Một từ trông tầm thường bây giờ có thể dẫn đến câu hay ba token sau; Beam Search giữ khả năng đó sống, trong khi greedy search giết chết nó ngay lập tức.

Phép ẩn dụ "Checkpoint Game": Tưởng tượng bạn chơi game platformer nơi một cú nhảy sai giết cả màn chơi. Greedy decoding là chơi không có checkpoint—nếu bạn chọn nhảy "tối ưu" cục bộ dẫn đến ngõ cụt, bạn phải chơi lại từ đầu. Beam Search là giữ $N$ bản save. Tại mỗi ngã rẽ, bạn thử cả $N$ hướng, chỉ giữ lại $N$ bản save hứa hẹn nhất. Bạn chưa nhìn thấy tương lai xa, nhưng bạn tránh được cạm bẫy "tốt ở đây, thảm họa ở đó".

Tác động đến bộ nhớ: Không giống greedy decoding chỉ cần một KV cache, Beam Search yêu cầu $N$ bộ cache riêng biệt. Với $N=5$ và độ dài chuỗi $L$, bạn đang lưu trữ gấp 5 lần tensor Key-Value. Đây là lý do beam width thường được giữ nhỏ (3–10) và tại sao nó kết hợp tự nhiên với PagedAttention để quản lý vụ nổ bộ nhớ này hiệu quả.

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

Mathematically, Beam Search tối đa hóa xác suất đồng thời của chuỗi: $P(y_1, ..., y_T) = \prod_{t=1}^T P(y_t | y_{<t})$. Bằng cách làm việc trong không gian log, tích thành tổng: $\sum \log P(y_t | ...)$, cho phép so sánh công bằng giữa các chuỗi dài ngắn khác nhau (thường kèm chuẩn hóa độ dài để tránh thiên vị câu ngắn).

Điểm then chốt là verification: Beam Search đảm bảo tìm được chuỗi tối ưu toàn cục (trong giới hạn beam width), không bị mắc kẹt trong cực trị cục bộ như greedy. Tuy nhiên, nó vẫn là xấp xỉ—tìm kiếm toàn diện qua không gian $V^L$ là bất khả thi về mặt tính toán.

Chi phí thực tế:

• Thời gian: $O(N \times V \times L)$ so với $O(V \times L)$ của greedy. Tuy nhiên, nghiên cứu từ Freitag & Al-Onaizan (2017) cho thấy flexible beam search (điều chỉnh độ rộng beam động) có thể đạt tăng tốc 43% so với beam cố định mà không mất chất lượng.
• Bộ nhớ: $O(N \times L)$ cho KV cache, tuyến tính với beam width. Điều này tạo áp lực lên VRAM, đặc biệt khi kết hợp với batch lớn trong Continuous Batching.

Ý nghĩa thực tế

Khi nào dùng Beam Search:

• Dịch máy: Google Translate từng dùng $k=5$. Các tác vụ đòi hỏi cấu trúc ngữ pháp chặt chẽ (dịch thuật, tổng hợp SQL) được hưởng lợi từ việc lập kế hoạch toàn cục.
• Sinh code: Khi cú pháp yêu cầu dấu ngoặc cân bằng hoặc indent đúng, Beam Search tránh được lỗi cú pháp "cận thị".
• Tóm tắt: Các mô hình T5/BART thế hệ đầu dùng beam search để đầu ra tập trung hơn.

Trade-offs quan trọng:

Lời nguyền Beam Search (Beam Search Curse): Nghịch lý là beam width quá lớn (ví dụ 50) thường cho văn bản tệ hơn cho tác vụ mở. Vì mô hình gán xác suất cao cho các mẫu generic ("the", "and"), beam lớn sẽ tối ưu hóa điên cuồng cho các chuỗi xác suất cao nhưng nhàm chán ("the company said the company said..."). Đây là lý do tại sao các tác vụ sáng tạo sử dụng temperature sampling thay thế—Beam Search sinh ra "an toàn" nhưng tẻ nhạt.

Bias độ dài: Beam Search có xu hướng thiên về câu ngắn hơn nếu không áp dụng chuẩn hóa độ dài (chia điểm cho $L^\alpha$).

Đào sâu hơn

• Paper gốc: Beam Search Strategies for Neural Machine Translation (Freitag & Al-Onaizan, 2017) — chiến lược beam width linh hoạt tăng tốc decoder 43% không mất chất lượng.
• Paper về giới hạn: A Call for Clarity in Beam Search (2022) — phơi bày sự khác biệt triển khai gây ra khoảng cách hiệu năng giữa các framework.

Cùng cụm

• KV Cache — Cơ chế lưu trữ bắt buộc phải nhân $N$ lần khi dùng beam search
• PagedAttention — Quản lý bộ nhớ cho nhiều beam hiệu quả qua cơ chế virtual memory
• Continuous Batching — Kết hợp beam search với batching động để tối ưu throughput
• Speculative Decoding — Kỹ thuật tăng tốc inference thay thế bằng cách dùng model nhỏ draft thay vì tìm kiếm beam

Đọc tiếp

• Flash Attention — Tối ưu hóa phép tính attention nền tảng cho cả beam search và greedy
• Quantization — Giảm áp lực bộ nhớ từ việc lưu $N$ bộ KV cache
• Inference Frontier (Level 2) — Các kỹ thuật suy luận tiên tiến như test-time compute vượt xa beam search truyền thống