Muon optimizer là gì? Vì sao Moonshot dùng để train Kimi K2 1T params

Cuối bài về GD, SGD, Adam mình có gài một câu: "optimizer chưa dừng ở Adam, có một thứ tên Muon đang tới." Bài này đáng lẽ phải viết ngay sau đó. Mình kéo dài tới giờ vì đào càng sâu càng thấy nó liên quan tới hai bài mình vừa viết là DeepSeek V4 và YOLO26.

V4 là một canh bạc kiến trúc. YOLO26 là một đợt dọn dẹp. Hai bài tưởng chừng không liên quan nhưng lại chung một quy luật: những thứ ta tưởng đã ổn định - kiến trúc transformer, phương pháp anchor-free detection - đang bị viết lại ngay dưới chân mình. Muon là quy luật đó áp vào không gian optimizer. Nó lấy Shampoo, lột bỏ những phần làm nó chậm, thêm momentum đúng chỗ, rồi kiểm chứng bằng 12 kỷ lục liên tiếp trên NanoGPT speedrun mà không một ai quay lại AdamW.

Paper "Muon is Scalable for LLM Training" của Moonshot AI tuyên bố chỉ cần 52% FLOPs để đạt cùng loss như AdamW.

Nói cho dễ hiểu: FLOPs là số phép toán mà GPU phải thực hiện, tỷ lệ thuận trực tiếp với tiền điện và thời gian thuê máy chủ. "Đạt cùng loss" nghĩa là model học xong ngang chất lượng - không thua kém, không bị cắt giảm. Vậy nên 52% có nghĩa là để có model ngang ngửa, bạn chỉ tốn chưa đến một nửa tiền và thời gian.

Mình hơi dè chừng với mấy con số kiểu này, vì paper về optimizer tuyên bố nhanh gấp đôi rồi biến mất là chuyện thường mấy năm gần đây. Nhưng cái khiến mình ngồi đọc tiếp là Kimi K2 - 1T tham số, 15.5T tokens - đã train xong bằng Muon mà không hề crash. DeepSeek V4 thì đổi luôn default sang Muon. Hai chỗ này không phải benchmark trong paper. Đây là production thật.

Trong khi đó ở Việt Nam, mình hỏi quanh vài bạn làm ML, không ai nghe tên. Mình viết bài này vì nghĩ điều đó không nên kéo dài.

AdamW đang bỏ quên điều gì

AdamW cập nhật từng phần tử của ma trận weight độc lập, như thể mỗi con số đứng riêng. Nhưng weights trong transformer không phải con số rời rạc, chúng là ma trận (operator). Khi update theo từng phần tử, bạn không tôn trọng cấu trúc SVD của operator, hệ quả là 90% ma trận trong model train bằng AdamW có SVD entropy thấp hơn Muon, update bị tụ vào vài hướng dominant.

AdamW cập nhật mỗi tham số theo công thức:

$$ w_{t+1} = w_t - \eta \cdot \frac{m_t}{\sqrt{v_t} + \epsilon} $$

Trong đó: - $m_t$ là hướng đi trung bình (trung bình động của gradient). - $v_t$ là độ lớn trung bình (trung bình động của bình phương gradient). - $\epsilon$ là một số rất nhỏ để tránh chia cho 0.

Nói nôm na: bạn đang leo núi trong sương mù. Gradient cho bạn biết hướng xuống dốc. AdamW không bước theo hướng đó ngay, mà tính xem hướng này có ổn định không ($m_t$), rồi chia cho độ dốc lớn hay nhỏ ($\sqrt{v_t}$). Nếu dốc hay thay đổi đột ngột, AdamW tự động rút ngắn bước chân. Nếu dốc êm đềm, bước chân dài hơn.

Điểm quan trọng: mỗi con số trong ma trận weight có cặp $m_t, v_t$ riêng, được cập nhật độc lập như một ông cụ tự lo việc của mình.

Vấn đề là weights trong transformer không phải con số đứng riêng. Chúng là ma trận. Một matrix shape $[d_\text{in}, d_\text{out}]$ là một operator: nhận vector vào, trả vector ra. Cấu trúc của operator nằm ở singular value decomposition - nếu $W = U \Sigma V^\top$, thì các giá trị suy biến trong $\Sigma$ cho bạn biết operator này kéo giãn không gian theo hướng nào, mạnh bao nhiêu.

