用损失梯度替代优化器更新做SVD,能更准确揭示模型参数中特征形成的集中区域。
Gradient-Direction Sensitivity Reveals Linear-Centroid Coupling Hidden by Optimizer Trajectories

- 改用损失梯度的滚动SVD,而非AdamW更新,提升诊断灵敏度1-2个数量级。
- 在多任务模型中,梯度法恢复出20-45倍的特征耦合强度,而更新法仅≤1倍。
- 约束注意力更新至秩3子空间可使模型快速理解(grokking)加速约2.3倍。
我们发现,将AdamW更新的滚动SVD替换为损失梯度的滚动SVD,可使诊断结果提升1-2个数量级。对损失梯度进行SVD而非AdamW更新,使四种单任务模运算中SED方向与线性中心假设(LCH)特征间的微扰耦合度从平均 $ \bar{R}_k \approx 3 $--$9\times$ 提升至 $100$--$330\times$,消除了原测量中的操作依赖性。在共享编码器的多任务Transformer上,基于更新的SED测得 $ \bar{R}_k \leq 1 $,显示诊断失效;而基于每任务梯度的梯度法恢复出所有四类操作下 $ \bar{R}_k = 20 $--$45\times$。跨任务梯度聚合是主要障碍;对每任务梯度进行SVD可解决此问题。因果干预表明,将注意力更新限制在任意秩-3子空间(无论是否由SED导出)均能使grokking速度加快约2.3倍,而移除该秩-3成分在正确梯度投影方法下影响可忽略。因此,SED-LCH耦合是参数空间中特征形成集中区域的强诊断指标,但非唯一因果路径:在当前超参数下,自然全秩AdamW注意力更新高度秩冗余。
原文摘要 · Abstract (English)
We show that replacing the rolling SVD of AdamW updates with a rolling SVD of loss gradients changes the diagnostic by 1-2 orders of magnitude. Performing SVD on the loss gradient instead of the AdamW update increases the measured perturbative coupling between SED directions and Linear Centroid Hypothesis (LCH) features from $ \bar{R}_k \approx 3 $--$9\times$ to $100$--$330\times$ across four single-task modular arithmetic operations, eliminating the apparent operation dependence in the original measurement. On a multitask transformer with a shared encoder, update-based SED gives $ \bar{R}_k \leq 1 $ -- an apparent failure of the diagnostic -- while per-operation gradient-based SED recovers $ \bar{R}_k = 20 $--$45\times$ across all four operations. Gradient aggregation across competing tasks is the main obstruction; performing SVD on per-task gradients resolves it. A causal intervention shows that constraining attention updates to any rank-3 subspace (whether SED-derived or random) accelerates grokking by approximately $2.3\times$ across random seeds and operations, while removing the rank-3 component has negligible effect under proper gradient-projection methodology. The SED-LCH coupling is therefore a strong diagnostic of where feature formation concentrates in parameter space, but it is not a unique causal pathway: the natural full-rank AdamW attention update is highly rank-redundant under our hyperparameters.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。