Differentiable Algebra Discovery Accelerates Grokking via a Non-Fourier Mechanism

Motivation

• Known: grokking 현상은 modular arithmetic 등에서 MLP가 지연 일반화(delayed generalization)를 보이며, Nanda et al. (2023)의 mechanistic account에 따르면 순환군(cyclic group)에 대해 sparse Fourier representation으로 수렴한다는 것이 알려져 있다. 또한 Grokfast 같은 gradient EMA 기법이 grokking을 가속하고, group-equivariant network나 tensor decomposition 연구도 별도로 존재한다.
• Gap: 기존 mechanistic grokking 연구는 순환군과 sparse Fourier representation에 국한되어 있어, 학습된 algebraic prior가 임의의 유한군(비가환군 포함)으로 일반화되고 수렴을 극적으로 가속할 수 있는지는 미해결 문제였다.
• Why: grokking의 대수적 구조 발견 메커니즘을 순환군을 넘어 임의의 유한군으로 확장하고, 이를 CP factorization이라는 수학적으로 엄밀한 틀로 정식화함으로써, implicit algebra discovery·물리 기반 불변량 보존·equivariant architecture라는 세 연구 흐름을 통합할 수 있는 일반적 inductive bias를 제시하기 때문이다.
• Approach: 저자들은 associativity, identity, inverse loss로 공동 훈련되는 rank-R bilinear product µ(a,b)=Σ_r W_r((U_r a)⊙(V_r b))를 제안하고, 이를 group multiplication tensor T_G의 CP factorization으로 수학적으로 특징짓는 6개의 명제를 증명한 뒤 다양한 유한군과 실제 물리/화학 태스크에서 검증한다.

Achievement

Figure 4. Grokking curves on four non-abelian groups (mean over 3 seeds, clipped at 15,000 steps). FORGE (blue) grokks w

1. grokking 가속 성능: Z/97Z에서 FORGE+Grokfast가 10.20× 속도향상(MLP+Grokfast는 오히려 0.86×로 느려짐)을 달성했고, S3·D4·A4에서는 매칭된 MLP가 9개 런 전부 val=0.000으로 실패한 반면 FORGE는 ~10³ step 내 grokking에 성공했다.
2. 비Fourier 메커니즘 발견: 6개 seed에 걸쳐 FORGE는 MLP와 동일한 정확도를 달성하면서도 sparse Fourier route가 아닌 질적으로 다른 경로(effective harmonic modes ≈23 vs ≈16, sign test p<0.016)를 사용함을 규명했다.
3. 보편적 non-abelian grokking: abelian, dihedral, alternating, symmetric, quaternionic 등 A7(order 2,520), Z1009(order 1,009)를 포함한 모든 테스트 유한군에서 grokking에 성공했으며, 최소 비가해군(non-solvable group)인 A5에서 MLP 대비 10.1×, MLP+Grokfast 대비 3.0× 속도향상을 보였다.
4. 보편 스케일링 법칙: grokking 시간이 732·|G|^0.170 (R²=0.785, 20개 군, 68 seed)이라는 sub-linear power law를 따름을 실증하여, 군의 order가 168배 증가해도 step 수는 2.4배만 증가함을 보였다.
5. 이론적 특성화: Strassen-refined rank bound(rank_CP(T_{S4})≤55<64), mechanistic Wedderburn recovery(1,639개 conjugacy-class multiplicity 중 약 1,630개 정확히 복원), causal axiom-emergence theorem(identity axiom이 일반화보다 583±80 step 앞서 발생, tautology hypothesis 반박) 등 6개 명제를 제시했다.
6. 도메인 일반화: Hamiltonian 시스템에서 invariant drift를 20–150,000× 감소시키고, QM9 U0 MAE를 15.3% 개선하여 group theory를 넘어선 general-purpose inductive bias임을 입증했다.

How

Figure 4. Grokking curves on four non-abelian groups (mean over 3 seeds, clipped at 15,000 steps). FORGE (blue) grokks w