Khi bạn cập nhật từng phần tử độc lập như AdamW, bạn không tôn trọng cấu trúc đó. Hệ quả thực tế, theo phân tích của paper Moonshot ở §3.4: train cùng một model với cùng một lượng compute, trên 90% các ma trận weight, AdamW cho ra SVD entropy thấp hơn Muon. Tức là update của AdamW bị tụ vào vài hướng dominant. Các hướng "hiếm" - giá trị suy biến nhỏ - bị nhấn chìm trong noise.

Keller Jordan kiểm tra điều này theo cách khác. Anh lấy momentum buffer của các params 2D trong transformer, tính SVD, và phát hiện chúng gần low-rank, condition number cao. Vài hướng có magnitude lớn nuốt hết phần còn lại. Những hướng có thể quan trọng cho việc học một concept mới nhưng đang còn yếu thì không bao giờ kịp lớn lên.

Đó là khoảng trống mà Muon nhắm vào.

Muon làm gì khác

Muon trực giao hóa toàn bộ ma trận momentum bằng phép lặp Newton-Schulz, đưa mọi singular value về 1, để mọi hướng update có magnitude bằng nhau. Hướng "hiếm" được kéo lên ngang hàng với hướng dominant. Embedding, lm_head, scalar params vẫn dùng AdamW, chỉ 2D weight matrices trong transformer dùng Muon.

Muon thay AdamW cho các 2D weight matrices trong transformer. Thay vì chia gradient theo từng phần tử, Muon trực giao hóa toàn bộ ma trận momentum bằng phép lặp Newton-Schulz (một chuỗi phép nhân ma trận lặp lại vài lần để xấp xỉ kết quả của SVD mà không cần làm SVD thật, chi tiết ở section dưới), trả lại $U V^\top$ từ SVD $U \Sigma V^\top$ - tức là đưa tất cả giá trị suy biến về 1. Embedding, lm_head, và scalar params vẫn dùng AdamW.

Ma trận trực giao là gì? Một ma trận vuông $Q$ được gọi là trực giao nếu $Q^\top Q = I$, tức chuyển vị của nó chính là nghịch đảo. Hình dung đơn giản: các cột của $Q$ là những vector đơn vị vuông góc với nhau từng đôi một - giống như trục $x, y, z$ trong không gian 3D, nhưng ở số chiều bất kỳ. Khi nhân một vector với ma trận trực giao, độ dài của vector không đổi, chỉ bị xoay hoặc phản xạ. Đó là lý do nó được gọi là "orthogonal" - giữ góc, giữ độ dài.

Trực giao hóa ma trận là gì? Trong ngữ cảnh Muon, đó là biến một ma trận bất kỳ $M$ thành một ma trận "gần trực giao" $UV^\top$ bằng cách bỏ đi ma trận đường chéo $\Sigma$ trong phân rã SVD $M = U\Sigma V^\top$. Kết quả $UV^\top$ là ma trận trực giao vì cả $U$ và $V$ đều là ma trận trực giao. Thao tác này đưa mọi singular value về 1, nghĩa là mọi hướng trong không gian vector đều được cập nhật với cùng một độ lớn - không hướng nào bị bỏ quên, không hướng nào bị nuốt chửng.

Nói nôm na: sau khi tính momentum buffer $M$, mình làm SVD $M = U \Sigma V^\top$, bỏ $\Sigma$ đi, thay bằng identity. Update cuối cùng là $U V^\top$. Mọi singular value bằng 1, mọi hướng update có magnitude như nhau, hướng hiếm được kéo lên ngang hàng với hướng dominant.

Còn toán học thì chỉ 3 dòng:

$$ \begin{aligned} M_t &= \beta \cdot M_{t-1} + \nabla L(W) \ O_t &= \mathrm{NewtonSchulz}(M_t) \quad \approx U V^\top \ W_t &= W_{t-1} - \eta \cdot O_t \end{aligned} $$

