Weight Decay Regimes in Grokking Transformers: Cheap Online Diagnostics

기술 요약: 그로킹 트랜스포머의 가중치 감쇠 체제

문제 제기
그로킹 (Grokking) 은 알고리즘적 작업 (특히 모듈러 산술) 으로 훈련된 트랜스포머가 훈련 손실이 포화 상태에 도달한 지 오래된 시점에서 암기에서 일반화로 급격히 전환되는 현상을 의미합니다. 기계적 연구는 내부 회로 (예: 푸리에 표현) 와 인덕션 헤드 (induction heads) 의 형성이 이러한 전환의 동인임을 규명했으나, 기존 문헌은 이러한 역학을 실시간으로 추적할 수 있는 '저렴한 온라인 진단 도구'가 부재합니다. 또한, 가중치 감쇠 ($\lambda$) 가 이러한 체제를 제어하는 매개변수로서 모델 규모 전반에 걸쳐 체계적으로 매핑되지 않았으며, 주의력 헤드 (attention-head) 의 조정이 가중치 크기 효과와 분리되어 특정 기여도를 규명하지 못했습니다.

방법론
저자들은 AdamW 옵티마이저를 사용하여 모듈러 산술 작업 (mod+, mod-, mod$\times$, mod$\div$) 에 디코더 전용 트랜스포머를 훈련합니다. 본 연구는 가중치 감쇠 값에 대한 밀집한 스윕 (dense sweep) 을 수행하며, 모델 규모 (0.82M 에서 85M 파라미터까지) 와 아키텍처 (트랜스포머, MLP, LSTM, Mamba) 를 변화시킵니다.

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

1. 온라인 순서 매개변수 (Online Order Parameters): 두 가지 스칼라 진단 지표가 모든 훈련 단계에서 주의력 활성화 (attention activations) 만을 사용하여 계산됩니다:
   • 평균 쌍별 주의력 헤드 코사인 유사도 ($\bar{s}(t)$): 주의력 헤드 간의 조정 (동기화) 정도를 측정합니다.
   • 엔트로피 표준 편차 ($\sigma_H(t)$): 헤드 간 샤논 엔트로피의 분산을 측정하여 분화 (differentiation) 추이를 추적합니다.
2. 유한 크기 스케일링 (Finite-Size Scaling): 저자들은 가중치 감쇠에 대한 그로킹 확률 ($P(\text{grok})$) 에 로지스틱 함수를 적합시켜 임계값 $\lambda_c$를 추정합니다. 또한 '그로킹까지의 시간' ($t_{\text{grok}}$) 에 멱법칙 (power law) 을 적합시켜 임계 지수 $\nu$를 결정합니다.
3. 인과적 개입 (Causal Interventions): 헤드 패턴 구조와 가중치 크기를 구분하기 위해, 저자들은 주의력 헤드를 재초기화하는 실험과 분화 단계의 정점에서 가중치 노름을 클리핑하는 실험을 짝을 지어 수행합니다.
4. 공식 검증 (Formal Verification): 진단 식별자 (경계 및 대수적 속성) 는 수학적 잘 정의됨을 보장하기 위해 Lean 4 에서 기계적으로 검증됩니다.

주요 기여

1. 정량적 임계값: 본 연구는 mod+ 작업에 대한 표준 4 층 8 헤드 트랜스포머를 위한 날카로운 가중치 감쇠 임계값을 규명합니다: $\lambda_c = 0.0158$ (95% CI [0.0109, 0.0200]).
2. 체제 다이어그램: 두 축 ($\lambda$, $N$) 지도는 세 가지 뚜렷한 체제를 보여줍니다:
   • 암기 (Memorization): $\lambda < \lambda_c$ (거의 제로에 가까운 그로킹).
   • 발달적 그로킹 (Developmental Grokking): $\lambda \ge \lambda_c$ (약 100% 그로킹률 도달, $\lambda$가 증가함에 따라 그로킹까지의 시간 감소).
   • 붕괴 (Collapse): $\lambda = 10$ (모든 헤드가 동일한 패턴으로 수렴, $\bar{s}=1.0$).