• DAD(Differentiable Algebra Discovery) module을 rank-R Hadamard decomposition µ(a,b)=Σ_r W_r(U_r a ⊙ V_r b)으로 parametrize하고, 이 module에 대해 Lassoc, Lident, Linv 3가지 algebra loss를 정의
• 학습된 identity element e와 inverse map g_φ를 두어 group axiom을 명시적으로 강제
• 전체 objective는 L = Ltask + λalg(t)(Lassoc+Lident+Linv)로, λalg(t)를 선형 감쇠시켜 초기에는 구조 학습을, 후반에는 태스크 학습을 우선시
• fiber embedding과 DAD 파라미터에 대해 weight decay를 분리한 split weight decay scheme을 도입해 "un-grokking" 현상(구조가 점진적으로 지워지는 현상)을 방지
• 4-way matched-budget ablation(MLP, MLP+Grokfast, FORGE, FORGE+Grokfast)으로 architecture의 필수성을 검증
• Fourier analysis를 통해 grokked embedding의 spectral mass 분포를 6개 독립 seed에 걸쳐 비교하고 sign test로 통계적 유의성 확인
• Strassen 및 Laderman의 matrix-multiplication algorithm을 활용해 S4의 CP rank 상한을 개선하는 이론적 증명 제시
• Hamiltonian dynamics(3개 시스템)와 QM9 분자 특성 예측(U0)에 적용하여 group theory 밖에서의 일반화 검증
• capacity-lift sweep(S4에서 R·d 변화)을 통해 grokking phase transition을 관찰하여 optimization hardness를 정량화

Originality

• 기존 grokking 연구가 순환군과 sparse Fourier representation에 국한된 것과 달리, 임의의 유한군(비가환군 포함)에 대해 학습 가능한 algebraic prior를 CP factorization으로 정식화한 최초의 접근
• MLP와 동일한 태스크 정확도를 달성하면서도 질적으로 다른 non-Fourier 표현 경로를 사용함을 실증적으로 규명, Nanda et al.(2023)의 mechanistic account를 "여러 대수적으로 올바른 경로가 존재하며 inductive bias가 어떤 경로를 찾을지 결정한다"는 방향으로 확장
• Strassen과 Laderman의 matrix-multiplication algorithm을 CP rank 상한 개선에 활용한 점은 대수적 복잡도 이론(algebraic complexity theory)과 mechanistic interpretability를 연결하는 독창적 시도
• causal axiom-emergence theorem을 통해 tautology hypothesis(axiom 검증이 이미 일반화된 후 사후적으로 나타난다는 가설)를 인과적으로 반박한 점
• group theory 태스크를 넘어 Hamiltonian dynamics, 분자 특성 예측까지 확장하여 differentiable algebra discovery를 general-purpose inductive bias로 제안

Limitation & Further Study

• 본문 발췌에 실험 세부사항(하이퍼파라미터 전체, 데이터 분할, 통계적 유의성 검정의 세부 방법론)이 충분히 제시되지 않아 재현성 검증이 어려움
• sign test 기반 유의성 검증이 6개 seed라는 비교적 작은 표본에 의존하고 있어, 더 많은 seed와 다양한 아키텍처 변형에 대한 강건성 검증이 필요
• rank R과 dimension d 등 hyperparameter 선택(capacity threshold ~256)이 태스크별로 어떻게 결정되는지, 더 큰 군이나 비군(non-group) 대수 구조로의 일반화 가능성에 대한 추가 연구가 필요
• Hamiltonian 및 QM9 실험이 group theory 실험만큼 포괄적인 ablation을 거쳤는지 불분명하며, 다른 도메인(자연어, 비전 등)으로의 확장 가능성에 대한 논의가 부족
• CP factorization의 이론적 보장(rank bound 등)이 특정 군(S4 등)에 한정되어 있어, 일반적 유한군에 대한 이론적 상한을 제시하는 후속 연구가 필요

Evaluation

총평: grokking mechanistic interpretability를 순환군을 넘어 임의의 유한군으로 확장하고 이를 CP factorization이라는 엄밀한 수학적 틀로 정식화한 뛰어난 연구로, 이론과 실증을 균형있게 결합했다는 점에서 높은 평가를 받을 만하지만, 재현성과 일반화 범위에 대한 추가 검증이 필요하다.