Fast Generalization after Interpolation via Critically Damped Momentum Optimization

기술 요약: 보간 후 임계 감쇠 모멘텀 최적화를 통한 빠른 일반화

문제 정의
유전학, 의학, 금융 분야에서 흔히 발생하는 고차원 저표본(high-dimensional, low-sample) 환경에서 머신러닝 모델은 종종 완벽에 가까운 훈련 성능(보간, interpolation)을 달로하면서도 미지의 데이터에는 실패하곤 한다. 이러한 격차는 손실 함수가 동일한 훈련 손실을 산출하지만 복잡도와 테스트 성능 면에서는 크게 다른 수많은 파라미터 구성들을 포함하는 최적화 지형(optimization landscape) 때문에 발생한다. 가중치 감쇠(weight decay)와 같은 규제 방법들이 전통적으로 최적화를 더 단순한 솔루션으로 유도하기 위해 사용되어 왔지만, "그로킹(grokking)" 현상은 시간적 단절을 드러낸다. 즉, 모델은 데이터를 먼저 암기(memorize)한 후 훨씬 나중에야 일반화 단계로 전이된다. 이러한 지연된 전이는 계산적으로 비효율적이며, 일반화를 위해 모델이 암기하는 데 필요한 에포크(epoch) 수가 훨씬 더 많이 소요된다. 핵심 문제는 모델이 보간 임계값(interpolation threshold)에 도달한 직후, 느린 암시적 역학(implicit dynamics)에 의존하는 대신 어떻게 효율적으로 저복잡도 보간 솔루션을 즉각적으로 선택할 것인가 하는 점이다.

방법론
저자들은 적합(fitting) 단계와 일반화 단계를 명시적으로 분리하도록 설계된 이상(biphasic) 최적화 전략인 GROKtimizer를 제안한다.

1. 이론적 프레임워크:

   • 보간 후 역학(Post-Interpolation Dynamics): 본 논문은 보간 후의 레짐(regime)을 국소 이차 동역학계(local quadratic dynamical system)로 모델링한다. 보간 솔 \ชัน 근처의 손실 지형이 헤시안 행렬($H$)과 $L_2$ 규제 항($\lambda$)을 포함하는 이차 형식(quadratic form)으로 근사될 수 있다는 가정하에, 평탄한 방향(flat directions, $H$의 영공간)을 따르는 역학은 감쇠 진동자(damped oscillator)와 유사하게 동작한다.
   • 임계 감쇠(Critical Damping): 저자들은 **임계 감쇠 모멘텀(Critically Damped Momentum, CDM)**이 이 평탄한 분지(basin) 내에서 파라미터 노름(norm)을 최소화하는 최적의 궤적을 제공함을 입증한다. 모멘텀 파라미터 $\beta$를 $\beta = 1 - 2\sqrt{\lambda\eta}$ (여기서 $\eta$는 학습률)로 조정함으로써, 옵티마이저는 임계 감쇠를 달성한다.
   • 수렴 속도 향상: 이론적으로 CDM은 이 레짐에서 고전적인 경사 하강법(Gradient Descent, GD)에 비해 이차적 속도 향상(quadratic speedup)을 제공한다. GD는 $\lambda^{-1}$에 비례하는 시간 척도로 수렴하는 반면, CDM은 $\lambda^{-1/2}$에 비례하는 시간 척도로 수렴한다. 이는 최소 노름 솔루션에 도달하는 데 필요한 에포크 수를 이론적으로 이차적으로 감소시킨다.
   • 학습률 추정: 이를 실제 구현하기 위해, 저자들은 전체 헤시안을 계산하지 않고도 헤시안의 최대 고윳값($\lambda_{max}$)을 추정할 수 있는 아이켄(Aitken) 보정 파워 반복법(Power Iteration method)을 제안한다. 이를 통해 안정성의 경계에 있는 최적의 학습률 $\eta_{max}$를 온라인으로 계산할 수 있다.

2. 알고리즘 구현 (GROKtimizer):

   • 단계 1 (보간): 옵티마이저는 보간 임계값에 빠르게 도달하기 위해 가중치 감쇠 없이(zero weight decay) 실행된다. 이 단계에서는 빠른 데이터 피팅을 위해 표준 모멘텀을 활용한다.
   • 단계 2 (일반화): 보간이 달성되면, 옵티마이저는 가중치 감쇠가 적용된 임계 감쇠 모멘텀으로 전환된다. 모멘텀 파라미터는 이론에서 도출된 임계 감쇠 조건에 맞춰 조정되며, 가중치 감쇠가 적용되어 파라미터를 보간 매니폴드 내의 최소 노름 솔루션으로 유도한다.