Một góc nhìn lý thuyết hơn (Bernstein và Newhouse, 2024): nếu coi optimizer là "steepest descent dưới một norm nào đó", thì AdamW = steepest descent dưới Max-of-Max norm (norm trên từng phần tử), còn Muon = steepest descent dưới spectral norm. Với weights là operators, spectral norm là norm tự nhiên hơn. Update có spectral norm bằng 1 thì step size dễ đoán, không bị nhạy cảm với cấu trúc ma trận.

Vì sao Newton-Schulz, không phải SVD

Newton-Schulz chỉ dùng matrix multiplication, chạy stable trên bfloat16 của GPU. SVD chính xác nhưng cần float32 và quá chậm. Phép lặp Newton liên hợp mà Shampoo dùng cũng cần float32, không thân thiện với GPU. Muon chọn Newton-Schulz vì matmul là thứ GPU làm nhanh nhất.

SVD bị loại đầu tiên vì quá chậm và đòi float32 với một loạt phép toán không phải matmul. Trên GPU thì đây là combo tệ nhất có thể.

Lựa chọn thứ hai là phép lặp Newton liên hợp, cái mà Shampoo dùng để tính inverse-fourth-root. Cũng float32, cũng không thân thiện với GPU.

Newton-Schulz là họ phép lặp thứ ba. Tên nghe lạ, nhưng ý tưởng đơn giản đến bất ngờ: thay vì tính SVD tường minh để biết $\Sigma$ rồi thay bằng identity, mình tìm một chuỗi phép biến đổi chỉ dùng matmul mà sau vài bước, kết quả tự động gần đúng với $UV^\top$, không cần đụng tới SVD lần nào.

Để hiểu vì sao điều đó khả thi, cần một quan sát nhỏ. Với ma trận $X = U \Sigma V^\top$, mình thử nhân $X$ với chuyển vị của chính nó rồi nhân lại với $X$:

$$X X^\top X = (U \Sigma V^\top)(V \Sigma U^\top)(U \Sigma V^\top) = U \Sigma^3 V^\top$$

Vì $U, V$ là ma trận trực giao nên $V^\top V = I$ và $U^\top U = I$. Cả hai biến mất ở giữa, $\Sigma$ "thoát ra" và bị nâng lên luỹ thừa 3, còn $U$ và $V$ không đổi. Tổng quát hơn, mọi đa thức bậc lẻ áp lên $X$ dưới dạng $aX + b(XX^\top)X + c(XX^\top)^2 X + \dots$ đều chỉ tác động lên $\Sigma$, để $U$ và $V$ yên.

Đó là chìa khoá. Nếu mình tìm được polynomial $\varphi(x) = ax + bx^3 + cx^5$ sao cho lặp đủ lần $\varphi^N(x) \to 1$ với mọi $x \in (0, 1]$, thì áp polynomial đó lên $X$ qua vài lần matmul, mình sẽ kéo mọi giá trị suy biến về 1, trong khi $U$ và $V$ được giữ nguyên. Output cuối cùng: $U \cdot I \cdot V^\top = UV^\top$. Đúng cái Muon cần, không cần SVD.

Quan trọng hơn nữa: cả phép lặp chỉ gồm matmul. Không có phép chia, không có căn bậc hai, không có decomposition. Matmul là thứ GPU làm nhanh nhất, chạy ổn định trên bfloat16, vừa vặn với phần cứng hiện đại.

Trick hay nhất của Keller

Coefficients mà người ta thường dùng là $(a, b, c) = (2, -1.5, 0.5)$. Polynomial này hội tụ đẹp về 1 nhưng rất chậm khi $x$ gần 0. Một singular value nhỏ phải lặp rất nhiều bước mới kéo lên được gần 1.

Keller Jordan làm một việc hơi điên: anh bỏ luôn ràng buộc phải hội tụ chính xác về 1. Thực nghiệm cho thấy nếu polynomial chỉ cần hội tụ về $1 \pm 0.3$ - tức giá trị suy biến cuối nằm trong khoảng $[0.7, 1.3]$ - loss không hề bị ảnh hưởng. Output không còn là $U V^\top$ nữa mà là $U S' V^\top$ với $S' \sim \mathrm{Uniform}(0.5, 1.5)$, nhưng model vẫn train tốt như nhau.

