🤖 machine learning

Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers

이 논문은 Muon으로 훈련된 트랜스포머가 표현-판독 인터페이스(representation-readout interface)에서의 불안정성으로 인해 모듈로 산술을 학습한 후 일반화가 실패하는 "포스트 그로킹 붕괴(post-grokking collapse)"를 보인다는 점을 입증하며, 이는 특정 모델 구성 요소를 동결하거나 과업에 정렬된 회로를 재조정함으로써 해결될 수 있는 발산하는 옵티마이저 역학 및 마스킹 효과로 특징지어지는 현상이다.

원저자: Ali Janati, Kaoutar El Maghraoui, Andrei Kanavalau, Anass Belfatmi

게시일 2026-08-10

📖 5 분 읽기🧠 심층 분석

CC BY 4.0

원저자: Ali Janati, Kaoutar El Maghraoui, Andrei Kanavalau, Anass Belfatmi

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 로봇에게 아주 구체적인 수학 퍼즐을 가르치고 있다고 상상해 보세요. 이 퍼즐은 두 숫자를 더한 뒤 큰 소수로 나눈 나머지를 찾는 것입니다. 처음에는 로봇이 마치 단어장을 통째로 외우는 학생처럼 답을 단순히 암기하는 것처럼 보입니다. 하지만 그러다 갑엔 '그로킹(grokking)'이라 불리는 마법 같은 일이 일어납니다. 긴 정체기 끝에, 로봇은 단순히 암기하는 것을 넘어 실제로 패턴을 '이해'하게 됩니다. 이제 로봇은 한 번도 본 적 없는 숫자들에 대해서도 완벽하게 문제를 풀 수 있습니다. 이는 과학자들에게 매우 중요한 일인데, 인공지능이 단순히 암기하는 것이 아니라 어떻게 '생각하는 법'을 배우는지 이해하는 데 도움을 주기 때문입니다.

하지만 여기 함정이 있습니다. 때때로 로봇이 해답을 찾아낸 후에도, 여전히 같은 퍼즐로 훈련을 받고 있음에도 불구하고 다시 그 답을 잊어버리기 시작합니다. 이는 마치 대수학을 드디어 이해한 학생이, 방금 풀었던 똑같은 문제에서 다시 실수를 연발하는 것과 같습니다. 과학자들은 로봇의 내면을 들여다보기 위해 특별한 도구들을 사용합니다. 그들은 로봇이 문제를 풀기 위해 사용하는 특정 패턴(마치 노래의 음표와 같은 것)을 관찰하여, 로봇이 여전히 올바른 음표를 사용하고 있는지 확인합니다. 여기서 핵심적인 질문은 이것입니다. 로봇은 왜 잊어버리는가? 로봇이 음표를 배우는 것을 멈췄기 때문인가, 아니면 다른 무언가가 일어나고 있는 것인가?

사라진 천재의 미스터리

이 연구에서 연구자들은 모듈로 덧셈(modular addition) 퍼즐을 풀도록 훈련된 '트랜스포머(Transformer)'라는 특별한 유형의 AI 로봇을 관찰했습니다. 그들은 로봇을 가르치기 위해 두 명의 서로 다른 '선생님(최적화 도구)'을 사용했습니다. AdamW라는 이름의 선생님은 표준적이고 신뢰할 수 있는 강사입니다. 또 다른 이름인 Muon은 더 새롭고 빠른 선생님으로, 로봇이 퍼즐을 훨씬 더 빠르게 이해하도록 돕는 것처럼 보입니다.

연구 결과, Muon은 확실히 속도광였습니다. Muon은 AdamW가 걸리는 시간의 절반 정도밖에 걸리지 않고도 로봇이 '그로킹'의 순간, 즉 깨달음의 순간에 도달하도록 도왔습니다. 하지만 반전이 있습니다. Muon이 찾아낸 해법은 믿을 수 없을 정도로 취약합니다. 로봇이 퍼즐을 해결하고 나면, 정확도가 100%에서 거의 0%로 떨어졌다가 다시 회복되는 과정을 반복하며 자주 실패하기 시작합니다. 이는 로봇이 여전히 같은 문제를 학습하고 있음에도 불구하고 계속해서 일어납니다. AdamW 역시 완전히 안전한 것은 아니지만, 훨씬 덜 자주, 그리고 훨씬 덜 심각하게 실패합니다.

