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

이 논문은 트랜스포머 아키텍처의 무제한 표현 크기 및 데이터 의존적 어텐션 경로와 같은 구조적 자유도를 제거하여 모듈러 덧셈 작업에서 '그로킹 (grokking)' 현상을 유발하는 기억 단계를 우회하고 일반화를 즉각적으로 달성할 수 있음을 보이며, 이는 작업의 고유한 대칭성과 아키텍처의 기하학적 편향이 정렬될 때 학습 역학이 크게 개선됨을 시사합니다.

핵심

이 논문은 인공지능 (AI) 이 문제를 풀 때 겪는 아주 특이한 현상인 **'그로킹 (Grokking)'**에 대해 연구한 것입니다.

한마디로 요약하면: **"AI 가 문제를 외우는 데는 시간이 걸리지만, 진짜 원리를 깨닫는 데는 훨씬 더 오래 걸립니다. 하지만 AI 의 '머리 구조'를 조금만 바꿔주면, 그 기다리는 시간을 20 배 이상 단축시킬 수 있다"**는 놀라운 발견을 담고 있습니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. '그로킹'이란 무엇인가요? (외우기 vs 깨닫기)

상상해 보세요. 어떤 학생이 수학 문제를 풀고 있습니다.

• 초반: 문제를 계속 풀다 보니, 정답을 외워버립니다. 시험지에는 모든 문제를 맞췄지만, 그건 단순히 암기한 것뿐이라서 새로운 문제가 나오면 전혀 못 풉니다. (이게 '기억 단계')
• 중반: 계속 공부하다가 갑자기 아하! 하는 순간이 옵니다. 문제의 **원리 (공식)**를 깨닫는 거죠.
• 결과: 원리를 깨닫는 순간부터는 새로운 문제도 척척 맞춥니다. (이게 '일반화' 또는 '그로킹' 단계)

문제는 이 '아하!' 순간이 오기까지 엄청나게 오래 걸린다는 것입니다. AI 는 수천 번의 실수를 반복하며 원리를 찾아내는데, 그 사이에 이미 외운 답을 계속 반복해서 틀린 답을 고집합니다.

2. 연구자의 질문: "왜 AI 는 이렇게 느릴까요?"

연구자는 "AI 가 원리를 깨닫기까지 왜 이렇게 오래 걸리는 걸까?"라고 의문을 품었습니다.
기존의 AI(트랜스포머) 는 아주 유연한 구조를 가지고 있습니다. 마치 자유로운 상상력을 가진 화가처럼, 그림을 그릴 때 색상의 농도 (크기) 를 마음대로 조절하고, 어떤 부분을 집중해서 볼지 (주목) 도 스스로 결정합니다.

연구자는 "아마도 이 너무나 자유로운 능력이 오히려 독이 되는 것 같다"고 추측했습니다.

• AI 가 원리를 찾기 전에, 일단 **무작정 외우는 방법 (기억)**을 먼저 찾아버리기 때문입니다.
• 마치 학생이 공식을 외우기 전에, 문제집의 답지를 암기해버리는 것과 같습니다.

3. 해결책: "AI 의 자유를 제한하자!" (두 가지 실험)

연구자는 AI 의 자유를 일부러 제한하는 두 가지 실험을 했습니다.

실험 A: "크기 제한" (구형 위상)

• 비유: AI 가 그림을 그릴 때, 붓의 세기나 색의 농도를 마음대로 조절하지 못하게 하고, 무조건 같은 크기의 원 안에만 그림을 그리게 했습니다.
• 효과: AI 는 더 이상 "크기"로 정보를 저장할 수 없게 되자, **방향 (각도)**으로만 정보를 표현해야 했습니다. 이는 수학적인 원리 (삼각함수) 와 딱 맞는 구조였습니다.
• 결과: AI 가 외우는 단계를 건너뛰고, 원리를 깨닫는 순간이 20 배 이상 빨라졌습니다. 마치 답지 암기를 금지하고 공식만 쓰게 했더니, 학생이 순식간에 공식을 이해한 것과 같습니다.

