联合仿射谱整形：超越仅权重Muon的权重与偏置更新耦合

# 联合仿射谱整形：超越仅权重Muon的权重与偏置更新耦合
来源：https://arxiv.org/html/2608.02991
\[1\]Honghai Liu

\[1\]State Key Laboratory of Robotics and Systems, Harbin Institute of Technology, Shenzhen, Shenzhen, 518055, China

###### 摘要

矩阵谱优化器重塑权重更新谱，但通常将向量值偏置委托给独立的优化器。我们研究这种分离是否中性（即无影响）。我们将每个仿射层表述为一个联合动量矩阵 \(A=[M_W,\alpha m_b]\) 并对整个矩阵应用截断正则化逆谱映射，同时生成权重和物理偏置更新。在从头训练的IMDb上对四层BERT-mini进行的严格五种子消融比较了精确SVD Muon、仅权重逆整形、仿射探针逆整形以及所提出的联合正则化逆（JRI）。仅权重逆整形使基于验证损失选择的测试准确率从84.903±0.242%提高到85.562±0.308%，并将选定的测试损失从0.3479降低到0.3345。允许偏置影响联合SVD但保留独立的Adam偏置更新，并不比仅权重逆整形有改进。联合使用变换后的偏置将选定测试准确率提高到85.738±0.180%，并将测试损失降低到0.3291，所有五个种子相对探针基线均有改进。在峰值性能窗口内，JRI保持了合格的权重更新范数，同时将偏置更新范数从0.02095降低到0.00301，将边界函数占比从86.58%降低到78.97%，并将权重诱导的边界运动与显式偏置之间的余弦从+0.030变为−0.137。独立的22种子复现得到85.743±0.203%的选定测试准确率。这些结果将联合仿射谱分配识别为对仅权重谱优化的一种小而一致的扩展。

###### 关键词：

矩阵优化、谱整形、Muon、仿射层、权重-偏置耦合、信息分配

## 1 引言

