기술 요약: 동역학계로서의 최적화: 잠재 ODE를 이용한 생성적 스케줄링

문제 정의

딥러닝은 고차원 비볼록(non-convex) 최적화에 크게 의존합니다. 업데이트 규칙은 단순하지만, 최적의 학습률 스케줄($\eta_t$)을 결정하는 것은 여전히 중요한 과제입니다. 전통적인 방식은 고정된 파라미터 스케줄(예: 상수, 코사인 감쇠, 스텝 감쇠, OneCycle)에 의존하거나 이러한 형태에 대한 하이퍼파라미터 탐색을 수행합니다. 그러나 이러한 방법들은 특정 학습 역학에 적응하지 못하는 경우가 많으며, 특정 모델-데이터-태스크 조합에 있어 무엇이 최적의 스케줄을 구성하는지에 대한 근본적인 이해가 부족합니다. 또한, 강화 학습(RL)이나 하이퍼그래디언트 디센트(hypergradient descent)와 같은 고급 방법들은 높은 계산 비용을 초래하거나, 네트워크 파라미터 정보가 필요하여 확장성이 제한되거나, 미래의 성능 궤적을 고려하지 못한다는 한계가 있습니다.

방법론

저자들은 훈련 과정을 하나의 동역학계(dynamical system)로 취급하는 메타 학습 접근 방식을 제안합니다. 이 방법의 핵심은 과거의 하이퍼파라미터 탐색 실행으로부터 학습의 잠재적 표현을 학습하는 잠재 상미분 방정식(Latent Ordinary Differential Equation, LODE) 모델입니다.

1. LODE 모델

이 시스템은 실험 중에 흔히 기록되는 시계열 지표인 훈련 손실($\ell$), 검증 성능 지표($\nu$), 그리고 학습률($\eta$)을 학습하도록 훈련됩니다.

• 인코딩(Encoding): RNN 인코더는 최근의 훈련 시계열 세그먼트 $x(t) = (\ell, \nu, \eta)$를 잠재 벡터 $z_0$로 매핑합니다. 이 벡터는 네트워크의 내부 파라미터에 접근할 필요 없이 결합된 훈련 상태를 나타내며, 이를 통해 이 방법은 임의의 모델 크기로 확장될 수 있습니다.
• 역학(Dynamics): 잠재 상태는 신경 ODE $f_\theta(z, t)$의 적분을 통해 시간의 흐름에 따라 전개됩니다. 경로 길이에 대한 $\ell_2$ 페널티는 시간 불변적 표현을 장려하여, 유사한 훈련 상태가 발생 시점과 관계없이 유사한 잠재 영역에 매핑되도록 보장합니다.
• 디코딩(Decoding): 디코더는 잠재 상태로부터 예상되는 미래의 손실, 검증 성능 및 학습률 궤적을 재구성합니다.

2. 동적 학습률 스케줄러

스케줄러는 새로운 "테스트" 모델을 훈련하는 동안 다음 $\mu$ 단계에 대한 최적의 학습률을 생성하기 위해 작동합니다. 과정은 다섯 단계로 구성됩니다:

1. 인코딩: 테스트 모델의 현재 훈련 이력을 잠재 벡터 $z_0$로 인코딩합니다.
2. 섭동(Perturbation): $z_0$에 가우시안 노이즈를 추가하여 잠재 벡터의 앙상블을 생성합니다.
3. 샘플 수락: 각 잠재 벡터는 앞으로 적분되고 디코딩됩니다. 샘플은 예측된 손실 궤적이 현재 테스트 모델의 손실과 특정 통계적 거리(2 표준 편차) 내에 머무는 경우에만 수락되어 관련성을 보장합니다.
4. 목표 조건화(Goal Conditioning): 수락된 샘플들에 대해, 시스템은 예측된 학습률 스케줄을 바탕으로 미래 시점 $\Delta t$에서의 검증 성능을 추정합니다.
5. 앙상블 예측: 최종 학습률은 가장 우수한 성능을 보이는 앙상블 멤버들(예측된 미래 검증 정확도가 가장 높은 멤버들)의 스케줄을 평균하여 계산됩니다.

주요 기여

