通过调整模型结构,让Transformer提前完成学习,突破延迟泛化难题。
The Geometric Inductive Bias of Grokking: Bypassing Phase Transitions via Architectural Topology
- 用球形拓扑约束模型,消除表示幅度自由度
- 移除注意力路由后,所有实验种子均100%泛化
- 结构设计需匹配任务对称性,否则无效
机械可解释性通常依赖训练后分析,本文采用干预式方法:通过修改架构拓扑,先验测试假设并观察训练动态。研究Transformer在循环模加法(Zp)任务中的grokking现象——延迟泛化,探究特定架构自由度是否延长记忆阶段。发现标准Transformer存在两个独立结构因素:无界表示幅度和数据依赖注意力路由。首先引入全有界球形拓扑,强制残差流中L2归一化,并使用固定温度的解嵌入矩阵,消除幅度自由度,使grokking启动时间缩短20倍以上,无需权重衰减。其次,采用均匀注意力消融,以均匀分布替代数据依赖查询-键路由,将注意力层变为连续词袋(CBOW)聚合器。尽管移除了自适应路由,这些模型在所有种子下均实现100%泛化,完全绕过grokking延迟。为检验加速是否仅为通用优化稳定器,使用非交换S5置换组合作为负控实验,对S5施加球形约束并未加速泛化。表明消除记忆阶段高度依赖架构先验与任务内在对称性的对齐。综合结果提供干预证据,证明架构自由度显著影响grokking,支持从结构角度预测训练动态的新视角。
原文摘要 · Abstract (English)
Mechanistic interpretability typically relies on post-hoc analysis of trained networks. We instead adopt an interventional approach: testing hypotheses a priori by modifying architectural topology to observe training dynamics. We study grokking - delayed generalization in Transformers trained on cyclic modular addition (Zp) - investigating if specific architectural degrees of freedom prolong the memorization phase. We identify two independent structural factors in standard Transformers: unbounded representational magnitude and data-dependent attention routing. First, we introduce a fully bounded spherical topology enforcing L2 normalization throughout the residual stream and an unembedding matrix with a fixed temperature scale. This removes magnitude-based degrees of freedom, reducing grokking onset time by over 20x without weight decay. Second, a Uniform Attention Ablation overrides data-dependent query-key routing with a uniform distribution, reducing the attention layer to a Continuous Bag-of-Words (CBOW) aggregator. Despite removing adaptive routing, these models achieve 100% generalization across all seeds and bypass the grokking delay entirely. To evaluate whether this acceleration is a task-specific geometric alignment rather than a generic optimization stabilizer, we use non-commutative S5 permutation composition as a negative control. Enforcing spherical constraints on S5 does not accelerate generalization. This suggests eliminating the memorization phase depends strongly on aligning architectural priors with the task's intrinsic symmetries. Together, these findings provide interventional evidence that architectural degrees of freedom substantially influence grokking, suggesting a predictive structural perspective on training dynamics.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。