깨진 악수 (The Broken Handshake)

그렇다면 무엇이 잘못된 것일까요? 연구자들은 문제가 로봇이 수학을 잊어버린 것이 아님을 발견했습니다. 로봇은 여전히 자신의 '두뇌' 속에 정답을 완벽하게 간직하고 있습니다. 문제는 로봇의 두 부분 사이의 깨진 악수입니다. 즉, '생각하는 부분(표현, representation)'과 '답을 내놓는 부분(출력 헤드, readout)' 사이의 문제입니다.

로봇의 뇌를 책이 가득한 도서관(생각)이라고 하고, 로봇의 입을 그 책을 소리 내어 읽는 사서(답)라고 상상해 보세요.

**도서관(은닉층, hidden layers)**은 매 초마다 책에 새로운 메모를 적어 넣습니다.
**사서(출력 헤드, output head)**는 그 책들을 읽고 답을 내놓으려 노력합니다.

Muon 선생님과 함께라면, 도서관은 업데이트될 때마다 책에 조금씩 다른 언어로 새로운 메모를 적어 내려가는데, 사서는 계속해서 예전 언어로 그 책들을 읽으려고 합니다. 처음에는 둘이 서로를 이해하며 잘 해나갑니다. 하지만 도서관이 사서가 적응하는 속도보다 더 빠르게 언어를 바꾸기 때문에, 결국 둘은 서로를 이해하지 못하게 됩니다. 도서관은 여전히 올바른 수학을 적고 있지만, 사서는 엉뚱한 단어를 읽고 있는 것입니다.

연구자들은 한쪽을 고정함으로써 이를 증명했습니다. 도서관이 언어를 바꾸는 것을 멈추자 로봇은 완벽한 상태를 유지했습니다. 사서가 읽는 스타일을 바꾸는 것을 멈추었을 때도 로봇은 완벽했습니다. 하지만 두 부분이 모두 자유롭게 움직이도록 허용하자, 둘은 서로 어긋났고 로봇은 실패하기 시작했습니다.

기계 속의 유령 (The Ghost in the Machine)

가장 혼란스러운 부분은 이것입니다. 로봇이 실패하고 있을 때조차, '도서관'에는 여전히 완벽한 해답이 들어있다는 점입니다. 연구자들은 수학을 알고 있는 부분만을 골라내기 위해 특별한 필터를 사용했습니다. 그렇게 했을 때, 해답은 여전히 100% 정확했습니다!

결국 로봇은 회로 마스킹(circuit masking) 현상을 겪고 있는 것입니다. 도서관이 올바른 답을 외치고 있지만, 로봇의 나머지 부분이 정확히 같은 볼륨으로 틀린 답을 외치고 있다고 상상해 보세요. 사서는 엉망진창인 소리를 듣고 잘못된 답을 내놓게 됩니다. 정답은 여전히 그곳에 있지만, 소음에 묻혀버린 것입니다.

연구자들은 다른 것은 바꾸지 않고 오직 정답이 있는 부분(수학 패밀리)의 볼륨만 높였을 때, 로봇이 즉시 다시 100% 정확도로 돌아간다는 것을 보여주었습니다. 수학이 사라진 것이 아니라, 단지 소음 속에 길을 잃었을 뿐이었습니다.

기존 도구들이 놓친 이유

보통 과학자들은 로봇이 문제를 푸는 데 사용하는 '음표'를 살펴봄으로써 로봇이 제대로 작동하는지 확인합니다. 그들은 올바른 음표가 존재하는지, 그리고 그 음표가 충분히 큰지를 체크합니다. 이 연구에서 그런 도구들은 로봇이 괜찮다고 말했습니다. 음표는 존재했고, 이전과 거의 똑같이 들렸습니다.

