Spherical transformer geometry speeds up grokking in arithmetic tasks
Dissecting Geometric Constraints in Grokking: Interaction-Dependent Acceleration in Modular Arithmetic
Machine LearningComputational Complexity
Summary
Some AI models suddenly learn a task very quickly after a long period of slow progress, a phenomenon called grokking. This paper looks at how changing certain parts of a specialized model called a spherical transformer affects how fast grokking happens in modular arithmetic tasks. The researchers found that combining several specific changes speeds up learning, but changing just one part isn't enough. The speedup depends on how these changes interact and doesn't happen in every task or model size tested.
What this means in practice
- •For machine learning engineers: Optimize training setups by tuning spherical transformer components to speed up learning in modular arithmetic models.
- •For ai system architects: Design modular arithmetic AI systems that reliably achieve faster learning by combining specific model modifications.
Authors
Zachary A. Hampton, Wasiq Khan, Bilal Khan
Abstract
Spherical transformer configurations can accelerate grokking through simultaneous modifications of several model properties. The present study investigated the contributions of these modifications in the published Sphere+Fourier configuration using 960 controlled modular-arithmetic runs. A complete factorial analysis varied residual projection, readout weight normalization, logit scale, and weight decay. Across four matched settings with weight decay, the successful-run mean crossing times without projection were 8-37 times those observed with projection. However, residual projection was insufficient to induce grokking without additional interventions. Various other unprojected configurations also generalized. Since projection altered the magnitude entering the readout, these comparisons could not separate geometry from output scale. A separate 100-run fixed-horizon study demonstrated faster first crossing for the complete configuration across five train/test partitions. Faster crossing did not necessarily correspond to greater persistence over extended training. Control experiments demonstrated rapid grokking without Fourier initialization. Task and depth sweeps also identified conditions in which the configuration failed to reach the criterion by 50K epochs. Accordingly, the observed acceleration depended on interactions among the components and did not extend uniformly across the tested tasks or depths.