提出低秩衰减机制,让模型在记忆后仍能持续简化结构并提前泛化。
Low-Rank Decay for Grokking in Scale-Invariant Transformers: A Spectral-Geometric View

- 用核范数类似正则化,保留权重空间的切向分量以持续优化
- 在模运算任务中使查询/键矩阵有效秩快速下降,提升泛化能力
- 适合研究模型泛化与权重动态的学者,尤其关注小样本学习
现代Transformer常使用归一化如RMSNorm和查询-键归一化,使模型部分呈现权重幅度的尺度不变性。此时标准的Frobenius范数权重衰减仅沿权重空间径向作用,无法直接简化归一化层的函数表示。本文通过小算法任务研究这一设定下的突现泛化(grokking),提出低秩衰减(Low-Rank Decay, LRD),一种类似核范数的谱正则化,其次梯度——极分解因子$UV^ op$——在尺度不变设置下仍保持切向分量。这带来明确的动力学差异:模型记忆训练集、任务梯度消失后,L2衰减无法再重塑权重谱,而LRD仍能以$\ ext{l}_1$式方式压缩奇异值。在模运算任务中,我们发现LRD促使查询/键矩阵的有效秩迅速坍缩,并将延迟泛化(grokking)发生的数据比例边界显著扩展。进一步通过‘针到扇’几何解释了核范数次微分在低秩流形附近的扩张行为。
原文摘要 · Abstract (English)
Modern Transformer architectures frequently employ normalization mechanisms such as RMSNorm and Query-Key Normalization, making parts of the model approximately scale-invariant with respect to weight magnitudes. In this regime, standard Frobenius-norm weight decay acts purely along the radial direction of the weight space and cannot directly simplify the function represented by the normalized layer. We study grokking in small algorithmic tasks through this lens and propose \emph{Low-Rank Decay} (LRD), a nuclear-norm-like spectral regularizer whose subgradient -- the polar factor $UV^\top$ -- retains a tangential component even in the scale-invariant setting. This distinction has a concrete dynamical consequence: after the model memorizes the training set and task gradients vanish, L2 decay can no longer reshape the weight spectrum, whereas LRD continues to compress singular values in an $\ell_1$-like fashion. On modular arithmetic tasks, we find that LRD induces rapid effective-rank collapse in Query/Key matrices and expands the data-fraction boundary at which delayed generalization (grokking) occurs. We further provide a spectral-geometric interpretation through the ``needle-to-fan'' expansion of the nuclear-norm subdifferential near low-rank strata.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。