주요 기여

1. 보간 후 레짐의 특성화: 저자들은 지연된 일반화 단계를 손실 지형의 헤시안 영공간(null space)을 따르는 감쇠 진동자 문제로 정식화하여, 그로킹 현상을 헤시안의 기하학적 구조와 연결했다.
2. 이론적 속도 향상: CDM이 표준 경사 하강법보다 저노름 보간 솔루션으로의 수렴을 이차적으로 더 빠르게 가속화함을 증명하여, 국소 이차 모델 하에서 1차 메서드(first-order methods) 중 최적성을 확립했다.
3. GROKtimizer: 보간 후의 노름 최소화 단계로 동적으로 전환되는 이상(bipathic) 최적화 스케줄을 도입했다.
4. 실증적 검증: 합성 그로킹 벤치마크 및 실제 고차원 데이터셋을 통해 검증하였으며, 표준 옵티마이저(Adam, AdamW, SGD, Muon) 대비 일관된 개선을 보여주었다.

실험 결과
논문은 다음과 같은 벤치마크에서 GROKtimizer를 평가한다:

• 합성 벤치마크: Modular Addition, Binary Addition, Sparse Parity와 같은 작업에서 GROKtimizer는 베이스라인을 크게 상회한다. 예를 들어, Gaussian 작업의 경우 가장 강력한 베이스라인에 비해 검증 손실을 여러 차수(orders of magnitude) 낮추었다.
• 실제 고차원 데이터: 생물 의학 분류 작업(Leukemia, TCGA) 및 분자 특성 예측(QM9)에서 GROKtimizer는 가장 높은 정확도와 가장 낮은 검증 손실을 달 기록했다. 이러한 이득은 특징(feature)의 수가 샘플 수보다 훨씬 많은 데이터 희소 환경에서 가장 두드러졌다.
• 샘플 효율성: 저데이터 MNIST 실험(클래스당 10~50개 샘플)에서 GROKtimizer는 가장 큰 이점을 보였으며, 이는 명시적인 노름 최소화가 결정 불가능한(underdetermined) 솔루션 공간에서 샘플 효율성을 개선함을 확인시켜 준다.
• 언어 모델 사전 학습: WikiText2 및 BabyLM 모델 실험에 따르면, 트랜스포머의 경우(스케일 대칭 및 적응형 업데이트로 인해) 파라미터 노름과 검증 손실 사이의 직접적인 상관관계가 덜 명확함에도 불구하고, 이상(biphasic) 패턴은 여전히 유효하다. 보간 후 노름 최소화 단계로 전환하는 것은 검증 손실을 추가로 감소시키는 결과를 낳았다.

의의 및 주장
본 논문은 포스트-인터폴레이션(post-interpolation) 훈련을 수동적인 대기 기간이 아닌 명시적인 최적화 문제로 다룸으로써 "그로킹" 문제에 대한 원칙적인 해결책을 제공한다고 주장한다. 손실 지형의 기하학을 활용함으로써, 이 방법은 저복잡도 솔루션의 선택을 가속화한다.

저자들은 자신들의 연구 결과를 **평탄한 최솟값 가설(flat-minima hypothesis)**과 결합하며, 헤시안의 영공간 차원이 평탄도와 복잡도의 견고한 척도로 기능한다고 주장한다. 이들은 이 영공간 내에서 노름을 최소화하는 것이 더 평탄하고 일반화 가능한 최솟값을 향해 항해하는 것과 동일하다고 제안한다.

한계점
저자들은 다음과 같은 한계점을 인정한다:

• 이 접근 방식은 주로 데이터가 제한된 환경을 위해 설계되었으며, 데이터가 풍부한 대규모 비전 또는 언어 사전 학습에서는 중요도가 낮을 수 있다.
• 이 방법은 모델이 보간에 도달할 수 있는 충분한 용량(capacity)을 갖추고 있다고 가정하며, 과소 매개변수화(under-parameterized)된 모델을 위한 해결책은 아니다.
• 이론적 분석은 손실 지형의 국소 이차 근사에 의존하지만, 실증적 결과는 이 가정이 편차를 보이더라도 방법론이 견고함을 시사한다.
• 파라미터 노름은 복잡도의 대리 지표(proxy)로 사용되었으며, 이는 모든 아키텍처에서 모델 일반화의 모든 측면을 포착하지 못할 수도 있다.