하지만 로봇은 그럼에도 불구하고 실패하고 있었습니다. 이는 도구들이 '음표'만 보고 있었을 뿐, 사서가 사용하는 '볼륨'이나 '언어'를 보지 못했기 때문입니다. 음표는 완벽했지만, 사서가 그것을 제대로 듣지 못하고 있었던 것입니다. 이는 라디오 방송이 선명하게 송출되고 있는지 확인하면서도, 정작 라디오가 엉뚱한 주파수에 맞춰져 있다는 사실은 알아채지 못하는 것과 같습니다.

심층 분석 (The Deep Dive)

연구자들은 또한 로봇을 더 '깊게'(더 많은 사고 층을 추가) 만들었을 때 어떤 일이 일 발생하는지도 테스트했습니다. 그들은 Muon이 학습 속도를 더욱 높인다는 것을 발견했지만, '깨진 악수'의 문제는 더 악화되었습니다. 더 깊은 모델에서는 로봇의 뇌가 한 섹션에서 수학을 작성하면, 답을 읽는 부분은 다른 섹션에 위치하게 되어 둘이 더 빠르게 어긋나게 됩니다.

또한 연구자들은 Muon 선생님의 특수한 '정규화(normalization)' 기술(업데이트를 안정적으로 유지하는 방법)을 제거하여 실험해 보았습니다. 이 기술을 제거하자 로봇은 수학을 매우 좁고 집중된 방식으로 학습했지만, 결국 완전히 무너져 내려 회복하지 못했습니다. 이는 이 특별한 기술이 로봇이 빠르게 학습하고 지식을 넓게 퍼뜨리도록 돕기도 하지만, 동시에 이러한 유형의 '드리프트(drift) 실패'에 더 취약하게 만든다는 것을 시사합니다.

핵심 요점

이 논문의 주요 교훈은 로봇의 각 부분이 동기화되지 않으면 빠른 학습은 위험할 수 있다는 것입니다. 로봇은 수학을 잊어버린 것이 아니라, 수학을 아는 부분과 그것을 말하는 부분 사이의 연결 고리를 잃어버린 것입니다.

연구자들은 학습이 끝난 후 '말하는' 부분(출력 및 임베딩)을 고정하면 로봇이 안정적으로 유지된다는 것을 발견했습니다. 이는 불안정성이 로봇이 수학을 못 해서 생기는 것이 아니라, 두 부분이 서로 다른 속도로 움직이며 공유하는 언어를 잃어버리기 때문에 발생한다는 것을 시사합니다.

요약하자면, 로봇은 천재이지만 갑자기 말을 더듬기 시작한 상태입니다. 왜냐서 뇌와 입이 서로 다른 방언을 사용하고 있기 때문입니다. 수학은 내면에 완벽하게 살아있지만, 두 방언을 일치시킬 방법이 없다면 그 천재는 자신의 능력을 제대로 보여줄 수 없습니다.

기술 요약: Muon으로 훈련된 트랜스포머의 표현-판독 인터페이스(Representation–Readout Interface)에서의 그로킹 후 붕괴

문제 정의
본 논문은 모듈로 산술(특히 모듈로 덧셈 및 뺄셈)에 대해 훈련된 트랜스포머 모델에서, "그로킹(grokking)"—훈련 정확도가 포화된 훨씬 이후에 일반화가 나타나는 현상—이 일어난 후 발생하는 특정 불안정성을 조사한다. 선행 연구는 Muon(직교화 옵티마이저)이 AdamW보다 그로킹을 가속화한다는 것을 입증했지만, 본 연구는 Muon으로 훈련된 모델이 그로킹 이후에 빈번하게 재발하는 급격한 일반화 성능 붕괴를 겪는다는 점을 보고한다. 결정적으로, 이러한 실패는 Muon에만 국한되지 않으며 AdamW 모델 또한 불안정성을 보인다. 본 논문은 두 옵티마이저 모두 5개의 시드(seed)에 대해 엄격하게 안정적이지 않음을 입증한다. 선택된 AdamW 베이스라인은 5개 중 4개의 시드에서 임계값 미만의 평가 결과(최저 27.59% 도달)를 기록한 반면, Muon은 5개 시드 모두에서 그러한 결과를 보였다. 핵심 문제는 모델이 훈련 세트를 여전히 해결하고 있음에도 불구하고 일반화 능력을 상실한다는 것이며, 표준적인 기계론적 진행 척도(스펙트럼 지원 및 전력 분포)는 실패를 감지하지 못하고, 정확도가 거의 무작위 수준(near-chance)으로 떨어지는 동안에도 회로가 온전하다고 보고한다.

