OmniOPD: Distill LLM không cần logit của teacher

Nếu bạn đang theo dõi mảng LLM training, chắc bạn biết cái bài toán kinh điển: làm sao để train một model nhỏ hơn mà vẫn giữ được chất lượng của model lớn? Knowledge distillation là câu trả lời quen thuộc, nhưng cái khó ở đây là on-policy distillation tức là distill dựa trên chính những gì student model tự generate ra lại đang có một điểm nghẽn rất lớn mà ít người nói đến.

Meta Research vừa release paper OmniOPD và mình thấy cái approach này khá thú vị, đáng để anh em ngồi phân tích kỹ.

Vấn đề gốc rễ của On-Policy Distillation truyền thống

Thông thường, khi nói đến việc train student model, chúng ta có ba hướng chính:

• SFT (Supervised Fine-Tuning): Train trên static dataset. Đơn giản, nhưng bị off-policy distribution shift tức là student học trên data của người khác generate, không phải data của chính nó.
• RL (Reinforcement Learning): Student tự explore, nhưng reward signal rất sparse. Cực kỳ tốn compute và unstable.
• OPD (On-Policy Distillation): Kết hợp cả hai student generate ra trajectories của chính nó, teacher cho feedback ở token level. Nghe có vẻ ngon.

Nhưng OPD truyền thống có hai vấn đề cực kỳ đau đầu.

Thứ nhất, nó cần access trực tiếp vào token-level logits của teacher. Nghĩa là teacher phải là open-weight model. Bạn muốn dùng Claude hay Gemini làm teacher? Quên đi, API không trả về logits.

Thứ hai, ngay cả khi bạn có logits rồi, cái signal này vẫn rất brittle. Vấn đề nằm ở chỗ teacher và student phải có overlap đủ lớn ở phần "next token plausible" nếu hai model có distribution quá khác nhau, logit matching trở thành noise nhiều hơn là signal. Và hệ quả là model hay rơi vào repetition loops cái lỗi mà ai làm LLM inference cũng từng gặp.

OmniOPD giải quyết bằng cách nào?

Thay vì match logit ở token level, OmniOPD chuyển sang chunk-level semantic similarity. Nghe có vẻ đơn giản nhưng cái shift này thay đổi khá nhiều thứ.

Cụ thể, framework này dùng Monte Carlo rollouts để estimate teacher's local preference. Thay vì hỏi "teacher nghĩ token tiếp theo là gì?", OmniOPD hỏi "nếu generate thêm một đoạn (chunk) từ điểm này, output của teacher và student có semantically tương đồng không?"

Cái metric để so sánh là continuous semantic similarity trên multi-token chunks không phải binary match hay KL divergence giữa hai distribution logit.

Mình thấy cái này hay ở chỗ: nó decoupled hoàn toàn khỏi internals của teacher model. Teacher chỉ cần generate text đúng thứ mà mọi API đều support.

Peak-Entropy Scheduler Cái trick thực sự thông minh

Nhưng nếu chỉ dùng semantic similarity thôi thì compute sẽ rất tốn vì bạn phải audit mọi token. OmniOPD giải quyết bằng peak-entropy scheduler.

ý tưởng: chỉ audit student ở những điểm mà nó đang uncertain nhất tức là những reasoning forks có entropy cao. Những chỗ student confident thì bỏ qua, không cần teacher feedback.

Theo kinh nghiệm của mình khi làm với các language model, những chỗ model uncertain thường là những chỗ quan trọng nhất branch points trong reasoning chain, chỗ chọn approach, chỗ quyết định format. Audit đúng những điểm này thay vì audit toàn bộ sequence là một cách rất pragmatic để optimize compute.

Chống policy collapse với Bayesian prior

Một vấn đề khác của OPD là policy collapse model dần dần forget những gì nó biết ở các token không được audit. OmniOPD dùng hai cơ chế:

1. Dirichlet-Multinomial Bayesian prior: Bound variance của discrete sampling, tránh model bị kéo quá mạnh về một hướng.
2. Base-model KL anchor: Giữ cho student không drift quá xa so với base model ở những token không được audit.

Cái KL anchor này mình thấy khá giống với cách RLHF dùng KL penalty để prevent reward hacking về bản chất là cùng một vấn đề: làm sao optimize một objective mà không sacrifice những thứ khác.

Kết quả thực tế

So sánh nhanh giữa các approach:

Con số +28.64% trên math benchmark so với standard OPD là khá ấn tượng. Nhưng cái mình quan tâm hơn là việc dùng Claude hay Gemini làm teacher mà vẫn outperform open-weight teacher đây mới là cái unlock thực sự.

Tại sao chunk-level lại beat token-level?

Paper lý giải rất thuyết phục: token-level logit có information density cao nhưng noise cũng cao. Khi bạn so sánh logit distribution giữa hai model có architecture khác nhau, vocabulary khác nhau, training data khác nhau bạn đang so sánh những thứ không hoàn toàn comparable.

Chunk-level semantic similarity thì ngược lại: information density thấp hơn một chút, nhưng noise thấp hơn nhiều. Và trong training, signal-to-noise ratio mới là thứ quyết định.

Anh em lưu ý điểm này: đây không phải lần đầu tiên chúng ta thấy pattern "coarser granularity = better learning signal" trong ML. Nó xuất hiện ở nhiều chỗ từ reward shaping trong RL cho đến contrastive learning trong vision models.

Implications cho team làm LLM

Nếu bạn đang build pipeline để fine-tune hoặc distill LLM, OmniOPD mở ra một hướng rất thực tế:

Trước đây: Muốn on-policy distillation → phải dùng open-weight teacher (LLaMA, Qwen, Mistral...) → bị giới hạn bởi capability của open models.

Bây giờ: Có thể dùng Claude, Gemini, GPT-4 làm teacher thông qua API → student có thể học từ những model mạnh nhất hiện tại mà không cần white-box access.

Về mặt practical, cái này có nghĩa là chi phí để build một strong student model giảm đáng kể bạn không cần self-host teacher model nặng hàng trăm GB, chỉ cần API calls là đủ.

Tất nhiên vẫn còn câu hỏi về cost của Monte Carlo rollouts generate nhiều samples từ teacher API để estimate preference không phải free. Nhưng so với việc phải maintain infrastructure cho một 70B+ teacher model thì trade-off này hoàn toàn hợp lý với nhiều team.

Mình đang chờ xem code implementation của OmniOPD được release chưa để thử integrate vào pipeline. Nếu anh em đã đọc paper đầy đủ hoặc có thêm góc nhìn về approach này, drop comment xuống bên dưới mình muốn nghe thêm về experience thực tế của mọi người với on-policy training.