실험 B: "주목 제한" (균일한 주의)

• 비유: AI 가 문제를 읽을 때, "이 단어가 중요해!", "저 단어가 중요해!"라고 스스로 판단하게 하지 않고, 모든 단어를 똑같이 중요하게 취급하게 했습니다. (예: "A + B = ?"에서 A 와 B 를 똑같이 1/3 씩만 봄)
• 효과: AI 가 특정 단어를 골라 기억하는 (암기) 능력을 없애버렸습니다.
• 결과: AI 는 더 이상 특정 조합을 외울 수 없게 되자, 원리 자체를 바로 깨닫고 문제를 풀기 시작했습니다.

4. 중요한 반전: "이게 모든 문제에 통할까요?" (S5 실험)

연구자는 "혹시 이 방법이 AI 를 안정화시키는 일반적인 마법약일까?"라고 의심했습니다. 그래서 비교군으로 '순열 조합 (S5)'이라는 아주 다른 문제를 풀게 했습니다. 이 문제는 덧셈처럼 단순한 원리가 아니라, 서로 순서가 바뀌면 결과가 달라지는 복잡한 규칙을 가지고 있습니다.

• 결과: 여기서 '크기 제한'을 적용하자 AI 는 완전히 망가졌습니다.
• 이유: 복잡한 문제는 AI 가 다양한 크기와 방향으로 정보를 표현해야 해결할 수 있는데, 연구자가 만든 '원형의 틀'이 너무 좁아서 AI 가 복잡한 구조를 만들지 못했기 때문입니다.

5. 결론: "맞춤형 키"가 필요합니다

이 논문의 핵심 메시지는 다음과 같습니다.

"AI 가 어떤 문제를 풀 때, 그 문제의 '수학적 성질'과 AI 의 '머리 구조'가 딱 맞아떨어져야 합니다."

• 단순한 규칙 (덧셈 등): AI 에게 자유를 제한하고, 문제의 원리 (원형 대칭) 에 맞춰 구조를 잡아주면, AI 는 외우는 시간을 건너뛰고 순식간에 깨닫습니다.
• 복잡한 규칙 (순열 등): 같은 제한을 주면 오히려 AI 가 답을 못 찾습니다.

요약

이 연구는 AI 가 왜 "아하!" 순간을 늦게 경험하는지 그 이유를 구조적 문제로 설명했습니다. 그리고 AI 를 단순히 "더 많이 훈련"시키는 것이 아니라, 문제에 맞는 구조로 설계해 주면, AI 는 훨씬 더 빠르고 효율적으로 지능을 발휘할 수 있음을 증명했습니다.

마치 비행기를 만들 때, 물고기처럼 지느러미를 달지 않고 날개를 달아야 하늘을 날 수 있는 것과 같습니다. 문제의 성질에 맞는 '날개 (구조)'를 설계해 주는 것이 AI 를 더 똑똑하게 만드는 지름길입니다.

기술 요약

1. 연구 배경 및 문제 제기 (Problem)

• 그로킹 (Grokking) 현상: 신경망이 훈련 데이터에 대한 정확도는 거의 완벽하게 달성하지만, 테스트 데이터에 대한 일반화 능력은 오랫동안 낮은 수준에 머무르다가 최적화 과정이 길어지면 갑자기 급격히 향상되는 지연된 일반화 현상입니다.
• 기존 연구의 한계: 기존의 기계적 해석 (Mechanistic Interpretability) 연구는 대부분 훈련이 완료된 모델의 가중치를 사후 (post-hoc) 분석하여 그로킹이 발생한 후의 메커니즘을 규명하는 데 집중했습니다.
• 핵심 가설: 표준 트랜스포머 아키텍처는 모듈러 덧셈 (Modular Addition, $Z_p$) 과 같은 교환적이고 주기적인 작업의 최소 대칭성 요구사항을 초과하는 **과도한 표현의 자유도 (Excess Degrees of Freedom)**를 가지고 있습니다. 이러한 과도한 자유도 (예: 무제한 벡터 크기, 데이터 의존적 어텐션 라우팅) 가 모델이 구조화된 해법 (푸리에 회로) 대신 기억 (Memorization) 중심의 해법을 선택하게 하여 그로킹 지연을 유발한다는 가설을 세웠습니다.

