배치 크기와 학습률의 토큰 예산 최적화: 확률적 조건부 그래디언트 방법의 새로운 이론

초록

본 논문은 μ‑KL 조건 하에서 모멘텀 기반 확률적 조건부 그래디언트(Scion 등) 알고리즘의 수렴성을 분석하고, 배치 크기 B, 시퀀스 길이 S, 학습률 β가 고정 토큰 예산 T에 미치는 영향을 정량화한다. BS = B·S에 대한 세 가지 스케일링 구간(노이즈‑지배, 중간, 대배치)과 최적 BST 스케일링 규칙 BS ≈ T^{2/3}을 제시하며, 이를 기반으로 토큰‑예산 인식 적응형 배치·시퀀스·학습률 스케줄을 설계한다. NanoGPT 실험을 통해 이론적 예측이 실제 학습 곡선과 일치함을 입증한다.

상세 분석

이 연구는 대규모 언어 모델 학습이 “토큰 예산 T”라는 제약 하에 진행된다는 사실에 착안한다. 토큰 예산은 업데이트 횟수 K, 배치 크기 B, 시퀀스 길이 S와 관계식 T = K·B·S 로 연결되며, 따라서 (B, S, β) 삼중항이 최적화 효율을 좌우한다. 기존 문헌에서는 배치 크기가 일정 수준을 넘어가면 일반화가 악화된다는 경험적 사실만을 제시했지만, 본 논문은 이를 μ‑KL(μ‑Kurdyka‑Łojasiewicz) 조건 하에서 정량화한다. μ‑KL 조건은 ‖∇f(x)‖_* ≥ μ·(f(x)‑f*) 형태로, 최적점과의 함수값 차이를 1차 미분량에 직접 연결한다. 이는 조건부 그래디언트(Linear Minimization Oracle, LMO) 기반 최적화에 자연스럽게 맞물리며, 특히 제한된 도메인에서 ζ‑QC와 동등하게 해석될 수 있다.

알고리즘 1은 모멘텀 파라미터 α와 스텝 사이즈 β, 학습률 η를 포함한다. 저자들은 (A1)‑(A4) 가정 하에, 특히 잡음 분산 σ²가 배치·시퀀스 곱 BS에 비례한다는 가정(σ² = σ_*²·BS)을 도입해 기대값 수렴 분석을 전개한다. 정리 4.1은 β = Θ(1/K), η = Θ(1/μ), α = min(1, Θ(ε·μ/(ρσ)²)) 등 파라미터 선택을 명시하고, K가 아래와 같은 복합 상한을 만족하면 E