Grokking의 끝은 대수적 표현 가능성인가

들어가며

Grokking은 신경망이 먼저 훈련 데이터를 외운 뒤, 한참 후에야 일반화 가능한 규칙을 갑자기 학습하는 현상으로 알려져 있다. 특히 모듈러 산술 과제는 이 현상을 연구하기 좋은 실험장으로 쓰여 왔다. 모델 용량이 작으면 암기조차 어렵고, 충분히 크면 일반화가 훨씬 빨리 나타난다는 점도 관찰되어 왔다.

이 논문은 그 스펙트럼의 극단을 묻는다. 만약 아키텍처가 표현할 수 있는 함수 집합이 자유롭게 커지는 대신, 유한 차원의 대수적 객체로 강하게 제한된다면 grokking은 여전히 나타날까?

핵심 내용

• 정확히 분석 가능한 설정: 저자들은 활성화 함수로 전정칙 단항식 σ(z)=z^k를 사용하는 2층 네트워크를 다룬다. 모듈러 산술 과제는 근의 단위를 통해 인코딩되며, 네트워크 출력은 이산 푸리에 문자로 기술할 수 있다.
• 폭을 키워도 표현 공간은 그대로: 은닉층 폭과 상관없이 출력은 ((\mathbb{Z}_p)^2) 위 문자들의 ((k+1))차원 부분공간에 머문다. 이는 전체 함수 공간 중 (O(k/p^2))에 해당하는 작은 조각이다.
• 표현 가능성의 정확한 판정식: 어떤 과제가 표현 가능하려면, 그 이산 푸리에 지지가 (u+v=k \pmod p) 대각선 위에 있어야 하며, 그 경우에만 가능하다. 선형 위상 목표에서는 이 조건이 (m+n=k)로 단순화된다.
• 실패는 최적화 문제가 아니다: 조건을 만족하지 않는 목표는 일반화에 실패하는 수준을 넘어, 훈련 세트에도 맞출 수 없다. 논문은 은닉 폭에 의존하지 않는 양의 훈련 손실 하한도 증명한다.
• 실험은 이분법적 결과를 보인다: 585회 실행에서 대수적 예측과 관측 결과가 99.8% 일치했다. 관찰된 것은 ‘암기 후 일반화’가 아니라, 즉시 성공하거나 완전히 실패하는 양상이다.

의미와 영향

이 연구는 grokking을 표현 가능성의 관점에서 극한까지 밀어붙인다. 일반적인 설명에서는 네트워크가 언제 암기에서 일반화로 전환되는지가 중요하다. 하지만 이 모델에서는 질문이 달라진다. 목표 함수가 애초에 아키텍처의 대수적 함수 클래스 안에 있는가가 전부에 가깝다.

물론 실제 표준 신경망이 항상 이렇게 흑백 논리로 움직인다는 뜻은 아니다. 논문은 병목 제거 실험을 통해 이 극단적 설정과 일반 네트워크를 연결한다. 표현 불가능 상태에서, 일반화 없는 암기, 그리고 용량 증가에 따라 간격이 줄어드는 grokking으로 이어지는 연속적인 경로를 제시한다.

따라서 이 모델은 이론적 현미경에 가깝다. 훈련 동역학 때문에 일반화가 늦게 나타나는 경우와, 아키텍처 자체가 표현할 수 없어 실패하는 경우를 분리해 보여준다. 합성 과제와 푸리에 구조, 용량 기반 일반화를 연구하는 데 중요한 기준점을 제공한다.

출처: arXiv