• 적응형 스케줄러: 훈련 손실, 검증 성능, 학습률 궤적을 예측하는 잠재 ODE 모델 기반의 학습률 스케줄러를 도입했습니다.
• 생성 능력: 단순히 고정된 파라미터 스케줄 세트에서 선택하거나 작은 곱셈 보정을 적용하는 방법과 달리, 이 방법은 장기적 성능을 최적화하기 위해 관찰된 훈련 데이터에서 크게 벗어나는 특화된 비파라미터 스케줄을 생성합니다.
• 확장성 및 효율성: 이 방법은 옵티마이저에 독립적이며 네트워크 파라미터 정보를 요구하지 않으므로 대규모 모델로의 확장이 용이합니다. 또한 RL 기반 스케줄러나 하이퍼그래디언트 디센트보다 계산 효율성이 높습니다.
• 평탄한 최솟값(Flat Minima)을 통한 일반화: 이 방법은 탐색 단계 동안 학습률을 전통적인 "안정성의 경계(edge of stability)" 너머로 일관되게 밀어붙여, 최적화 도구가 손실 지형의 더 평탄한 최솟값에 안착할 수 있게 함으로써 더 나은 일반화 성능을 유도합니다.

실험 결과

저자들은 LODE 스케le듈을 네 가지 비적응형 베이스라인(constant, cosine-OneCycle, cosine-decay, step-decay), 하이퍼그래디언트 디센트, RL 기반 컨트롤러, 그리고 스케줄 프리(schedule-free) 최적화와 비교 평가했습니다.

• 데이터셋 및 모델: Fashion-MNIST, CIFAR-100, ImageNet(CNN, ResNet18, ResNet34 사용) 및 단편 소설 코퍼스를 사용한 트랜스포머 모델의 다음 토큰 예측 태스크에서 실험을 수행했습니다.
• 성능: LODE 스케줄러는 모든 모델 아키텍처와 데이터셋에서 모든 베이스라인을 능가했습니다. 모든 랜덤 시드에 대해 일관되게 낮은 분산을 유지하며 가장 높은 테스트 정확도를 달-성했습니다.
  • CIFAR-100 (ResNet18)에서 LODE는 OneCycle의 74.0%, RL-컨트롤러의 73.9% 대비 74.9%의 정확도를 달성했습니다.
  • ImageNet (ResNet34)에서 LODE는 RL-컨트롤러의 73.9% 대비 74.5%를 달성했습니다.
  • 트랜스포머 태스크에서 LODE는 59.8%의 다음 토ک 예측 정확도를 기록하여, 스케줄 프리 방법(56.3%) 및 기타 방법들을 앞질렀습니다.
• 손실 지형 분석: 헤시안(Hessian)의 최대 고윳값($\lambda_{max}$) 분석 결과, LODE 스케줄러로 훈련된 모델은 훨씬 더 평탄한 최솟값(예: CIFAR-100에서 LODE의 $\lambda_{max} \approx 0.62$ vs OneCycle의 $1.22$)으로 수렴함을 확인했습니다. 이는 우수한 일반화 특성을 나타냅니다.
• 계산 비용: LODE 스케줄러의 추론 비용은 단순 파라미터 스케줄보다 약 25% 높지만, 하이퍼그래디언트 디센트나 RL 스케줄러보다는 현저히 낮습니다. 또한 RL 방법과 달리 모델 크기에 따라 비용이 증가하지 않습니다.

의의 및 주장

본 논문은 훈련을 동역학계로 모델링함으로써 LODE 스케줄러가 훈련의 뚜렷한 단계(그래디언트 수집, 탐색, 수렴, 과적합)를 인식하고 그에 따라 학습률을 조정할 수 있다고 주장합니다.

• 파라미터 형태를 넘어서: 이 방법은 최적의 스케줄이 반드시 매끄러운 파라미터 함수일 필요는 없음을 보여줍니다. 즉, 전통적인 스케줄러가 놓치는 급격한 전환이나 비단조적(non-monotonic) 동작을 포함할 수 있습니다.
• 안정성의 경계: 스케줄러는 탐색 단계 동안 전통적인 "안정성의 경계" 제한을 초과하는 학습률을 효과적으로 활용합니다. 이를 통해 손실 지형을 가로질러 더 크고 평탄한 분지(basin)를 찾은 뒤 수렴 단계로 진입하게 하며, 이러한 동작은 더 나은 장기적 결과를 가져옵다.
• 실용성: 이 접근 방식은 기존의 실험 추적 플랫폼(예: Weights & Biases, MLFlow) 위에 레이어로 구현할 수 있는 실용적인 도구로 제시됩니다. 이를 통해 메타 러너를 위한 광범한 새로운 훈련 실행 없이도, 기존의 하이퍼파라미터 탐색 데이터를 활용하여 뛰어난 스케줄을 생성함으로써 훈련 과정을 효율화할 수 있습니다.

저자들은 결론적으로, 자신들의 방법이 복잡한 모델(특히 이점이 가장 두드러지는 트랜스포머와 같은 경우)에 대해 계산 효율적이고 수치적으로 안정적이며 효과적인 방식으로 비볼록 손실 지형을 탐색할 수 있는 방법을 제공한다고 밝힙니다.