Differentiable Algebra Discovery Accelerates Grokking via a Non-Fourier Mechanism
Abstract
Neural networks can implicitly discover algebraic structure through the grokking phenomenon, but prior mechanistic accounts are limited to cyclic groups and sparse Fourier representations; whether a learned algebraic prior can generalize to arbitrary finite groups and dramatically accelerate convergence remains open. We introduce FORGE: a rank-R bilinear product μ(a,b)=Σᵣ Wᵣ((Uᵣa)⊙(Vᵣb)) trained jointly with associativity, identity, and inverse algebra losses, which we prove is a CP factorization of the group multiplication tensor TG. Six propositions characterize the mechanism: a Strassen-refined rank bound (rankCP(T_{S₄})≤55 < 64 = Σρ dim(ρ)³, improving the Wedderburn upper bound via Strassen's and Laderman's matrix-multiplication algorithms); mechanistic Wedderburn recovery; CP–isotypic alignment and Frobenius–Schur discrimination; a causal axiom-emergence theorem; and a universal sub-linear grokking-time scaling law. A four-way matched-budget ablation isolates the architecture as essential: on ℤ/97, MLP+Grokfast achieves 0.86× (a slowdown), while FORGE+Grokfast achieves 10.20×; matched MLPs fail on S₃, D₄, A₄ in all 9 runs (val=0.000), while FORGE groks all three in ∼10³ steps via a qualitatively distinct non-Fourier route (6 seeds, sign test p<0.016): effective harmonic modes ≈23 vs. ≈16. FORGE groks every finite group tested—abelian, dihedral, alternating, symmetric, and quaternionic—through A₇ (order 2,520) and ℤ{1009} (order 1,009); on A₅ (smallest non-solvable group) FORGE achieves 10.1× speedup over MLP (1,800 vs. 18,267 steps) and 3.0× over MLP+Grokfast (5,467 steps). Grokking time follows a universal power law 732·|G|^{0.170} (R²=0.785, 20 groups, 68 seeds): a 168× increase in group order costs only 2.4× more steps. Mechanistic analysis recovers ≈1,630 of 1,639 Wedderburn conjugacy-class multiplicities exactly; identity axiom emergence causally precedes generalization by 583±80 steps in all 12 seed-runs, refuting the tautology hypothesis. Beyond group theory, FORGE reduces Hamiltonian invariant drift by 20–150,000× and improves QM9 U₀ MAE by 15.3%, establishing differentiable algebra discovery as a general-purpose inductive bias.