We develop a gradient flow on the space of probability measures defined on matrix-valued parameters induced by regularized Muon, an analytically smoothed version of the idealized Muon optimizer. The key observation is that the regularized orthogonalization map is the gradient of a smooth Fenchel-dual smoothing of the nuclear norm. This identifies the (regularized) Muon update as a mirror/prox step in the update variable, with momentum acting as the dual coordinate. We use this structure to lift Muon from a single matrix parameter to finite-particle probability objectives of the form $J(ρ)=R\left(\int F d ρ\right)$, a setting motivated by mean-field descriptions of neural-network training, and derive the inertial continuous-time limit. Using this structure, we derive the finite-particle continuous-time limit under the inertial scaling of step size and momentum, and then pass to a phase-space mean-field equation over probability laws on parameter-momentum pairs. The resulting flow can be shown to be a damped Hamiltonian probability dynamics whose kinetic energy is induced by the regularized Muon mirror potential. We prove an exact Hamiltonian dissipation identity, showing that the Hamiltonian energy decreases monotonically. While the target objective itself need not be monotone along the inertial Muon dynamics, under additional gradient-dominance, bounded-momentum, and curvature/alignment assumptions, we obtain continuous and discrete-time exponential convergence rates for the objective gap. We also study the well-posedness of the mean-field limit equation and establish propagation of chaos guarantees for the interacting particle system. Finally, we extend the formulation to Hilbert-valued feature maps on product matrix spaces, yielding a blockwise Muon probability flow applicable to smooth transformer mixture-of-experts models.


翻译:我们针对由正则化Muon(理想化Muon优化器的一种解析平滑版本)诱导的矩阵值参数空间上的概率测度,发展了一种梯度流。关键发现是:正则化正交映射是核范数光滑Fenchel对偶光滑化的梯度。这揭示了(正则化)Muon更新实质上是更新变量上的镜像/近端步,动量作为对偶坐标。利用该结构,我们将Muon从单矩阵参数提升至形式为$J(ρ)=R\left(\int F d ρ\right)$的有限粒子概率目标(该设定受神经网络训练的均场描述启发),并推导了惯性连续时间极限。基于该结构,我们推导了步长与动量惯性标度下的有限粒子连续时间极限,进而过渡至参数-动量对概率律上的相空间均场方程。所得流可视为阻尼哈密顿概率动力学,其动能由正则化Muon镜像势诱导。我们证明了精确的哈密顿耗散恒等式,表明哈密顿能量单调递减。尽管目标函数本身沿惯性Muon动力学未必单调,但在额外梯度优势、有界动量及曲率/对齐假设下,我们获得了目标间隙的连续与离散时间指数收敛率。我们还研究了均场极限方程的良好适定性,并建立了相互作用粒子系统的混沌传播保证。最终,我们将该框架推广至积矩阵空间上的Hilbert值特征映射,得到适用于平滑Transformer混合专家模型的块状Muon概率流。

0
下载
关闭预览

相关内容

【2023新书】光滑流形上的优化引论,368页pdf
专知会员服务
56+阅读 · 2023年8月7日
通过条件梯度进行结构化机器学习训练,50页ppt与视频
专知会员服务
13+阅读 · 2021年2月25日
【Nature论文】深度网络中的梯度下降复杂度控制
专知会员服务
41+阅读 · 2020年3月9日
【干货书】凸随机优化,320页pdf
专知
12+阅读 · 2022年9月16日
从泰勒展开来看梯度下降算法
深度学习每日摘要
13+阅读 · 2019年4月9日
博客 | 机器学习中的数学基础(凸优化)
AI研习社
14+阅读 · 2018年12月16日
2018年深度学习优化算法最新综述
计算机视觉战队
10+阅读 · 2018年12月11日
干货|掌握机器学习数学基础之优化[1](重点知识)
机器学习研究会
10+阅读 · 2017年11月19日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
6+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
VIP会员
最新内容
学习数据的几何:形状空间分析数学综述
专知会员服务
4+阅读 · 6月17日
定向能反无人机系统最新发展动态
专知会员服务
6+阅读 · 6月17日
从燃煤战舰到算法战争:水面指挥的永恒要求
专知会员服务
3+阅读 · 6月17日
相关VIP内容
【2023新书】光滑流形上的优化引论,368页pdf
专知会员服务
56+阅读 · 2023年8月7日
通过条件梯度进行结构化机器学习训练,50页ppt与视频
专知会员服务
13+阅读 · 2021年2月25日
【Nature论文】深度网络中的梯度下降复杂度控制
专知会员服务
41+阅读 · 2020年3月9日
相关基金
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
6+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
Top
微信扫码咨询专知VIP会员