A Stochastic--Geometric Theory of Scaling Laws in Grokking

비디오 게임 캐릭터가 까다로운 퍼즐을 풀려고 애쓰는 모습을 보고 있다고 상상해 보세요. 처음에는 캐릭터가 방금 지나온 경로를 하나하나 다 외우려고 미친 듯이 사방을 뛰어다닙니다. 순식간에 점수를 0으로 만들지만, 규칙을 실제로 배운 것은 아닙니다. 그저 특정 지도를 암기했을 뿐이죠. 이것을 **암기(memorization)**라고 부릅니다.

그러다 마법 같은 일이 일つ 벌어집니다. 아무것도 변하지 않는 것 같은 길고 지루한 시간이 흐른 뒤, 캐릭터가 갑 Pap! 하고 멈춰 서더니, 판을 쳐다보고는 깨달음을 얻습니다. 이제 캐릭터는 처음 보는 새로운 퍼즐들도 완벽하게 풀어내기 시작합니다. 이 갑작스러운 "아하!" 모먼트가 바로 연구자들이 **그로킹(grokking)**이라고 부르는 현상입니다.

오랫동안 사람들은 왜 이런 지연 현상이 발생하는지 알지 못했습니다. 마법이었을까요? 아니면 글리치(오류)였을까요? 이 논문에서 저자들은 컴퓨터의 '두뇌'(신경망) 내부에서 어떤 일이 일어나고 있는지 시각화할 수 있는 새로운 방법, 즉 모양과 거리의 지도를 제안합니다.

솔루션의 양파 (The Onion of Solutions)

저자들은 컴퓨터의 '두뇌'가 존재할 수 있는 공간이 마치 양파나 과녁처럼 중첩된 구(sphere)들의 집합체와 같다고 제안합니다.

1. 바깥쪽 껍질 (초기화): 컴퓨터가 시작될 때, 설정값들은 얇은 바깥쪽 껍질 위에 흩어져 있습니다. 이는 거대한 풍선의 아주 가장자리에 다트를 던져 맞히는 것과 같습니다.
2. 중간 껍질 (암기): 컴퓨터가 학습함에 따라, 설정값들은 빠르게 중간 층으로 미끄러져 내려갑니다. 여기에서 컴퓨터는 훈련 데이터를 완벽하게 암기합니다. 이는 규칙은 이해하지 못한 채 지도만 외운 우리 게임 캐릭터와 같습니다. 컴퓨터는 여기서 오랫동안 갇혀 있게 됩니다.
3. 핵심부 (일반화): 양파의 깊은 중심부에는 "일반화"의 핵심이 자리 잡고 있습니다. 이곳은 컴퓨터가 실제로 규칙을 이해하고 새로운 문제를 해결할 수 있는 최적의 지점입니다.

"그로킹" 현상은 단순히 중간 껍질에서 빈 공간을 지나 핵심부로 이동하는 여정입니다. 미스터리는 이것이였습니다: 왜 여기까지 가는 데 그렇게 오래 걸리는가?

취객의 걸음걸이와 자기력 (The Drunk Walk and the Magnetic Pull)

저자들은 컴퓨터의 학습 과정(Adam이라는 옵티마이저를 사용함)이 두 가지 힘의 조합이라고 설명합니다.

• 취객의 걸음걸이 (확산, Diffusion): 컴퓨터는 작은 데이터 배치(batch)로부터 학습하기 때문에, 그 경로는 마치 직선을 걸으려는 취객처럼 약간 비틀거립니다. 이 비틀거림은 무작위적입니다.
• 자기력 (표류, Drift): 또한, 컴퓨터를 중심부로 끌어당기는 꾸준한 힘도 존재합니다. 특히 $\ell_2$ 정규화(설정값을 작게 유지하려는 부드러운 자석 역할을 함)라는 규칙 때문입니다.

저자들은 컴퓨터가 중간 껍질에 갇혀 있는 이유가 "취객의 걸음걸이"가 그 층을 빠져나가기에는 너무 약하고, "자기력"은 즉시 중심부로 끌어당길 만큼 강하지 않기 때문이라고 제안합니다. 컴퓨터는 무작위적인 비틀거림이 자신을 충분히 멀리 밀어내어 핵심부로 떨어지게 할 때까지 중간 껍질을 한참 동안 배회해야 합니다.

게임의 규칙 (스케일링 법칙, Scaling Laws)

저자들은 단순히 추측한 것이 아니라, 수학(구체적으로는 정지 시간 이론(stopping-time theory)과 확률 미분 방정식(stochastic differential equations))을 사용하여 이 지연 시간이 정확히 얼마나 걸릴지 예측했습니다. 그들은 이 여정의 속도를 조절하는 세 가지 주요 "노브(조절 손잡이)"를 찾아냈습니다.