방법론
저자들은 은닉 가중치 행렬(hidden weight matrices)은 Muon에 의해 업데이트하고, 토큰/위치 임베딩과 출력 언임베딩(readout)은 AdamW에 의해 업데이트하는 분할 옵티마이저(split-optimizer) 훈련 체계를 채택하였다. 연구에는 다양한 깊이(1, 2, 4)의 디코더 전용 트랜스포머를 사용하였으며, 모듈로 산술 작업( $a+b \pmod p$ 및 $a-b \pmod p$ , 여기서 $p=113$ 및 변형 $p=97$ )을 활용하였다.

주요 방법론적 구성 요소는 다음과 같다:

매칭 브랜칭(Matched Branching): 인과적 요인을 격리하기 위해 비트 단위로 동일한 파라미터 상태로부터 궤적을 분기한다. 한 브랜치는 표준 최적화를 계속하는 반면, 다른 브랜치들은 특정 파라미터 그룹(은닉 행렬, 임베딩 또는 판독층)을 동결하여 필요성을 테스트한다.
푸리에 분해(Fourier Decomposition): 저자들은 모델의 내부 잔차 상태(residual states)를 피연산자 그리드(operand grid)에 대한 2D 푸리에 모드로 분해한다. 이들은 특정 "패밀리(families)"의 모드(예: 덧셈 패밀리 $(k, k)$ , 뺄셈 패밀리 $(k, -k)$ )를 정의하고, 필터링 개입을 사용하여 이러한 패밀리들이 작업 수행에 미치는 충분성과 필요성을 측정한다.
개입 스위트(Intervention Suite): 본 연구는 학습된 알고리즘의 정확한 성격과 불일치의 근원을 파악하기 위해 위상 스크램블링(phase scrambling), 주파수 재배치(frequency relocation), 교차 판독 치환(cross-readout substitution)과 같은 개입을 적용한다.
절제 실험(Ablation): 관찰된 스펙트럼 분산과 불안정성의 역할을 테스트하기 위해 Muon의 정규화 직교 경로(normalized-orthogonalization path)를 절제한다.

주요 결과 및 발견