3. 이중 단계 역학: 훈련 궤적은 다음과 같은 두 단계로 특징지어집니다:
   • 1 단계 (동기화): 헤드가 수렴하며 ($\bar{s}$가 약 0.93 에서 0.995 로 상승) 그로킹 시작과 일치합니다.
   • 2 단계 (분화): 헤드가 발산하며 ($\bar{s}$가 약 0.88 로 하락, $\sigma_H$ 증가) 테스트 정확도는 높은 상태를 유지합니다.
4. 헤드당 차원 (Per-Head Dimension) 을 진폭 조절기로: 분화의 진폭 ($\sigma_H$) 은 경험적으로 헤드당 차원 ($d/H$) 에 의해 지배되며, $d/H$가 증가함에 따라 단조적으로 포화됩니다.
5. 교차 아키텍처 범위: 가중치 감쇠로 제어되는 전환은 주의력이 없는 아키텍처 (MLP, LSTM, Mamba) 에서도 관찰되지만, 임계값 $\lambda_c$는 아키텍처에 따라 크게 이동합니다 (약 한 자릿수 범위).

결과

• 스케일링 지수: 그로킹까지의 시간은 멱법칙 $t_{\text{grok}} \propto (\lambda - \lambda_c)^{-\nu}$를 따르며, 경험적 지수는 $\nu = 0.757$ (95% CI [0.725, 0.799]) 입니다. 저자들은 현재 그리드 밀도 하에서 이 값이 표준 참조 지수 (예: $\nu=1/2$ 또는 3D 이징 $\nu \approx 0.63$) 와 일치하지 않는다고 명시적으로 언급하며, 보편성 클래스 식별은 향후 연구에 미루고 있습니다.
• 장기적 취약성: 확장된 시간 범위 (20,000 에포크) 에서 일부 궤적에서 시드 의존적인 '반그로킹 (anti-grokking)' 붕괴가 관찰되며, 이는 정확도 감퇴로 끝나는 5 단계 패턴으로 특징지어집니다. 이는 보편적 주기라기보다 취약성으로 식별됩니다.
• 인과적 격리: 주의력 헤드를 재초기화하면 대조군에 비해 2 단계 진폭 ($\sigma_H$) 이 유의하게 감소하는 반면, 가중치 노름 클리핑은 그렇지 않습니다. 이는 효과를 가중치 크기가 아닌 헤드 패턴 구조로 격리시킵니다.
• 다중 작업 복제: 가중치 감쇠 제어 패턴은 네 가지 모듈러 연산과 두 가지 모델 규모 전반에 걸쳐 유지되어, 덧셈을 넘어선 체제 지도의 견고성을 확인합니다.

의의 및 주장
본 논문은 정제된 지역 학습 계수 (rLLC) 와 같은 더 비싼 손실 지형 진단을 보완하는, 주의력 활성화만을 사용하여 그로킹 역학을 추적할 수 있는 저렴하고 온라인 계산 가능한 진단 프레임워크를 제공한다고 주장합니다. 이는 암기, 발달적 그로킹, 붕괴 체제를 분리하는 주요 경험적 제어 매개변수로서 가중치 감쇠를 확립합니다.

저자들은 그들의 범위에 대해 겸손합니다:

• 관찰된 지수 $\nu \approx 0.76$에 대해 보편성 클래스를 주장하지 않으며, 테스트된 참조 지수들이 신뢰 구간 밖이라는 점을 지적합니다.
• 가중치 감쇠가 모든 아키텍처에서 일반화에 필수적이라고 주장하지 않으며, 오직 표준 주의력 아키텍처 내에서 제어 매개변수로 작용한다고만 주장합니다.
• 분화가 발생하는 이유에 대한 인과적 메커니즘 주장은 명시적으로 후속 연구에 미루고, 대신 언제 그리고 어떻게 나타나는지에 초점을 맞춥니다.
• 발견 사항은 소형 트랜스포머 주의력 모델의 모듈러 산술로 제한되며, 대형 언어 모델이나 더 넓은 보편성에 대한 주장은 범위를 벗어납니다.

본 연구는 유한 크기 스케일링, 리아푸노프 수축, 유리 완화 (glass relaxation) 에 대한 동시 이론적 연구를 보완하는 '체제 지도'와 '순서 매개변수'를 제공하는 현상론적 및 진단적 기여로 자신을 위치시킵니다.