Bỏ ràng buộc đó cho phép anh maximize $a$ - slope tại 0. Slope tại 0 quyết định tốc độ kéo các giá trị suy biến nhỏ lên. Kết quả sau khi tinh chỉnh: $(a, b, c) = (3.4445, -4.7750, 2.0315)$. Chỉ cần 5 lần lặp là đủ.

Trái: 1 bước polynomial. Tuned (đỏ) dốc hơn nhiều ở gần 0. Phải: sau 5 bước. Baseline vẫn còn lẹt đẹt với giá trị suy biến nhỏ. Tuned đã đẩy gần như mọi $x \in (0, 1]$ vào vùng $[0.7, 1.3]$ chấp nhận được.

Đoạn code thực hiện đúng 3 dòng (trích muon.py của Keller, đầy đủ, không cắt):

A = X @ X.mT              # XX^T, để lát nữa nhân lại với X tạo bậc lẻ
B = b * A + c * A @ A     # bXX^T + c(XX^T)^2, phần "đa thức" của polynomial
X = a * X + B @ X         # aX + bXX^TX + c(XX^T)^2 X = aX + bX^3 + cX^5 (theo singular value)

Mỗi lần chạy 3 dòng này là một bước lặp. Lặp 5 lần là xong, $X$ đã gần $UV^\top$.

Mối liên hệ với Shampoo

Muon không phải optimizer mới. Muon là Shampoo bị lột bỏ phần tích luỹ preconditioner qua thời gian, đổi coupled-Newton (float32) thành Newton-Schulz (bf16, chỉ matmul), cộng thêm momentum trước khi trực giao hóa. Cùng một bản chất toán học, nhưng chạy được trên GPU hiện đại.

Trước khi vào liên hệ, nói nhanh Shampoo là gì cho ai chưa gặp. Shampoo (Gupta, Koren, Singer 2018) là một second-order optimizer kiểu Adagrad nhưng cho ma trận: thay vì giữ một preconditioner duy nhất, nó tích luỹ hai preconditioner $L_t = \sum_s G_s G_s^\top$ và $R_t = \sum_s G_s^\top G_s$ ở hai phía gradient, rồi update bằng $W_{t+1} = W_t - \eta L_t^{-1/4} G_t R_t^{-1/4}$. Nó mạnh hơn Adam vì khai thác cấu trúc 2 chiều của weight matrix, nhưng đắt: phải tính inverse-fourth-root của ma trận lớn ở float32, và phải tích luỹ $L, R$ qua thời gian.

Muon đi đường tắt. Nếu bạn lấy Shampoo nhưng bỏ phần tích luỹ đi, tức là chỉ dùng $L = G G^\top$ và $R = G^\top G$ ngay tại bước hiện tại, thì với $G = U S V^\top$, update rule rút gọn lại thành:

$$W_{t+1} = W_t - \eta (US^{-1/2}U^\top)(USV^\top)(VS^{-1/2}V^\top) = W_t - \eta \cdot UV^\top$$

Tức là Shampoo không có accumulation chính là gradient được trực giao hóa. Cộng thêm momentum trước khi trực giao hóa, đổi coupled-Newton thành Newton-Schulz, bạn có Muon.

Muon không phải optimizer mới phát minh ra. Muon là Shampoo bị lột bỏ những phần làm nó chậm và không thân thiện với GPU, rồi thêm momentum đúng chỗ. Trong một thế giới mà mọi paper optimizer đều cố gắng nghĩ ra thuật toán "mới", Muon đi ngược lại: nó thanh lọc một thuật toán cũ về dạng tối giản nhất chạy được trên GPU.

Phần Moonshot AI sửa để Muon scale được

Muon gốc của Keller bỏ weight decay. Ở quy mô speedrun NanoGPT thì không sao. Nhưng khi Moonshot scale lên 800M params trên 100B tokens, weights bắt đầu vượt range của bf16, output các layer phình lên theo. Training không crash nhưng cũng không healthy.

Họ sửa hai thứ trong paper:

Thêm weight decay, đơn giản như AdamW: $W_t = W_{t-1} - \eta (O_t + \lambda W_{t-1})$. Với $\lambda = 0.1$.

Match update RMS với AdamW. Đây là chỗ tinh tế hơn. Họ chứng minh một lemma: với matrix shape $[A, B]$ full-rank, theoretical Muon update có RMS bằng $\sqrt{1/\max(A, B)}$. Phụ thuộc shape.

Trong khi AdamW update RMS thường ổn định khoảng $0.2$ tới $0.4$. Nghĩa là nếu bạn dùng cùng learning rate cho cả hai, bạn đang cho mỗi optimizer một step size effectively khác nhau.

Họ scale lại: $\eta \cdot (0.2 \cdot O_t \cdot \sqrt{\max(A, B)} + \lambda \cdot W_{t-1})$. Constant $0.2$ và scale $\sqrt{\max(A, B)}$ đưa update RMS của Muon về cùng range với AdamW. Hệ quả thực tế: bạn dùng lại nguyên LR và WD đã tune cho AdamW, không cần grid search lại từ đầu. Đây là phần khiến Muon scale được trong production thực, không chỉ trên bài test speedrun.

Phản biện: chỗ Muon chưa đáng tin

Mình đã cố tìm chỗ Muon thua. Có một chỗ rõ ràng.

Paper Moonshot có Table 6 trong §3.5.1 Ablation Studies mà nếu bạn đang tính chuyển ngay sang Muon thì nên đọc kỹ. Họ test bốn tổ hợp pretrain-finetune (AdamW hay Muon ở mỗi giai đoạn), và kết quả là chỉ có cặp Muon-Muon mới thắng tất cả. Nếu pretrain bằng AdamW rồi SFT bằng Muon (hoặc ngược lại), Muon không còn ưu thế đáng kể.

Cụ thể hơn, khi SFT Qwen2.5-7B (vốn pretrain bằng AdamW) bằng Muon, kết quả thua SFT bằng Adam trên gần như mọi benchmark: MMLU 70.8 so với 71.4, GSM8K 85.8 so với 89.8 (Table 7 trong paper). Paper thừa nhận thẳng đây là open problem, không che giấu.

Keller cũng list 3 câu hỏi anh chưa trả lời được trong bài blog gốc:

1. Muon có giữ được ở 20B+ params, 1T+ tokens không? (Kimi K2 đã trả lời phần này: có, ở 1T params.)
2. Newton-Schulz có distribute được tử tế trên large clusters không?
3. Chỉ work cho pretrain, hay cũng work cho fine-tune và RL?

Câu thứ ba quan trọng nhất với người đang làm RLHF / RL post-training. Hiện tại chưa ai có câu trả lời chắc.

Còn một câu nữa từ Keller mà mình tôn trọng. Khi được hỏi tại sao orthogonalization work, anh thừa nhận một phần valid answer là "It just is OK" (quote Shazeer 2020). Tức là cộng đồng deep learning đang ship một thuật toán work mà chưa hiểu hết tại sao nó work. Sự thẳng thắn đó đáng tôn trọng, nhưng nó cũng là dấu hiệu đáng ngại nếu bạn cần explanation đầy đủ trước khi adopt.

Lập luận then chốt không nằm ở quy luật mở rộng quy mô

Paper Moonshot có quy luật mở rộng quy mô đẹp. Cụ thể:

• 5 dense Llama models từ 399M tới 1.5B, fit curve: $L_\text{Muon} = 2.506 \cdot C^{-0.052}$, $L_\text{AdamW} = 2.608 \cdot C^{-0.054}$. Tỉ số FLOPs để đạt cùng loss: ~52% Muon so với 100% AdamW.
• Moonlight (16B MoE / 5.7T tokens) đánh bại DeepSeek-V2-Lite, Qwen2.5-3B trên gần như mọi benchmark dù chỉ train 5.7T tokens (Qwen2.5-3B train 18T).
• Overhead của optimizer chỉ 1-3% tổng wallclock. NS chạy bf16, không ảnh hưởng numerical.

