发现嵌入层是模型延迟泛化的关键,通过优化学习率可加速训练。
Mechanistic Insights into Grokking from the Embedding Layer
- 嵌入层更新慢导致稀有词停滞,引发延迟泛化
- 嵌入与权重的双线性耦合引入鞍点,加剧初始化敏感性
- 自适应学习率比能有效缓解耦合,适合优化Transformer
Grokking指神经网络在达到完美训练性能后才出现延迟泛化现象,已在Transformer和MLP中观察到,但其驱动机制尚不明确。本文揭示嵌入层在其中的核心作用:将嵌入引入MLP后,模块化算术任务中出现延迟泛化,而无嵌入的MLP则立即泛化。分析发现两个关键机制:(1) 嵌入更新动态中,稀有词因梯度更新稀疏和权重衰减而停滞;(2) 双线性耦合使嵌入与下游权重交互产生鞍点,增加对初始化的敏感性。为验证,研究了频率感知采样以平衡词元更新、减少梯度方差,并推导出嵌入专用学习率。证明自适应学习率比 η_E/η_W ∝ σ_max(E)/σ_max(W) · f_W/f_E 可缓解双线性耦合,加速收敛。该方法不仅改善grokking行为,还拓展至Transformer优化中的广泛挑战。
原文摘要 · Abstract (English)
Grokking, a delayed generalization in neural networks after perfect training performance, has been observed in Transformers and MLPs, but the components driving it remain underexplored. We show that embeddings are central to grokking: introducing them into MLPs induces delayed generalization in modular arithmetic tasks, whereas MLPs without embeddings can generalize immediately. Our analysis identifies two key mechanisms: (1) Embedding update dynamics, where rare tokens stagnate due to sparse gradient updates and weight decay, and (2) Bilinear coupling, where the interaction between embeddings and downstream weights introduces saddle points and increases sensitivity to initialization. To confirm these mechanisms, we investigate frequency-aware sampling, which balances token updates by minimizing gradient variance, and embedding-specific learning rates, derived from the asymmetric curvature of the bilinear loss landscape. We prove that an adaptive learning rate ratio, \(\frac{η_E}{η_W} \propto \frac{σ_{\max}(E)}{σ_{\max}(W)} \cdot \frac{f_W}{f_E}\), mitigates bilinear coupling effects, accelerating convergence. Our methods not only improve grokking dynamics but also extend to broader challenges in Transformer optimization, where bilinear interactions hinder efficient training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。