1. 학습률 ($\eta$): 이것은 컴퓨터가 내딛는 보폭의 크기입니다. 논문은 만약 보폭을 너무 작게 만들면 중간 껍질을 탈출하는 데 영원히 걸릴 것이라고 보여줍니다. 반대로 너무 크게 만들면 목표를 지나칠 수 있습니다. 여기에는 "골디락스(적당한)" 존이 존재합니다.
2. 배치 크기 ($b$): 이것은 컴퓨터가 한 걸음을 내딛기 전 살펴보는 예시의 개수입니다. 논문은 배치가 커질수록 "취객의 걸음걸이"가 덜 비틀거리게 되며, 이는 오히려 중간 껍질에서의 탈출을 늦춘다는 것을 시사합니다.
3. 정규화 ($\lambda$): 이것은 "자기력"의 강도입니다. 논문은 더 강한 자기력이 컴퓨터가 중간 껍질을 더 빨리 탈출하도록 돕지만, 특정 지점까지만 그렇다는 것을 밝혀냈습니다.

그들은 이 노브들이 탈출 시간에 어떻게 영향을 미치는지에 대한 구체적인 공식(스케일링 법칙)을 도출했습니다. 예를 들어, 암기에서 일반화로 넘어가는 데 걸리는 시간은 대략 $1/(\eta \lambda)$에 비례합니다. 이는 학습률이나 정규화를 두 배로 높이면 지연 시간이 절반으로 줄어든다는 것을 의미합니다.

제외된 가설들

이 논문은 자신들이 주장하는 것이 무엇이 아닌지 매우 신중하게 밝히고 있습니다. 저자들은 네트워크가 갑자기 "회로를 찾아냈다"거나 생물학적인 각성과 같은 신비로운 깨달음을 얻었다고 말하는 것이 아닙니다. 대신, 그로킹은 순수하게 솔루션 공간의 기하학적 구조와 최적화 과정의 무작위성의 결과라고 주장합니다. 또한 이것이 단순한 우연이 아니라는 점도 분명히 합니다. 그들의 수학은 이것이 Adam이 이러한 특정 유형의 문제를 최적화할 때 나타나는 예측 가능하고 구조적인 특징임을 보여줍니다.

얼마나 확신하는가?

저자들은 자신의 이론에 상당히 자신감을 가지고 있지만, 증명한 것과 측정한 것을 구분하는 데 신중합니다.

• 수학: 그들은 연속 시간 모델을 기반으로 한 엄밀한 수학적 증명을 통해 이러한 스케일링 법칙을 도출했습니다. 또한 심볼릭 대수 시스템(수식을 계산하는 컴퓨터 프로그램)을 사용하여 공식의 정확성을 검증함으로써 방정식에 대한 높은 신뢰도를 확보했습니다.
• 증명: 그들은 두 가지 특정 유형의 퍼즐, 즉 군론 학습(구체적으로는 대칭군 $S_5$ 상의 학습)과 모듈로 산술(구체적으로는 $Z_{127}$ 상의 정수)을 통해 아이디어를 테스트했습니다.
• 결과: 실험에서 컴퓨터의 행동은 그들의 예측과 일치했습니다. 예를 들어, 학습률이나 배치 크기를 변경했을 때 그로킹에 걸리는 시간이 그들의 공식이 예측한 대로 변했습니다. 또한 그들의 이론이 예측한 대로 암기 반경의 "U자형" 곡선도 관찰되었습니다.

하지만 저자들은 자신의 수학적 모델이 작은 학습률과 큰 배치 크기와 같은 특정 조건에 의존한다는 점을 언급합니다. 그들은 이 이론이 모든 가능한 신경망에서 발생하는 모든 종류의 그로킹을 설명한다고 주장하는 것이 아니라, "껍질-핵심(shell-core)" 기하학이 존재하는 이 특정 구조적 작업들에서 발생하는 현상을 설명한다고 말합니다.

결론

그로킹은 마법이 아니라 기하학적인 여정입니다. 컴퓨터는 바깥쪽에서 시작하여 "암기 껍질"에 갇혔다가, 학습 과정의 무작위한 노이즈가 자신을 "일반화 핵심"으로 밀어 넣을 때까지 배회해야 합니다. 이 점프에 걸리는 시간은 얼마나 빨리 움직이는지(학습률), 걸음걸이가 얼마나 일정한지(배치 크기), 그리고 얼마나 강력하게 중심부로 끌려가는지(정규화)에 달려 있습니다. 저자들은 수학으로 이 규칙들을 그려내고 실험으로 확인함으로써, 왜 이러한 신경망들이 때때로 긴 잠을 자는 것처럼 보이다가 갑자기 깨어나는지에 대한 명확한 그림을 제시했습니다.