조심스러운 옵티마이저: 한 줄 코드로 학습 효율을 높이다

초록

본 논문은 기존 모멘텀 기반 옵티마이저(Ada​mW, Lion 등)에 “업데이트 방향과 현재 그래디언트가 일치할 때만 적용”하는 한 줄 코드를 추가하는 방법을 제안한다. 이 마스크 연산은 손실 감소를 보장하면서 원래 옵티마이저의 수렴 이론을 그대로 유지한다. 실험 결과, 대규모 언어 모델 사전학습 및 이미지 분류에서 동일한 하이퍼파라미터 설정으로도 학습 속도와 최종 정확도가 향상됨을 보여준다.

상세 분석

논문은 모멘텀 기반 옵티마이저가 업데이트 방향 uₜ와 실제 그래디언트 gₜ 사이에 부호 불일치가 발생하면 손실이 일시적으로 증가하고 수렴이 느려지는 현상을 지적한다. 이를 해결하기 위해 저자들은 요소별 마스크 ϕ(uₜ ∘ gₜ) = α(uₜ ∘ gₜ)·I(uₜ ∘ gₜ > 0) 를 도입한다. 여기서 α 은 “양의 스케일링 팩터”로, 마스크에 의해 감소된 업데이트 크기를 보정한다. 핵심 아이디어는 부호가 일치하지 않는 차원에 대해 업데이트를 0으로 만들고, 일치하는 차원에서는 원래 업데이트에 비례하여 적용함으로써 uₜ·gₜ ≥ 0 를 강제하는 것이다.

이러한 변형은 연속시간 해밀턴-디센트 프레임워크에 자연스럽게 매핑된다. 기존의 해밀턴 시스템은 H(w,s)=L(w)+K(s) 라는 라플라스 함수가 비증가하도록 설계되지만, L(w) 자체는 순간적으로 증가할 수 있다. 논문은 수정된 동역학 (5) 에서 xₜ = ∇L(wₜ) ∘ ∇K(sₜ) 를 정의하고, ϕ 가 xᵢ>0이면 ϕᵢ≥1, xᵢ<0이면 ϕᵢ≤0 을 만족하도록 하면 dH/dt ≤ ‑ΔH, dL/dt ≤ ‑ΔL 가 동시에 성립함을 증명한다. 즉, 원래 시스템보다 더 빠르게 전체 에너지와 실제 손실이 감소한다.

이론적 보장은 두 단계에서 제공된다. 첫째, 연속시간 분석에서 라플라스 함수와 손실 모두가 단조 감소함을 보이며, 두 번째로 이산시간 분석에서는 마스크가 적용된 업데이트가 원래 업데이트보다 손실 감소량이 크거나 최소한 손실을 증가시키지 않는 조건을 제시한다(정리 2.3, 2.4). 특히, μ‑smooth 손실에 대해 단계 크기 ε 가 특정 상한 이하일 때 마스크 적용이 보장된다.

실험 부분에서는 2‑D 이차형 손실 L(w)=κw₁²+w₂² (κ=4) 에 대해 GD, GDM, C‑GDM을 비교한다. C‑GDM은 진동과 오버슈팅을 크게 줄이며, 라플라스와 손실 모두를 더 빠르게 감소시킨다. 대규모 실험에서는 LLM 사전학습(예: 7B 파라미터 모델)과 이미지 분류(CIFAR‑100, ImageNet)에서 C‑AdamW와 C‑Lion이 각각 +2~3% 의 토큰 효율 향상 및 +0.5% 정도의 정확도 상승을 기록한다. 중요한 점은 기존 하이퍼파라미터를 그대로 사용했으며, 추가 튜닝이 거의 필요 없다는 것이다.

비판적으로 보면, 마스크 연산이 α 값에 크게 의존한다. 논문은 α = dim/nnz + 1 을 기본값으로 제시하지만, 실제 고차원 모델에서는 nnz 비율이 학습 단계마다 크게 변동할 수 있다. 또한, 마스크가 완전한 0 업데이트를 초래하는 경우(특히 uₜ 와 gₜ 가 거의 반대인 경우) 학습이 일시적으로 정지할 위험이 있다. 저자들은 모멘텀 누적이 결국 정렬을 회복한다는 논증을 제시하지만, 급격히 변하는 손실 지형에서는 수렴 속도가 오히려 느려질 가능성도 존재한다.

마지막으로, “한 줄 코드”라는 주장에도 불구하고 실제 구현에서는 ϕ 의 스케일링과 ε 의 재조정이 필요하며, 이는 코드 라인 수와는 별개로 실험 설계에 추가적인 복잡성을 도입한다. 그럼에도 불구하고, 이 접근법은 기존 옵티마이저에 대한 최소 침습적 개선을 제공한다는 점에서 실무 적용 가능성이 높다.