2. 방법론 (Methodology)

저자는 훈련 후 분석이 아닌, 훈련 전 아키텍처 구조를 개입 (Interventional Approach) 하여 수정하는 방식을 취했습니다. 두 가지 독립적인 구조적 요인을 제거하거나 제한하여 그로킹 동역학에 미치는 영향을 관찰했습니다.

A. 개입 1: 구형 잔여 스트림 (Spherical Residual Stream) - 크기 자유도 제한

• 문제: 표준 트랜스포머에서는 잔여 스트림 (Residual Stream) 의 벡터 방향뿐만 아니라 **크기 (Magnitude)**에도 정보가 인코딩될 수 있습니다. 이는 무제한의 크기 성장을 허용하여 '피자 (Pizza)'와 같은 조각난 기억 기반 해법을 가능하게 합니다.
• 해결책: **완전 경계 구형 토폴로지 (Fully Bounded Spherical Topology)**를 도입했습니다.
  • 모든 잔여 스트림에 L2 정규화를 적용하여 벡터 크기를 1 로 고정합니다.
  • 언임베딩 (Unembedding) 행렬도 정규화하고, 로그트 (Logits) 를 스케일된 코사인 유사도로 계산하여 크기 기반의 자유도를 완전히 제거합니다.
  • 이를 통해 모델이 벡터의 방향 (각도) 관계만으로 정보를 인코딩하도록 강제하여, 작업의 본질적인 주기적 대칭성 (원형 기하학) 과 정렬되도록 합니다.

B. 개입 2: 균일 어텐션 제거 (Uniform Attention Ablation) - 라우팅 자유도 제한

• 문제: 표준 어텐션 메커니즘은 데이터에 의존하는 복잡한 쿼리 - 키 (Query-Key) 라우팅을 학습할 수 있습니다. 그러나 모듈러 덧셈은 이론적으로 균일한 '토큰 가방 (Bag-of-Tokens)' 집계만으로도 해결 가능합니다.
• 해결책: 균일 어텐션 제거를 적용했습니다.
  • 학습된 어텐션 점수를 무시하고, 시퀀스 내 모든 토큰에 대해 고정된 **균일 분포 (예: [1/3, 1/3, 1/3])**를 강제합니다.
  • 이를 통해 어텐션 레이어를 적응형 라우팅이 없는 연속형 단어 주머니 (CBOW) 집계기로 축소하여, 대칭성을 깨는 라우팅 경로를 제거했습니다.

C. 부정적 통제 (Negative Control): $S_5$ 순열 합성

• 모듈러 덧셈 ($Z_p$, 교환법칙 성립) 에서의 가속화가 특정 작업의 기하학적 정렬 때문인지, 아니면 일반적인 최적화 안정화 효과인지 구분하기 위해 비교환적 (Non-commutative) 인 대칭군 $S_5$ 작업에 동일한 구형 제약을 적용했습니다.

3. 주요 결과 (Key Results)

A. 모듈러 덧셈 ($Z_p$) 실험 결과

1. 그로킹 지연의 극적인 단축:
   • 기저선 (LayerNorm/RMSNorm): 평균 약 54,160 에포크 후에 일반화가 시작되었습니다.
   • 완전 경계 구형 토폴로지: 가중치 감쇠 (Weight Decay) 없이도 평균 2,100 에포크 만에 100% 일반화에 도달했습니다. 이는 20 배 이상의 가속화입니다.
   • 균일 어텐션: 데이터 의존적 라우팅을 제거한 모델 또한 모든 시드에서 그로킹 지연 없이 즉시 일반화했습니다.
