incorporate the extra scale on the muon update proposed by kimi team with conjunction with Jianlin Su

Author: lucidrains

README.md

@@ -99,3 +99,15 @@ for _ in range(100):
     url          = {https://kellerjordan.github.io/posts/muon/}
 }
 ```
+
+```bibtex
+@misc{kimiteam2025kimik2openagentic,
+    title   = {Kimi K2: Open Agentic Intelligence}, 
+    author  = {Kimi Team and Yifan Bai and Yiping Bao and Guanduo Chen and Jiahao Chen and Ningxin Chen and Ruijue Chen and Yanru Chen and Yuankun Chen and Yutian Chen and Zhuofu Chen and Jialei Cui and Hao Ding and Mengnan Dong and Angang Du and Chenzhuang Du and Dikang Du and Yulun Du and Yu Fan and Yichen Feng and Kelin Fu and Bofei Gao and Hongcheng Gao and Peizhong Gao and Tong Gao and Xinran Gu and Longyu Guan and Haiqing Guo and Jianhang Guo and Hao Hu and Xiaoru Hao and Tianhong He and Weiran He and Wenyang He and Chao Hong and Yangyang Hu and Zhenxing Hu and Weixiao Huang and Zhiqi Huang and Zihao Huang and Tao Jiang and Zhejun Jiang and Xinyi Jin and Yongsheng Kang and Guokun Lai and Cheng Li and Fang Li and Haoyang Li and Ming Li and Wentao Li and Yanhao Li and Yiwei Li and Zhaowei Li and Zheming Li and Hongzhan Lin and Xiaohan Lin and Zongyu Lin and Chengyin Liu and Chenyu Liu and Hongzhang Liu and Jingyuan Liu and Junqi Liu and Liang Liu and Shaowei Liu and T. Y. Liu and Tianwei Liu and Weizhou Liu and Yangyang Liu and Yibo Liu and Yiping Liu and Yue Liu and Zhengying Liu and Enzhe Lu and Lijun Lu and Shengling Ma and Xinyu Ma and Yingwei Ma and Shaoguang Mao and Jie Mei and Xin Men and Yibo Miao and Siyuan Pan and Yebo Peng and Ruoyu Qin and Bowen Qu and Zeyu Shang and Lidong Shi and Shengyuan Shi and Feifan Song and Jianlin Su and Zhengyuan Su and Xinjie Sun and Flood Sung and Heyi Tang and Jiawen Tao and Qifeng Teng and Chensi Wang and Dinglu Wang and Feng Wang and Haiming Wang and Jianzhou Wang and Jiaxing Wang and Jinhong Wang and Shengjie Wang and Shuyi Wang and Yao Wang and Yejie Wang and Yiqin Wang and Yuxin Wang and Yuzhi Wang and Zhaoji Wang and Zhengtao Wang and Zhexu Wang and Chu Wei and Qianqian Wei and Wenhao Wu and Xingzhe Wu and Yuxin Wu and Chenjun Xiao and Xiaotong Xie and Weimin Xiong and Boyu Xu and Jing Xu and Jinjing Xu and L. H. Xu and Lin Xu and Suting Xu and Weixin Xu and Xinran Xu and Yangchuan Xu and Ziyao Xu and Junjie Yan and Yuzi Yan and Xiaofei Yang and Ying Yang and Zhen Yang and Zhilin Yang and Zonghan Yang and Haotian Yao and Xingcheng Yao and Wenjie Ye and Zhuorui Ye and Bohong Yin and Longhui Yu and Enming Yuan and Hongbang Yuan and Mengjie Yuan and Haobing Zhan and Dehao Zhang and Hao Zhang and Wanlu Zhang and Xiaobin Zhang and Yangkun Zhang and Yizhi Zhang and Yongting Zhang and Yu Zhang and Yutao Zhang and Yutong Zhang and Zheng Zhang and Haotian Zhao and Yikai Zhao and Huabin Zheng and Shaojie Zheng and Jianren Zhou and Xinyu Zhou and Zaida Zhou and Zhen Zhu and Weiyu Zhuang and Xinxing Zu},
+    year    = {2025},
+    eprint  = {2507.20534},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2507.20534}, 
+}
+```

adam_atan2_pytorch/muon_adam_atan2.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+import math
 from typing import Callable
 
 import torch
@@ -66,6 +67,7 @@ def __init__(
         cautious_factor = 1., # set to 0. for zeroing out any updates not in same direction as gradient as in https://arxiv.org/abs/2411.16085
         a = 1.27,
         b = 1.,
+        muon_rms_factor = 0.2,
         muon_steps = 5,
         muon_beta1 = 0.95,
         muon_newton_schulz5_coefs = (3.4445, -4.7750, 2.0315),
@@ -101,14 +103,15 @@ def __init__(
             muon_steps = muon_steps,
             muon_newton_schulz5_coefs = muon_newton_schulz5_coefs,
             muon_eps = muon_eps,
+            muon_rms_factor = muon_rms_factor,
         )
 
         if remove_muon_params_from_params:
             params = list(set(params) - set(muon_params))
 
         param_groups = [
             dict(params = params, lr = lr),
-            dict(params = muon_params, lr = muon_lr, beta1 = muon_beta1, use_muon = True)
+            dict(params = muon_params, lr = muon_lr, beta1 = muon_beta1, rms_factor = muon_rms_factor, use_muon = True)
         ]
 
         super().__init__(param_groups, defaults)
@@ -134,6 +137,11 @@ def step(
 
                 param_init_lr = init_lr if not use_muon else init_muon_lr
 
+                # set lr scale to 1. if muon update
+
+                if use_muon:
+                    a = 1.
+
                 # maybe decoupled weight decay
 
                 if self.decoupled_wd:
@@ -193,7 +201,7 @@ def step(
                     # maybe cautious update - algorithm 2 in https://arxiv.org/abs/2411.16085
                 else:
 
-                    muon_steps, muon_coefs, muon_eps = group['muon_steps'], group['muon_newton_schulz5_coefs'], group['muon_eps']
+                    muon_steps, muon_coefs, muon_eps, muon_rms_factor = group['muon_steps'], group['muon_newton_schulz5_coefs'], group['muon_eps'], group['muon_rms_factor']
 
                     # Muon from Keller Jordan
                     # https://kellerjordan.github.io/posts/muon/
@@ -205,6 +213,13 @@ def step(
                         eps = muon_eps
                     )
 
+                    # incorporate the match adam RMS from Kimi team
+                    # https://kexue.fm/archives/11267
+
+                    muon_update_scale = math.sqrt(max(exp_avg.shape[-2:])) * muon_rms_factor
+
+                    update = update * muon_update_scale
+
                 if cautious_factor < 1.:
                     align_mask = (update * grad) > 0
                     scale = torch.where(align_mask, torch.ones_like(grad), cautious_factor)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "adam-atan2-pytorch"
-version = "0.2.2"
+version = "0.2.3"
 description = "Adam-atan2 for Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }