기술 요약: 그로킹(Grokking)을 길들이는 법: 제어 신호로서의 표현 기하학(Representation Geometry)

문제 정의

그로킹은 과잉 매개변수화된 신경망에서 발생하는 현상으로, 모델이 초기에 훈련 데이터를 암기한 후 장기간의 최적화를 거친 뒤에야 강력한 일반화 성능을 보이는 현상을 말한다. 최근의 연구들은 가중치 감쇠(weight decay), 특징 압축(feature compression), 회로 형성(circuit formation), 그리고 표현 학습 역학(representation learning dynamics)을 그로킹과 연결 지었으나, 지연된 일반화의 발생 시점과 출현을 결정짓는 구체적인 요인들은 여전히 완전히 이해되지 않은 상태이다. 기존 문헌의 공통된 주제는 학습된 표현이 점진적으로 저차원 구조로 압축된다는 것이다. 그러나 이전의 연구들은 표현 기하학을 직접적인 제어 메커니즘이라기보다 창발적인 속성으로 주로 취급해 왔다. 본 논문은 은닉 표현의 기하학적 구조를 명시적으로 조작하는 것이 그로킹의 시작과 속도에 영향을 미치는 제어 신호로 기능할 수 있는지 조사한다.

방법론

저자들은 은닉 표현의 유효 차원을 훈련 중에 직접적으로 제어하도록 설계된 표현 수준의 스펙트럼 정규화 도구인 **기하학적 차원 정규화(Geometric Dimensionality Regularization, GeomDR)**를 도입한다.

기하학적 특성 규명

은닉층 표현 행렬 $Z \in \mathbb{R}^{N \times d}$에 대해, 방법론은 먼저 데이터를 중심화하여 $\tilde{Z}$를 얻는다. 그 다음 공분산 행렬 $C = \frac{1}{N-1}\tilde{Z}^\top\tilde{Z}$를 계산한다. 유효 차원($D_{eff}$)은 공분산 행렬 $C$의 고윳값($\mu_i$)에 대한 참여 비율(participation ratio)을 사용하여 다음과 같이 계산된다:
$$D_{eff} = \frac{(\sum_i \mu_i)^2}{\sum_i \mu_i^2}$$
또한, 국소 이웃 거리($\rho$)는 정규화된 표현들의 $k$-최근접 이웃까지의 평균 유클리드 거리로 정의된다.

GeomDR 목적 함수

GeomDR는 목표 차원 $d^*$를 초과하는 방향에 포함된 분산을 억제하는 페널티를 부과한다. 내림차순으로 정렬된 고윳값($\mu_1 \ge \mu_2 \ge \dots \ge \mu_d$)이 주어졌을 때, 층별 손실 함수는 다음과 같다:
$$L_{GeomDR}^{(\ell)} = \sum_{j=d^*+1}^{d} \mu_j^{(\ell)}$$
전체 훈련 목적 함수는 태스크 손실($L_{task}$)과 기하학적 정규화 항의 결합이다:
$$L = L_{task} + \lambda(t)L_{GeomDR}$$
결정적으로, 이 정규화 항은 처음부터 활성화되지 않는다. 이는 초기 무제약 훈련 단계(활성화 단계 $t_s$) 이후에 도입되며, 코사인 램프(cosine ramp)를 사용하여 강도 $\lambda(t)$가 점진적으로 증가한다. 이러한 설계는 암기가 일반화에 앞서 발생하는 자연스러운 그로킹 궤적을 모사한다.

실험 프로토콜

본 프레임워크는 그로킹 현상을 보이는 것으로 알려진 세 가지 알고리즘 태스크에 대해 평가되었다:

1. 모듈러 덧셈(Modular Addition): $(a + b) \mod 97$ 예측.
2. 모듈러 나눗셈(Modular Division): $(a \cdot b^{-1}) \mod 97$ 예측.
3. 치환 합성(Permutation Composition): $S_5$에서의 두 치환의 합성 예측.

실험에는 MLP와 트랜스포머(Transformer) 두 가지 아키텍처가 사용되었다. 연구는 개입 스케줄(강도 $\lambda_{max}$ 및 램프 기간 $T_{ramp}$), 목표 차원($d^*$), 그리고 개입 시작 시점($t_s$)을 체계적으로 변화시켰다.

주요 기여

