Training Deep Learning Models with Norm-Constrained LMOs

In this work, we study optimization methods that leverage the linear minimization oracle (LMO) over a norm-ball. We propose a new stochastic family of algorithms that uses the LMO to adapt to the geometry of the problem and, perhaps surprisingly, show that they can be applied to unconstrained problems. The resulting update rule unifies several existing optimization methods under a single framework. Furthermore, we propose an explicit choice of norm for deep architectures, which, as a side benefit, leads to the transferability of hyperparameters across model sizes. Experimentally, we demonstrate significant speedups on nanoGPT training using our algorithm, Scion, without any reliance on Adam. The proposed method is memory-efficient, requiring only one set of model weights and one set of gradients, which can be stored in half-precision. The code is available at https://github.com/LIONS-EPFL/scion .

한 줄 요약

Scion이라는 새로운 LMO 기반 최적화 알고리즘을 제안하여 nanoGPT 학습 속도를 향상시키고 메모리 효율성을 확보한다.

핵심 기여도

LMO를 사용한 새로운 확률적 알고리즘 가족을 제안하여 문제의 기하학적 구조에 적응.
기존 최적화 방법을 하나의 프레임워크로 통합.
딥러닝 아키텍처에 대한 명시적인 노름 선택을 제안, 하이퍼파라미터 전이 가능.
nanoGPT 학습에서 Adam 없이 1.6× 가속 성능 달성.

핵심 아이디어

기존 최적화 알고리즘은 학습 중 기하 구조를 동적으로 조정하지만, 이는 네트워크 구조에 대한 사전 정보를 무시한다. 본 연구는 LMO(Linear Minimization Oracle)를 사용하여 사전에 문제의 기하 구조에 맞게 최적화기를 설계하는 새로운 접근법을 제안한다. 특히, LMO는 노름-볼 내에서 선형 최소화를 수행하며, 이는 문제의 스케일에 관계없이 방향만 고려하는 특성을 가진다. 이 아이디어는 기존 Conditional Gradient 방법과 유사하지만, 벡터 ℓ∞-노름과 스펙트럼 노름과 같은 비표준 노름에서도 적용 가능하다는 점에서 차별화된다. 또한, Newton-Schultz 반복을 활용한 효율적인 LMO 계산을 제안하여, 실제 딥러닝 학습에 적용 가능성을 높인다.

기술적 접근법

**LMO 기반 알고리즘**: 문제의 노름-볼 내에서 선형 최소화를 수행하는 LMO를 핵심 연산으로 사용.
**확률적 알고리즘**: LMO를 사용하여 기하 구조에 적응하는 확률적 알고리즘 가족 제안.
**하이퍼파라미터 전이**: 딥러닝 아키텍처에 대한 명시적인 노름 선택을 통해 모델 크기와 관계없이 하이퍼파라미터가 전이 가능.
**반복 알고리즘**: Newton-Schultz 반복을 사용하여 스펙트럼 노름 기반 LMO를 효율적으로 계산.
**메모리 효율**: 모델 가중치와 그라디언트를 각각 하나의 세트만 저장하여, 반정밀도(half-precision)로 저장 가능.

주요 결과

**nanoGPT 학습 가속**: Scion 알고리즘을 사용한 nanoGPT 학습에서 Adam 없이 1.6× 가속 성능 달성.
**메모리 효율성**: 기존 최적화 알고리즘 대비 메모리 사용량 감소, 반정밀도 저장 가능.
**하이퍼파라미터 전이**: 제안된 노름 선택을 통해 모델 크기와 관계없이 하이퍼파라미터가 전이 가능.

의의 및 한계

**의의**: LMO 기반 최적화는 기존 Adaptive Optimizer(예: Adam)와 달리 네트워크 구조에 대한 사전 정보를 활용하여 학습 효율성을 높인다. 특히, Newton-Schultz 반복을 활용한 LMO 계산은 실제 딥러닝 학습에 실용적이다.
**한계**: 본 연구는 주로 스펙트럼 노름과 ℓ∞-노름을 기반으로 한 LMO를 다루며, 다른 노름 구조에 대한 일반화 가능성은 명시되지 않음. 또한, 실험은 nanoGPT에만 제한되어 있어 다른 모델에 대한 확장성은 추가 연구가 필요.

실용적 활용

Scion 알고리즘은 Adam 없이도 빠른 학습 속도와 메모리 효율성을 제공하므로, 특히 메모리 제한이 있는 장치(예: 모바일, 임베디드)에서 딥러닝 모델 학습에 유용하게 활용될 수 있다. 또한, 하이퍼파라미터 전이 가능성은 다양한 모델 크기에서 동일한 최적화 전략을 적용할 수 있도록 지원한다.