GrokAlign：Grokking的几何表征与加速

机器学习社区面临的一项关键挑战是理解和加速深度网络的训练动力学，这种动力学导致了延迟泛化和输入扰动的涌现鲁棒性，也称为grokking。先前的研究已将延迟泛化等现象与深度网络从线性到特征学习阶段的转变联系起来，并将涌现鲁棒性与网络功能几何的变化联系起来，特别是采用连续分段仿射非线性的深度网络中所谓的线性区域的排列。本文解释了grokking如何在深度网络的雅可比矩阵中实现，并证明将网络的雅可比矩阵与训练数据对齐（在余弦相似度意义上）可在低秩雅可比假设下确保grokking。我们的结果为雅可比正则化在优化深度网络中的使用提供了强有力的理论动机——我们将其作为一种方法引入，称为GrokAlign——我们通过实证证明，与权重衰减等更常规的正则化器相比，GrokAlign能更早地诱导grokking。此外，我们引入质心对齐作为雅可比对齐的一种可处理且可解释的简化方法，能够有效识别和追踪深度网络训练动力学的各个阶段。配套网页（https://thomaswalker1.github.io/blog/grokalign.html）和代码（https://github.com/ThomasWalker1/grokalign）。