Nhưng lập luận then chốt thật sự, theo mình, không phải các con số trên. Nó nằm ở lập luận của Keller về tiêu chuẩn bằng chứng cho optimizer.

Anh chỉ ra một vấn đề kinh điển: hầu hết paper optimizer đều tuyên bố đánh bại AdamW, nhưng không ai thực sự dùng chúng. Lý do là benchmark trong paper thường có AdamW baseline tune kém, hoặc setup được chọn sao cho optimizer mới tỏa sáng. Bạn không có cách nào kiểm chứng từ bên ngoài.

Keller đề xuất một tiêu chuẩn khác: competitive speedrun. Muon lần đầu tiên xuất hiện ngày 15/10/2024, lập kỷ lục NanoGPT với cải thiện 35% step-time so với AdamW. Sau đó 12 kỷ lục liên tiếp được lập bởi 7 researchers khác nhau, không ai bỏ Muon quay lại AdamW.

Lập luận của anh: bạn không cần tin benchmark của ai cả. Bạn chỉ cần tin rằng tồn tại những researchers biết tune AdamW và muốn lập kỷ lục NanoGPT. Nếu Muon là giả, người tiếp theo lập kỷ lục sẽ quay lại AdamW. 12 lần liên tiếp không ai quay lại. Đó là chuỗi bằng chứng mà không paper nào replicate được vì nó đòi hỏi một cộng đồng thật, đua thật, không có động cơ để giả.

Đây là cách nghĩ khác hẳn paper-driven research. Optimizer được kiểm chứng không phải bằng bảng số mà bằng lựa chọn thật sự của những người thực sự muốn thắng.

286 dòng code, 1T params

Repo gốc KellerJordan/Muon có toàn bộ implementation trong 286 dòng Python. Core algorithm (Newton-Schulz + muon_update) chỉ 35 dòng. Để so sánh: PyTorch AdamW khoảng 400 dòng, FSDP vượt 10K dòng.

35 dòng đó đang chạy trên Kimi K2 (1T params, 15.5T tokens, không crash). 35 dòng đó đang khiến DeepSeek đổi default. 35 dòng đó đang là frontier.

Có một điều mình vẫn nghĩ tới sau khi đọc xong. Nếu một thứ đơn giản tới mức này, đã được kiểm chứng bằng 12 kỷ lục liên tiếp từ tháng 10/2024, mà tới giờ mình mới biết, thì còn bao nhiêu thứ tương tự đang chạy ngoài kia mà mình đang bỏ lỡ chỉ vì không đọc đúng nguồn. Câu trả lời có lẽ không vui.

Câu hỏi thường gặp về Muon optimizer

Muon optimizer là gì?

Muon là optimizer cho 2D weight matrices trong transformer, ra mắt tháng 10/2024 bởi Keller Jordan. Nó trực giao hóa momentum buffer bằng phép lặp Newton-Schulz (5 vòng matmul trên bf16), đưa mọi singular value của update về 1, để hướng "hiếm" được học ngang hàng với hướng dominant. Bản chất toán học là Shampoo bị lột bỏ accumulation, cộng momentum.

Muon nhanh hơn AdamW bao nhiêu?

Theo paper "Muon is Scalable for LLM Training" của Moonshot AI, Muon chỉ cần khoảng 52% FLOPs để đạt cùng loss như AdamW trên các Llama dense models từ 399M tới 1.5B params. Overhead của Newton-Schulz chỉ chiếm 1-3% tổng wallclock training. Tức cùng một model output, bạn tiết kiệm hơn nửa tiền điện và thời gian thuê GPU.

Newton-Schulz iteration là gì?

Newton-Schulz là chuỗi phép biến đổi chỉ dùng matrix multiplication, lặp 5 lần để xấp xỉ kết quả của trực giao hóa $UV^\top$ mà không cần làm SVD. Polynomial dạng $aX + b(XX^\top)X + c(XX^\top)^2 X$ với coefficients $(3.4445, -4.7750, 2.0315)$ do Keller Jordan tinh chỉnh. Vì chỉ gồm matmul, nó chạy stable trên bf16 của GPU, khác với SVD (cần float32, quá chậm).