矩阵感知优化器重新激发了对神经网络更新几何结构的兴趣。Muon将动量矩阵 \(M=U\Sigma V^\top\) 替换为其极因子 \(UV^\top\)，在应用更新前将非零奇异谱展平。其近期的大规模实现表明，当与合适的更新缩放和权重衰减结合时，这种矩阵处理对语言模型训练是实用的[1 (https://arxiv.org/html/2608.02991#bib.bib1)]。后续工作对该操作进行了推广。DynMuon研究 \(U\Sigma^p V^\top\) 并在训练过程中调度 \(p\) [2 (https://arxiv.org/html/2608.02991#bib.bib2)]；Freon涵盖Schatten范数和准范数区域，包括倒置谱[3 (https://arxiv.org/html/2608.02991#bib.bib3)]；Muonp保留原始谱的分数幂[4 (https://arxiv.org/html/2608.02991#bib.bib4)]。相关方法抑制主导谱尖峰或裁剪谱范数，而不是完全展平谱[5 (https://arxiv.org/html/2608.02991#bib.bib5),6 (https://arxiv.org/html/2608.02991#bib.bib6)]。总之，这些研究表明奇异谱是一个分配接口：优化器可以将更新预算集中在主导模式上，在模式之间展平，或将其重新导向较弱的模式。

然而，大多数谱优化器仅将此接口应用于矩阵值权重。偏置向量、归一化参数、嵌入和其他张量通常分配给Adam或AdamW[7 (https://arxiv.org/html/2608.02991#bib.bib7)]。这种实现约定很方便，但对于仿射映射而言并非功能中性的。

对于均值为 \(\mu\) 的输入，一阶函数更新可写为

\[
\Delta z=\Delta W(x-\mu)+\left(\Delta W\mu+\Delta b\right). \quad (2)
\]
第一项随中心化样本变化，并改变表示的形状。第二项在所有样本间共享，并移动仿射边界。因此，即使显式偏置被单独优化，权重更新也包含类似偏置的分量 \(\Delta W\mu\)。早期工作测量了Transformer中的这种隐藏边界运动，并表明权重诱导的边界位移可以支配显式偏置通路[8 (https://arxiv.org/html/2608.02991#bib.bib8)]。更一般地，优化的信息分配视角将参数路径和坐标视为有限更新预算的竞争性目的地[9 (https://arxiv.org/html/2608.02991#bib.bib9)]。

公式(2 (https://arxiv.org/html/2608.02991#S1.E2))提出了一个具体问题：*如果权重和偏置共同决定一个仿射函数，那么它们是否也应共享谱更新预算？* 简单地将偏置与权重矩阵拼接本身并非新事。齐次坐标提供了经典表示 \([W,b][x^\top,1]^\top\)，而逐层曲率方法如K-FAC通过增广激活联合预处理权重和偏置[10 (https://arxiv.org/html/2608.02991#bib.bib10)]。我们的问题更窄。我们问的是动量对 \((M_W,m_b)\) 是否应在奇异值映射下被视为单一对象，以及变换后的偏置列是否应用作实际的物理偏置更新。

我们通过受控的四路消融来研究这个问题。第一种方法对权重应用精确SVD Muon，对偏置应用Adam。第二种方法在保留Adam偏置的同时，用正则化逆谱替换平坦的Muon谱。第三种方法允许缩放后的偏置动量参与联合SVD，但丢弃变换后的偏置列，仍使用Adam偏置；我们称此为*仿射探针*条件。第四种方法对两个分量应用相同的联合谱映射，并使用生成的列更新两个参数。我们称此方法为*联合正则化逆*（JRI）。

实证结果分离了三种效应。首先，在所研究的BERT-mini/IMDb设置中，逆谱分配显著优于精确SVD Muon，即使控制了权重更新范数和SVD实现也是如此。其次，允许偏置揭示联合谱而不使用变换后的偏置不会产生可靠的收益。第三，完整的联合路径在验证选择的准确率、测试损失和校准方面提供了较小但一致的改进。功能诊断表明，额外的改进与更小且轻微补偿性的显式偏置更新相关，而非与更大的偏置通路相关。

我们的贡献在范围上有意限定：

1. 我们定义了一个联合仿射谱对象 \([M_W,\alpha m_b]\)，以及一个直接输出 \(\Delta W\) 和 \(\Delta b\) 的谱更新。
2. 我们提供了一个严格的消融，区分仅权重逆整形、偏置作为谱探针、以及偏置作为联合变换的物理更新。
3. 我们将观察到的性能差异与形状–边界分配、更新范数和模块级偏置行为联系起来，同时明确区分了先前关于逆谱工作的贡献。

我们不声称逆谱普遍最优，不声称拼接权重和偏置是史无前例的，也不声称JRI已可用于大规模部署。本工作是对矩阵谱优化中一个缺失接口的机制研究。

## 2 相关工作

### 2.1 矩阵更新的谱整形

Muon将 \(M=U\Sigma V^\top\) 映射到 \(UV^\top\)，从而将所有非零奇异值替换为公共值。实践工作表明，当仔细选择逐层更新尺度时，这种更新可以扩展到大型语言模型[1 (https://arxiv.org/html/2608.02991#bib.bib1)]。最近的分析表明，平坦谱只是一个更大家族中的一个点。DynMuon使用 \(U\Sigma^p V^\top\)，并主张训练早期使用正幂、后期使用轻微负幂[2 (https://arxiv.org/html/2608.02991#bib.bib2)]。Freon导出了一个基于Schatten的家族，在标准梯度更新和Muon之间插值，并扩展到大型奇异值被相对抑制的准范数区域[3 (https://arxiv.org/html/2608.02991#bib.bib3)]。Muonp专注于保留原始谱部分信息的分数正幂[4 (https://arxiv.org/html/2608.02991#bib.bib4)]。

其他方法更直接地针对谱各向异性。Spectra将梯度谱建模为低秩尖峰加长尾，并抑制尖峰子空间而不放大尾部[5 (https://arxiv.org/html/2608.02991#bib.bib5)]。SPECTRA引入谱前和谱后裁剪，以限制稀疏梯度尖峰和优化器更新的谱范数[6 (https://arxiv.org/html/2608.02991#bib.bib6)]。这些方法与当前实现密切相关，因为这里使用的截断逆映射在大多数弱模式上饱和。因此，JRI的行为更像主导模式抑制加平坦尾部，而非无约束的矩阵伪逆。

因此，本文的新颖性不在于使用逆奇异值映射。相关区别在于映射所应用的对象，以及变换后的偏置列是否被保留为更新。

### 2.2 逐层矩阵预条件

矩阵和张量预条件器先于Muon存在。K-FAC通过Kronecker积近似逐层Fisher块，并通过增广常数坐标联合处理仿射权重和偏置[10 (https://arxiv.org/html/2608.02991#bib.bib10)]。Shampoo沿着张量维度维护矩阵统计量以获得非对角预条件[11 (https://arxiv.org/html/2608.02991#bib.bib11)]。这些方法将层视为结构化优化单元，而非不相关的标量坐标。

JRI不是曲率近似。它不估计Fisher矩阵、Hessian或Kronecker因子。相反，它重塑当前动量矩阵的奇异值。其与K-FAC的关系是概念性的：两者都拒绝在层级别完全分离处理权重和偏置。它们的数学对象和预期几何不同。

### 2.3 权重-偏置功能耦合与信息分配

对于非中心化输入，权重和偏置在功能上并非独立，因为 \(\Delta W\mu\) 是样本独立的。早期工作将此术语形式化为隐藏边界运动，并构造了显式的形状–边界正交更新[8 (https://arxiv.org/html/2608.02991#bib.bib8)]。该方法在函数空间中分解梯度，并补偿边界位移。JRI则保持参数化不变，通过共享谱变换耦合两条路径。

更广泛的信息分配视角通过训练期间错误信号写入的位置来表征优化：跨参数路径、坐标、层和样本区域[9 (https://arxiv.org/html/2608.02991#bib.bib9)]。从这个角度看，奇异值映射控制跨矩阵模式的分配，而联合仿射构造控制权重诱导路径与显式边界路径之间的分配。这里使用这种解释作为说明框架；实证主张仍针对所测量的具体设置。

## 3 方法

### 3.1 动量与仿射谱对象

考虑一个合格的仿射层，其权重 \(W\in\mathbb{R}^{m\times n}\)，偏置 \(b\in\mathbb{R}^m\)。设 \(G_W^{(t)}\) 和 \(g_b^{(t)}\) 分别为第 \(t\) 步的梯度。我们维护动量状态

\[
\begin{aligned}
V_W^{(t)} &= \beta V_W^{(t-1)} + (1-\beta) G_W^{(t)}, \quad (3) \\
v_b^{(t)} &= \beta v_b^{(t-1)} + (1-\beta) g_b^{(t)}, \quad (4)
\end{aligned}
\]
其中 \(\beta=0.95\)。进入谱变换的方向使用Nesterov式插值

\[
\begin{aligned}
M_W^{(t)} &= (1-\beta) G_W^{(t)} + \beta V_W^{(t)}, \quad (5) \\
m_b^{(t)} &= (1-\beta) g_b^{(t)} + \beta v_b^{(t)}. \quad (6)
\end{aligned}
\]
JRI形成增广矩阵

\[
A^{(t)} = \left[ M_W^{(t)},\; \alpha m_b^{(t)} \right] \in \mathbb{R}^{m\times(n+1)}, \quad (7)
\]
其中

\[
\alpha = \begin{cases}
16 & \text{若 } n=256, \\
32 & \text{若 } n=1024.
\end{cases} \quad (8)
\]
系数 \(\