2. 최적화 안정성:
   • 구형 제약만 적용하고 가중치 감쇠를 유지한 경우 ($\lambda=1.0$) 는 로그트 크기 증가와 가중치 감쇠 간의 충돌로 인해 최적화가 불안정해졌습니다.
   • 반면, **완전 경계 토폴로지 (언임베딩 정규화 + 가중치 감쇠 0)**는 안정적인 최적화를 가능하게 하여 그로킹을 완전히 우회했습니다.
3. 푸리에 회로 검증:
   • 스펙트럼 분석 (FFT) 을 통해 가속화된 모델들도 기존 연구에서 확인된 **푸리에 회로 (Fourier Circuit)**를 사용함을 확인했습니다. 즉, 새로운 단축경로를 찾은 것이 아니라, 구조적 제약이 모델이 더 일찍 올바른 기하학적 해법에 도달하도록 유도한 것입니다.

B. $S_5$ 순열 합성 실험 결과 (부정적 통제)

• $Z_p$에서는 구형 제약이 일반화를 가속시켰지만, 비교환적 인 $S_5$ 작업에서는 구형 제약이 적용된 모든 모델이 100,000 에포크 동안 일반화에 실패했습니다.
• 이는 구형 제약이 일반적인 최적화 도구가 아니라, 작업의 내재적 대칭성 (주기성/교환성) 과 정렬될 때만 유효한 기하학적 귀납적 편향임을 증명합니다. $S_5$는 고차원 비가환적 구조가 필요하므로, 1 차원 원형 기하학으로 제한된 모델은 이를 표현할 수 없었습니다.

4. 주요 기여 (Key Contributions)

1. 개입적 해석 방법론의 제시: 사후 분석에서 벗어나, 아키텍처의 구조적 자유도를 사전에 제어하여 가설을 검증하는 개입적 (Interventional) 접근법을 제시했습니다.
2. 그로킹의 기하학적 본질 규명: 그로킹 지연이 최적화의 불가피한 단계가 아니라, 과도한 아키텍처 자유도로 인한 기억 기반 해법과 구조화된 해법 사이의 전환 지연임을 입증했습니다.
3. 작업 특이적 구조 정렬 (Task-Specific Structural Alignment): 아키텍처의 기하학적 제약 (구형 토폴로지) 이 작업의 수학적 대칭성 (주기성) 과 정렬될 때만 그로킹이 제거됨을 보여주었습니다. 이는 "하나의 만능 아키텍처"가 아니라 "작업에 맞는 구조적 정렬"의 중요성을 강조합니다.
4. 예측적 아키텍처 설계: 기계적 해석을 통해 작업의 본질적 구조를 파악하고, 이를 아키텍처 설계에 반영하여 학습 동역학을 예측하고 제어할 수 있음을 보였습니다.

5. 의의 및 결론 (Significance)

이 연구는 그로킹이 단순한 최적화 현상이 아니라, 표현 기하학의 재배열 과정임을 보여줍니다. 표준 트랜스포머가 가진 무제한 크기 표현과 복잡한 어텐션 라우팅이 오히려 기억 (Memorization) 을 유도하여 일반화를 지연시킨다는 점을 규명했습니다.

• 실용적 시사점: 알고리즘적 추론이나 주기적 패턴이 중요한 작업 (예: 시간 계열 예측, 수학적 추론) 에서는 아키텍처에 명시적인 기하학적 제약 (예: 구형 정규화, 주기적 어텐션) 을 도입함으로써 학습 효율을 극대화하고 그로킹을 방지할 수 있습니다.
• 이론적 확장: 이는 "The Bitter Lesson" (데이터와 연산량에 의존하는 것) 과 모순되지 않으면서도, 지식 (작업의 수학적 구조) 을 아키텍처에 인코딩함으로써 더 효율적이고 해석 가능한 모델을 설계할 수 있는 새로운 패러다임을 제시합니다.

결론적으로, 이 논문은 아키텍처 토폴로지가 학습 동역학의 핵심 결정 요인이며, 작업의 대칭성과 아키텍처의 기하학을 정렬시키는 것이 지연된 일반화를 해결하는 열쇠임을 실험적으로 증명했습니다.