붕괴의 본질: 실패는 "은닉-판독 불일치(hidden–readout misalignment)"이다. 훈련 손실이 0에 가까워짐에 따라 그래디언트는 소멸하지만, 두 옵티마이저는 다르게 반응한다. Muon의 업데이트 스텝 사이즈는 그래디언트 크기와 거의 무관해지는 반면(탄성 $\approx -0.03$ ), AdamW는 이에 비례한다(탄성 $\approx +1.5$ ). 이로 인해 은닉 표현과 판독층은 파라미터당 8.0배의 속도로 서로 멀어진다. 훈련 세트는 여전히 해결된 상태이지만, 판독층은 자신이 함께 훈련되었던 표현을 더 이상 디코딩할 수 없게 된다.
마스킹 vs. 회로 실패: 논문은 두 가지 실패 모드를 구분한다. *회로 실패(circuit failure)*에서는 푸리에 패밀리 자체가 작업을 수행하는 것을 멈춘다. *회로 마스킹(circuit masking)*에서는 작업에 정렬된 패밀리(예: 덧셈 패밀리)가 고립된 상태에서는 작업을 완벽하게 수행(정확도 100%)하지만, 전체 모델 내에서 적대적인 나머지 부분(adversarial remainder)에 의해 압도된다. 저자들은 작업에 정렬된 성분을 재조정하는 것만으로도 전체 모델의 정확도를 복구할 수 있음을 보여줌으로써, 회로는 온전했으나 "투표에서 밀려났음(outvoted)"을 증명한다.
옵티마이저 차이: Muon은 AdamW보다 훨씬 빠르게 그로킹에 도달하지만(중앙값 13,011 스텝 vs. 30,486 스텝), 모든 테스트된 구성에서 보편적으로 불안정하다. AdamW 역시 5개 시드에 대해 불안정하지만, 붕계의 심각도(임계값 미만 평가 횟수로 측정)는 Muon보다 두 자릿수 낮다. 단일 시드 스윕(single-seed sweep)에서 AdamW 구성은 학습률에 따라 심각도가 변하는 조건부 불안정성을 보였으나, 5개 시드에 대한 반복 실험을 통해 선택된 "안정적" 베이스라인이 시드 간에 엄격하게 안정적이지는 않음을 확인했다. Muon의 불안정성은 표현을 넓은 스펙트럼에 분산시키는 업데이트 규칙(효과적인 공액 쌍 326개 vs. AdamW의 약 5개)에 기인하며, 이로 인해 작업 정렬 패밀리가 차지하는 총 전력의 비중이 작아진다.
패밀리의 작업 선택: 옵티마이저가 아닌 작업(task)이 푸리에 패밀리를 선택한다. 덧셈으로 훈련된 모델은 $(k, k)$ 모드를 사용하고, 뺄셈으로 훈련된 모델은 $(k, -k)$ 모드를 사용한다. 패밀리 내의 기저(basis)는 정확하며 공유되지 않는다. 동일한 시드에서 시작된 두 모델은 서로 다른 주파수와 위상을 사용하도록 발산하며, 이들의 판독층은 서로 거의 직교하여 상대방의 표현을 디코딩하는 데 실패한다.
깊이와 정렬: 더 깊은 모델에서 푸리에 코드는 마지막 전 블록의 MLP에 의해 작성되지만, 최종 블록에서만 언임베딩(unembedding)에 의해 "해독 가능(legible)"해진다. 불일치는 최종 블록의 잔차와 언임베딩 사이의 인터페이스에서 구체적으로 발생한다.
예방: 회로가 형성된 후 은닉 행렬이나 비은닉 파라미터(임베딩 및 판독층)를 동결하면 붕괴를 완전히 방지할 수 있다. 그로킹 후 임베딩과 판독층을 동결하는 "Stable Muon"은 451,400 스텝의 포스트 그로킹 단계와 5개 시드에 대해 엄격한 안정성을 달성한 반면, 표준 Muon은 반복적으로 실패했다.

의의 및 주장
본 논문은 분할 옵티마이저 훈련 및 그로킹 후 역학의 근본적인 한계를 식별했다고 주장한다: 잔차 스트림은 특권적인 기저(privileged basis)를 가지지 않으며, 손실 표면이 평탄해질 때 손실은 표현과 판독 사이의 특정 정렬을 유지하기 위한 압력을 제공하지 않는다. 저자들은 다음과 같이 주장한다:

표준 척도는 불충분하다: 스펙트럼 척도(지원 및 전력 분포)는 붕괴 중에도 불변할 수 있는데, 이는 실패가 알고리즘 자체의 상실이 아니라 기저의 변화(불일치)이기 때문이다.
그로킹은 가역적 마스킹이다: 일반화에서 붕계로의 전이는 암기에서 일반화로의 전이의 역과 같으며, 두 과정 모두 작업 정렬 패밀리의 표현 전력 비중이 적대적 나머지 부분에 비해 상승하거나 하락하는 과정을 포함한다.
인터페이스가 실패의 지점이다: 불안정성은 은닉 표현이나 판독층 단독의 속성이 아니라, 그들의 결합된 식별(joint identification)의 문제이다. 인터페이스의 한쪽을 고정하면 시스템이 안정화된다.
Muon의 속도는 대가를 수반한다: Muon은 회로의 발견을 가속화하지만, 그 업데이트 규칙은 표준 훈련 조건 하에서 판독과의 불일치에 더 취약한 표현을 생성하므로 장기적 안정성을 위해 개입(동결)이 필요하다.

이 연구는 새로운 응용을 제안하기보다는 기계론적 해석 가능성(mechanistic interpretability)의 특정 실패 모드에 대한 기계론적 설명을 제공하며, 모델이 올바른 회로를 보유하고 있더라도 표현 드리프트(representational drift)로 인해 이를 활용하지 못할 수 있음을 강조한다.