1. 선행되는 붕괴 관찰: 저자들은 베이스라인 설정에서 차원 붕괴(dimensionality collapse)가 암기에서 일반화로의 전이보다 일관되게 앞서 발생함을 입증한다. 이는 기하학적 압축이 단순한 상관관계가 아니라 전이를 이끄는 기계적 동인임을 시사한다.
2. GeomDR 도입: 타겟 부분 공간 외부의 공분산 스펙트럼을 직접 억제함으로써 은닉 표현의 유효 차원을 제어하는 스펙트럼 정규화 도구인 GeomDR을 제안한다.
3. 역학의 체계적 제어: 기하학적 개입이 지연된 일반화 역학을 실질적으로 변화시킬 수 있음을 보여주는 포괄적인 경험적 연구를 제공한다. 개입 스케줄과 목표 차원에 따라 GeomDR은 그로킹을 가속화하거나, 지연시키거나, 혹은 완전히 억제할 수 있다.

결과

그로킹의 가속화

평가된 모든 설정에서 GeomDR은 그로킹 역학을 일관되게 변화시켰다:

• MLP 성능: 모듈러 덧셈에서 GeomDR은 그로킹을 최대 52배 가속화했다(단계 수를 약 362K에서 7K로 단축). 모듈러 나눗셈(39x)과 치환 합성(23x)에서도 유사한 개선이 관찰되었다. 모든 MLP 사례에서 GeomDR은 베이스라인이 불안정하거나 실패했던 지점에서 100% 성공률을 달 Ach성했다.
• 트랜스포머 성능: 트랜스포머의 베이스라인 그로킹 역학은 이미 MLP보다 빨랐지만, 그럼에도 불구하고 GeomDR은 일관된 가속(예: 모듈러 덧셈에서 ~2x)을 제공하며 100% 성공률을 유지했다.

하이퍼파라미터의 영향

• 목표 차원 ($d^*$): 중간에서 큰 규모의 목표 차원($16 \le d^* \le 64$)이 가장 강력한 가속을 보였다. 공격적인 압축($d^* \le 8$)은 종종 학습을 늦추거나 불안정하게 만들었다.
• 개입 시점 ($t_s$): "임계 시간 창(critical temporal window)"이 확인되었다. 유용한 표현이 형성되기 전(너무 이른 시점)이나 암기가 고착된 후(너무 늦은 시점)에 적용된 개입은 효과가 떨어졌다. 중간 단계의 활성화 시점(예: MLP의 경우 $t_s = 2000$)이 가장 빠르고 안정적인 그로킹을 생성했다.
• 스케줄: 중간 정도의 정규화 강도와 램프 기간이 일반적으로 가장 좋은 성능을 보였다. 지나치게 강한 개입은 훈련을 불안정하게 만들 수 있다.

기하학적 상관관계

연구는 그로킹을 가속화하는 개입이 차원 붕괴의 발생 또한 가속화함을 확인했다. GeomDR 모델에서는 베이스라인과 비교했을 때 차원 붕애와 성공적인 일반화 사이의 시간적 간격이 유의미하게 줄어들었으며, 이는 기하학적 재구성과 일반화의 출현 사이의 긴밀한 결합을 시사한다.

확장성 및 효율성

저자들은 공분산 스펙트럼을 표현의 무작위 서브셋(예: 데이터의 5%25%)으로부터 추정하더라도 GeomDR이 효과적으로 작동함을 입증하였으며, 이는 그로킹 가속에 영향을 주지 않으면서 실행 시간을 2.4배2.7배 단축시킨다.

의의 및 주장

본 논문은 표현 기하학이 그로킹을 위한 효과적인 제어 신호 역할을 한다고 주장한다. 은닉 활성화의 공분산 스펙트럼을 직접 수정함으로써, 연구자들은 네트워크가 언제, 얼마나 빠르게 암만에서 일반화로 전환될지를 체계적으로 조절할 수 있다.

저자들은 이 작업이 기하학적 개입이 지연된 일반화에 대한 실용적인 접근법임을 보여주는 증거를 제공한다고 상정한다. 표현의 변화를 수동적으로 관찰하는 대신, 이 방법론은 기하학적 속성을 직접 조작하여 학습 역학에 대한 가설을 테스트할 수 있게 한다. 결과는 표현 기하학과 그로킹 사이의 관계가 다양한 태스크와 아키텍처 전반에 걸쳐 견고하다는 것을 시사하지만, 그 개선의 폭은 변동이 있다(트랜스포터의 경우 더 강력한 귀납적 편향을 가지고 있기 때문에 MLP보다 개선 폭이 적게 나타남).

결론적으로, 본 논문은 표현 차원을 유일한 인과적 메커니즘으로 확립하지는 않지만, 기하학적 개입이 그로킹의 타이밍을 변화시키기에 충분하다는 것을 입증함으로써 최적화 및 일반화 연구를 위한 새로운 실험 방법론을 제시한다.