- Nhóm nghiên cứu Apple – MILA – Flatiron vừa công bố HedgeMamba: công thức 2 bước distill Transformer thành Mamba thuần (không hybrid).
- Chỉ với 10B token (~2,7% ngân sách huấn luyện teacher), mô hình học trò đạt perplexity 14.11 so với 13.86 của Pythia-1B — gần như giữ nguyên hiệu năng mà vẫn thừa hưởng O(1) memory của SSM.
TL;DR
Một nhóm nghiên cứu từ Apple, MILA và Flatiron Institute (Abhinav Moudgil, Ningyuan Huang, Eeshan Gunesh Dhekane, Pau Rodríguez, Aswathy Balagopalan, Luca Zappella, Federico Danieli) vừa nộp lên ICLR 2026 bài "Attention to Mamba: A Recipe for Cross-Architecture Distillation". Họ chứng minh có thể distill một Transformer pretrained thành Mamba thuần — không block Attention, không hybrid — và vẫn giữ gần như nguyên hiệu năng teacher, chỉ tốn ~2,7% ngân sách token pretraining.
Kết quả chính: HedgeMamba-1B đạt perplexity 14.11, rất gần với 13.86 của Pythia-1B teacher, chỉ dùng 10B token distillation so với 334B token pretraining gốc.
What's new
Trước đây, distill trực tiếp Transformer sang Mamba fail rất nặng (perplexity nhảy lên trên 100). Các công trình trước (Wang et al. 2024, Bick et al. 2024) buộc phải dùng kiến trúc hybrid — xen kẽ block Attention với block SSM — để cứu hiệu năng. Nhóm Apple/MILA/Flatiron lập luận rằng vấn đề không nằm ở kiến trúc, mà ở khởi tạo trọng số.
Công thức mới chia thành 2 giai đoạn:
- Stage 1 — Softmax → Linear Attention: thay block Attention của teacher bằng linear attention dạng Hedgehog (Zhang et al. 2024), dùng kernel trick để xấp xỉ softmax bằng tích vô hướng của các feature map học được. Đóng băng MLP, LayerNorm, embedding của teacher; chỉ học feature map qua cosine-matching loss với output từng layer.
- Stage 2 — Linear Attention → Mamba: dùng trọng số Linear Attention đã học làm khởi tạo cho mô hình Mamba đã điều chỉnh (gọi là HedgeMamba). Mở thêm các thành phần Mamba: SSM mixer (Λ, B, C), short 1D convolution, và gate branch với SiLU. Fine-tune toàn bộ bằng cross-entropy với ground truth.
Điểm khác biệt cốt lõi: đây là lần đầu tiên một SSM thuần (không Attention) giữ được hiệu năng teacher Transformer ở quy mô 1B, nhờ bắc cầu qua Linear Attention làm "ngôn ngữ trung gian".
Technical facts
Toàn bộ thử nghiệm chính dùng Pythia-1B làm teacher, OpenWebText làm dữ liệu distillation, tokenizer GPT-NeoX. Ngân sách 10B token (~1.1 epoch OpenWebText). Bảng kết quả downstream (zero-shot, lm-eval-harness):
| Model (1B) | PPL ↓ | ARC-C | ARC-E | PIQA | Lambada | BoolQ | HSwag | WinoG |
|---|---|---|---|---|---|---|---|---|
| Pythia (teacher) | 13.86 | 27.04 | 56.98 | 70.72 | 42.07 | 60.82 | 47.16 | 53.43 |
| Hedgehog baseline | 14.89 | 26.45 | 52.74 | 68.01 | 30.60 | 54.80 | 40.79 | 50.91 |
| HedgeMamba | 14.11 | 27.13 | 53.66 | 68.72 | 32.31 | 55.20 | 41.87 | 52.17 |
| Naive direct Mamba | >100 | — | — | — | — | — | — | — |
Một số con số đáng chú ý:
- Ngân sách token cực tiết kiệm: 10B token distillation ≈ 2,7% của 334B token pretraining Pythia gốc. Stage 1 dùng 1B token, batch 48, seq length 1024, 20K steps. Stage 2 dùng 9B token, 180K steps.
- Phân bổ token tối ưu giữa 2 stage là 10/90 (Stage 1 / Stage 2) cho PPL 14.11. Các lựa chọn khác: 100/0 = 25.71, 50/50 = 14.58, 0/100 = 17.08. Nghĩa là Stage 1 không cần nhiều token, nhưng không thể bỏ.
- Scaling theo ngân sách token: 1B → 16.56 PPL, 2B → 15.61, 3B → 15.15, 10B → 14.11 — giảm đều, chưa chạm đáy.
- Scaling theo kích thước: ở 160M teacher (PPL 39.38), HedgeMamba đạt 26.84 (thực tế vượt teacher do teacher undertrain). Ở 410M: 16.50 → 16.48. Ở 1B: 13.86 → 14.11.
- Ablation: thêm SSM hoặc Conv riêng lẻ không cải thiện PPL (vẫn 14.89). Nhưng thêm gate branch SiLU kéo PPL xuống 14.58 — đây là thành phần Mamba đóng góp lớn nhất.
- Hidden state dim 2048 (gấp 8 lần giới hạn 256 của kernel Mamba tối ưu) → phải serial hoá fused parallel_scan trong training, làm wall-clock bị "thổi phồng". Authors thừa nhận đây là rào cản kỹ thuật cần CUDA kernel mới để mở khoá tốc độ thực tế.
Comparison
So với các hướng đang song hành:
- MOHAWK / Llamba (Cartesia AI): 3-stage recipe, distill Llama-3.1-8B thành Llamba-8B thuần Mamba-2 với 12B token (<0,1% của 15T token teacher). Average downstream 68.8 vs 69.4 của Llama teacher. Llamba-3B thậm chí vượt Llama-3.2-3B (63.9 vs 61.9). Teacher hiện đại hơn (Llama-3), output cũng thực dụng hơn.
- Hybrid distillation (Wang et al. 2024, Bick et al. 2024): xen kẽ block SSM với block Attention. Đạt được hiệu năng tốt nhưng hy sinh tính "thuần recurrent" — vẫn còn chi phí quadratic ở các block Attention giữ lại.
- HedgeMamba: điểm độc đáo là Mamba thuần ở output, giá trị lý thuyết cao hơn — chứng minh Linear Attention đủ làm cầu nối nếu khởi tạo đúng cách.
Use cases
Ai hưởng lợi từ công thức này?
- Team có Transformer pretrained sẵn: không muốn huấn luyện SSM từ đầu (quá đắt), nhưng vẫn muốn ưu điểm O(1) memory per token của recurrent architecture. Distillation recipe cho phép "chuyển đổi" mô hình đã có với chi phí chỉ 2–3% ngân sách gốc.
- Inference long-context: Transformer scale quadratic theo sequence length khi generate autoregressive, trong khi Mamba giữ chi phí hằng số trên mỗi token mới. Với input dài (RAG, long chat, code hiểu cả repo), Mamba distilled có thể rẻ và nhanh hơn đáng kể.
- On-device / edge deployment: memory footprint hằng số hợp với phone, laptop, embedded. Cartesia đã ship Llamba MLX cho iPhone/MacBook; cùng hướng này.
- Nghiên cứu distillation: phát hiện 10/90 token split, tác dụng của gate branch, và các ablation scaling tokens/model size là input có giá trị cho các công thức tiếp theo.
Limitations & pricing
Không có gì hoàn hảo:
- Teacher cũ: chỉ thử trên Pythia (160M/410M/1B). Pythia không phải SoTA, nên trần của HedgeMamba bị giới hạn bởi trần của teacher. Chưa chứng minh được trên Llama-3, Qwen-2.5, Mistral.
- Lambada drop mạnh nhất: từ 42.07 → 32.31 (−9,8 điểm). Các task cần long-range reasoning vẫn bị tổn thương dù PPL chung phục hồi tốt.
- Wall-clock chưa chứng minh: hidden state 2048 > giới hạn 256 của kernel Mamba tối ưu → training phải serial hoá, inference cũng chưa có benchmark tốc độ đàng hoàng. Reviewer ICLR coi đây là điểm yếu đáng kể.
- Novelty bị chất vấn: meta-review ICLR 2026 nghiêng về reject vì "method dựa nhiều vào techniques đã có" (Hedgehog + Mamba). Đóng góp chủ yếu là công thức kết hợp + ablation, không phải primitive mới.
- Chưa public weights: paper-only, không có checkpoint, không có API. Giá bằng 0 nhưng cũng chưa dùng được ngay.
What's next
Roadmap ngầm từ rebuttal của tác giả:
- Viết CUDA kernel riêng cho hidden state 2048 để mở khoá tốc độ wall-clock thật sự — đây là blocker lớn nhất trước khi công thức có impact thực tế.
- Mở rộng lên teacher hiện đại hơn (Llama-3, Qwen-2.5). Nếu tỉ lệ 2,7% token giữ nguyên, distill một Llama-8B có thể chỉ tốn ~400B token — vẫn rẻ hơn pretraining.
- Xu hướng chung: post-training conversion đang thành một hướng đi nghiêm túc. Ngoài HedgeMamba, đã có MOHAWK (2024), Llamba (02/2025), Attention-Bridge (10/2025). Khả năng cao trong 2026 sẽ xuất hiện nhiều công thức tương tự với các kiến trúc recurrent khác (RWKV, RetNet, GLA).
Nguồn: arXiv 2604.14191, OpenReview ICLR 2026, Goomba Lab MOHAWK, Llamba (Cartesia AI).
