Where Reusable Computation Becomes Detectable: Solution-Frame Path Triage for Modular-Arithmetic Grokking
Ruian Lei ⋅ Ruoxi Jiang ⋅ Marley M Vellasco ⋅ M. Tanveer ⋅ Zenglin Xu
Abstract
Mechanistic studies of grokking often begin after an endpoint circuit has been identified. We study the preceding retrospective triage problem: given a solved final checkpoint and a long saved trajectory, which path and window should be inspected first? We propose solution-frame path triage. The final checkpoint fixes a solution frame; earlier checkpoints are then scored path by path. The Geometric Coherence Score (GCS) measures coherence in this procedure: it asks whether neighboring inputs in the final frame undergo similar local Jacobian transformations through the attention, MLP, or full block measurement maps. Centered Kernel Alignment (CKA) and $L_2$ distance to the final path state provide the complementary proximity scores. Across 104 modular arithmetic Transformer runs, GCS reveals a repeatable path schedule that simple norm and dimensionality summaries miss: attention coherence rises near grokking and often turns over, while MLP maps diversify and later reconverge. After selection, held-out checks show that the resulting windows are enriched for accuracy and algebraic changes, align with Fourier progress on modular addition, and identify attention-to-MLP interface states with high replacement cost. The method returns a short list of path/window hypotheses for endpoint analysis. It is retrospective by design: the final checkpoint is part of the procedure.
Chat is not available.
Successful Page Load