Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

The Geometric Inductive Bias of Grokking: Bypassing Phase Transitions via Architectural Topology

Alper Yıldırım

TL;DR

Together, these findings provide interventional evidence that architectural degrees of freedom substantially influence grokking, suggesting a predictive structural perspective on training dynamics.

Abstract

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.

The Geometric Inductive Bias of Grokking: Bypassing Phase Transitions via Architectural Topology

TL;DR

Together, these findings provide interventional evidence that architectural degrees of freedom substantially influence grokking, suggesting a predictive structural perspective on training dynamics.

Abstract

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.
Paper Structure (35 sections, 6 equations, 2 figures, 3 tables)

This paper contains 35 sections, 6 equations, 2 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Path A: The Magnitude Degree of Freedom.
  3. Path B: The Routing Degree of Freedom.
  4. Summary of Contributions.
  5. Related Work
  6. Grokking as Geometric Reorganization
  7. Periodic and Circular Inductive Biases
  8. Fourier Features and Structural Continuity
  9. Formal Problem Setting and Methodology
  10. The Standard Residual Stream (Baselines)
  11. The Spherical Residual Stream and Fully Bounded Topology (Intervention A)
  12. Fully Bounded Topology.
  13. Uniform Attention Routing (Intervention B)
  14. Task 1: The Cyclic Group $Z_p$ (Geometric Alignment and Attention)
  15. Task 2: The Symmetric Group $S_5$ (Negative Control)
  16. ...and 20 more sections

Figures (2)

  • Figure 1: Test accuracy across architectural configurations. (a) Standard normalizations exhibit the classic, delayed grokking profile over 100,000 epochs. (b) A zoomed-in view of the first 5,000 epochs reveals that the topologically bounded models bypass this delay entirely, exhibiting immediate and stable convergence while the baselines remain trapped at chance accuracy.
  • Figure 2: Training dynamics under Uniform Attention Ablation (Seed 2). (a) The LayerNorm baseline immediately generalizes when data-dependent routing is removed, bypassing grokking. (b) Imposing a spherical constraint while retaining weight decay introduces optimization instability, delaying generalization. (c) The Fully Bounded topology (zero weight decay) stabilizes the geometry, resulting in immediate and flawless generalization.