Có nên dùng Muon để fine-tune model pretrained bằng AdamW không?

Hiện tại không nên. Paper Moonshot Table 6 và Table 7 cho thấy SFT Qwen2.5-7B (vốn pretrain bằng AdamW) bằng Muon thua AdamW trên gần như mọi benchmark: MMLU 70.8 vs 71.4, GSM8K 85.8 vs 89.8. Pretrain-SFT optimizer mismatch là open problem chưa ai giải. Chỉ Muon-Muon (cả pretrain và SFT đều Muon) mới thắng đều.

Muon có scale được lên model lớn không?

Có, đã được kiểm chứng tới 1T params. Kimi K2 của Moonshot AI (1T params, 15.5T tokens) train hoàn toàn bằng Muon mà không crash. DeepSeek V4 đổi default sang Muon. Moonshot đã bổ sung weight decay và RMS matching với AdamW để Muon scale được trong production thực, không chỉ trên speedrun NanoGPT.

Muon có dùng cho RLHF / RL post-training được không?

Chưa có data công khai. Đây là câu hỏi thứ ba trong 3 open questions mà chính Keller Jordan list ra trong bài blog gốc. Hiện chưa ai publish kết quả Muon trên RL post-training, nên nếu bạn đang làm RLHF, đây là vùng đất chưa khám phá.

Code Muon ở đâu?

Implementation gốc của Keller Jordan ở github.com/KellerJordan/Muon, toàn bộ 286 dòng Python, core algorithm chỉ 35 dòng. Class MuonWithAuxAdam đã handle sẵn việc split params (2D matrix dùng Muon, embedding/lm_head/scalar dùng AdamW). Setup khoảng 5 phút nếu bạn đã có training loop AdamW.

Tại sao Muon không dùng SVD trực tiếp?

SVD chính xác nhưng cần float32 và một loạt phép toán không phải matmul, combo tệ nhất trên GPU. Phép lặp Newton liên hợp mà Shampoo dùng cũng cần float32. Newton-Schulz chỉ dùng matrix multiplication, GPU làm cực nhanh, chạy stable trên bf16. Đây là lý do Muon scale được còn Shampoo thì không.

Bạn nên làm gì

Nếu bạn đang train model > 1B params từ đầu và optimizer là AdamW, một weekend thử Muon là đáng. Code có sẵn ở github.com/KellerJordan/Muon. Cách dùng đơn giản: 2D matrix params dùng Muon, embedding/lm_head/scalar dùng AdamW. Class MuonWithAuxAdam đã handle phần split sẵn.

Nếu bạn đang fine-tune từ checkpoint pretrained bằng AdamW (Llama, Qwen, Mistral...), đợi thêm. Pretrain-SFT mismatch là vấn đề thật, paper Moonshot thừa nhận, chưa ai có giải pháp.

Nếu bạn đang làm RLHF hay RL post-training, đây là vùng chưa có data công khai. Có thể bạn là người đầu tiên ghi lại trải nghiệm.

Mình thì đang đọc paper Essential AI ("Practical Efficiency of Muon for Pretraining", arXiv 2505.02222) để xem họ trả lời câu hỏi distributed Newton-Schulz thế nào. Có gì sẽ viết tiếp.

Nếu bạn đã thử Muon và thấy nó không work cho case của mình, mình muốn nghe. Đặc biệt nếu bạn nghĩ lập luận "speedrun là bằng chứng" của Keller có lỗ hổng, mình muốn nghe gấp đôi.

Hãy học toán để khi xạo chó có thể thở ra mấy từ để tăng trust

Bình.

Nguồn tham khảo:

• Liu et al. (Moonshot AI, 2025), "Muon is Scalable for LLM Training", arXiv:2502.16982
• Jordan et al. (2024), "Muon: An optimizer for hidden layers in neural networks", https://kellerjordan.github.io/posts/muon/
• Code: https://github.com/KellerJordan/Muon
• Bernstein & Newhouse (2024), spectral norm analysis of Shampoo
• Kimi K2 release: https://github.com/moonshotai/Kimi-K2
• Essential AI (2025), "Practical Efficiency of Muon for Pretraining", arXiv:2505.02222