A New Understanding of Momentum: Approximating Gradient Descent at the Feature Level - Scientific Spaces
一个以动量为状态变量的优化器，基本形式如下： Mt= βMt−1+(1−β)Gt Wt= ϕ(Wt−1,Mt,Gt,t) 不同的优化器差别主要体现在更新函数ϕ上，如SGDM、SignSGD、Muon等，新的探索也基本围绕ϕ进行，因为第一个式子的动量太简单，简单到大家不认为它有什么修改空间。 而本文的主角正是动量。我们将给动量机制带来一个新的理解视角：动量不仅仅可以看作梯度的平均，它还可以被视为一个在线回归问题的解。由此出发，我们可以自然地引出最近的一些工作。 跟Muon一样，本文主要考虑线性层的矩阵参数：设有线性层Y=XW，其中X∈Rb×din是输入，W∈Rdin×dout是权重，Y∈Rb×dout是输出。损失函数记为L(Y)=L(XW)，那么 G= ∂L ∂W =X⊤ ∂L ∂Y 最简单的优化器是（参数层面的）梯度下降 W←W−η ∂L ∂W =W−ηX⊤ ∂L ∂Y 然而，根据《为什么我们偏爱各向同性？基于最速下降的理解》的思想，我们认为参数本质上是模型的副产品，模型特征层面的变化，才是跟模型效果最相关的。理想情况下，我们希望实现特征层面的梯度下降 Y←Y−η ∂L ∂Y 问题在于，Y不是可以随意修改的变量，我们能直接修改的只有W，所以只能设法通过修改W来间接实现这一效果。 怎么间接实现呢？设最终的更新规则为W←W−ηΦ，那么Y←Y−ηXΦ，我们希望它能尽可能接近式(4)的效果，即希望XΦ≈ ∂L ∂Y ，所以我们考虑最小化 min Φ 1 2 ‖XΦ− ∂L ∂Y ‖ 2 F + λ 2 ‖Φ‖ 2 F 其中λ>0是正则系数。事实上，这只不过是一个线性回归问题，可以直接解得 Φ∗=(X⊤X+λI)−1X⊤ ∂L ∂Y 留意到X⊤ ∂L ∂Y 正好是G= ∂L ∂W ，(X⊤X+λI)−1是基于输入数据的Preconditioner，为了综合不同batch的贡献，我们考虑对X⊤X和X⊤ ∂L ∂Y 都做EMA，这将得到一个SGDM变体： Mt= βMt−1+(1−β)Gt Zt= βZt−1+(1−β)(X ⊤ t Xt+λI) Wt= Wt−1−ηZ −1 t Mt 现在我们知道，通过校正的梯度(X⊤X+λI)−1G，可以实现特征层面的梯度下降，那么我们可以将它视为某种“更靠谱的梯度