반응형 옵티마이저 (Reactive Optimizers)
요약 (TL;DR) 우리는 에이전트를 옵티마이저의 핵심에 배치했습니다. 이 에이전트는 학습 동역학(training dynamics)을 관찰하고, 이에 반응하며, 학습이 진행됨에 따라 옵티마이저의 소스 코드를 재작성합니다. 특정 작업에 대해 아직 최적화 레시피가 조정되지 않은 경우에도, 이미지 및 언어 작업을 아우르는 상용 제품에서 선택할 수 있는 최고의 레시피와 맞먹거나 능가하는 성능을 보여줍니다. 나아가, 이 시스템이 손실 급증(loss spikes)을 방지하고 과적합(overfitting)을 탈출하기 위해 자발적으로 정교한 스케줄과 복잡한 전략을 고안해낸다는 것을 입증했습니다. OPAL 패키지를 사용하여 직접 학습 실행에 적용해 보세요.
오늘날 딥 신경망(deep neural network)을 최적화하는 데 사용되는 레시피는 보통 학습이 시작되기도 전에 완전히 결정됩니다. 학습 과정이 어떻게 진행되어야 하는지에 대한 방대한 인간의 지혜가 투입되었습니다: 적응형 스텝 크기(adaptive step sizes) [1], 사전 조건화(preconditioning) [2, 3], 직교화된 업데이트(orthogonalized updates) [4, 5], 학습 단계에 연결된 스케줄(schedules tied to the training step) [6, 7], 모델 크기에 따라 조정되는 학습률(learning rates scaled with model size) [8] 등이 있습니다. 자동 발견(Automated discovery) [9, 10, 11, 12]은 더 나은 옵티마이저를 찾기 위해 수많은 실행을 검색합니다. 이 옵티마이저는 수작업으로 제작되었든 자동으로 발견되었든 관계없이, 사전에 작성된 규칙에 따라 변화하는지 여부와 관계없이 학습 끝까지 변함없이 작동합니다. 이러한 범위는 다음과 같은 질문들에 대한 답을 남겨둡니다: 옵티마이저는 예상치 못한 손실 급증이나 정체(stall)에 어떻게 반응해야 할까요? 만약 한 모듈의 그래디언트가 진동하기 시작하는 동안 다른 모듈들은 부드럽게 개선된다면 어떨까요? 그리고 새로운 유형의 아키텍처가 Adam이나 Muon으로는 제대로 학습되지 않는다면, 그 아키텍처 자체가 나쁜 것일까요, 아니면 아직 자신에게 맞는 옵티마이저를 만나지 못한 것일까요?
본 게시물에서는 기존 옵티마이저의 범위를 확장하여, 모델 자체의 학습 과정에 반응하도록 만들고, 이를 통해 단일 학습 실행(single training run) 내에서 끊임없이 적응하는 옵티마이저를 구현하고자 합니다. 이 반응형 옵티마이저는 현장에서 추론하며 각 파라미터 그룹에 대한 가중치 업데이트 규칙을 작성합니다. 최적화 코드를 직접 작성함으로써, 복잡한 스케줄을 구현하고, 학습 중 발생하는 사고(training accidents)에 대응하며, 존재하는 단일 학습 에피소드에 맞춰 옵티마이저를 커스터마이징할 자유도를 갖게 됩니다. 이전에 학습 중 옵티마이저를 적응시키려는 시도[12, 13, 14]와 모델 스스로 파인튜닝을 지시하게 하는 것[15]에 기반하여 발전된 반응형 옵티마이저는 완전히 제로샷(zero-shot)이며, 반복 실행에서 학습하거나 유지하는 것이 전혀 없고, 옵티마이저의 전체 알고리즘을 재작성할 수 있도록 합니다.
우리는 이미지 생성 및 언어 모델링과 같은 실제 상황에서 반응형 옵티마이저가 정적(static)인 옵티마이저보다 훨씬 더 잘 수행할 수 있는 사례부터 시작합니다. 이후에는 정교한 스케줄링, 손실 급증 억제(suppressing loss spikes), 과적합 탈출(escaping overfitting)을 포함하여 반응형 옵티마이저의 창발적 특성(emergent properties)으로 넘어갑니다. 마지막으로, 이 모든 것이 얼마나 많은 비용이 들고 대규모로 얼마나 실용적인지 질문합니다. 우리는 반응형 옵티마이저가 에이전트 세대(generation of agents)가 거듭될수록 더욱 강력해질 것으로 예상하며, 아직 레시피가 존재하지 않는 곳(새로운 아키텍처나 새로운 데이터셋의 경우)에서 가장 유용할 것이며, 딥러닝 최적화의 퍼즐을 푸는 한 걸음이 될 것이라고 기대합니다.
1. 레시피 없는 학습
한 박사 과정 학생이 자신만의 설계 모델이나 트랜스포머를 학습시킬 새로운 데이터셋이라는 아이디어를 가지고 있다고 가정해 봅시다. 그에게는 단 하루 동안 사용할 수 있는 하나의 GPU와 소규모의 LLM API 토큰만 주어져 있어, 진지하게 수행할 수 있는 학습 실행은 오직 한 번뿐입니다. 학생은 무엇을 해야 할까요? 가장 당연한 행동 방침은 가장 표준적인 학습 레시피나 최근 논문에서 제시된 레시피를 선택하고, 그것이 자신의 시나리오에 적합하지 않을 가능성(예: 손실 스파이크/정체 등)을 감수하는 것일 겁니다.
근본적인 질문은 더 광범위합니다. GPU 시간, 연구원 시간 또는 벽시계 시간(wall-clock time)으로 제한될 때, 최적의 최적화 레시피를 사전에 알지 못한 채 어떻게 딥러닝 모델을 최고의 성능으로 학습시킬 수 있을까요? 새로운 아키텍처, 모달리티, 언어 및 과학 데이터 모두 이 설정에 속하며, 소수의 사람만이 대규모로 탐색(sweep at scale)할 여유가 있습니다.
반응형 옵티마이저(reactive optimizer)는 그 예산(budget)을 사용할 또 다른 방법을 제공합니다. 바로 GPU/인간 시간을 대신하여 LLM 토큰을 사용하는 것입니다. 최적화 도구가 학습 과정의 문제에 반응하고 스스로 조정함으로써, 레시피 튜닝의 부담 대부분을 줄일 수 있습니다.
간략한 결과(Results in brief)
우리는 모디파이드-nanogpt [16]의 최적화 트랙(Track 3)에서 시작합니다. 이는 언어 모델 사전 학습 작업(FineWeb으로 162M GPT를 학습시키는 것)이며, 이 단일 설정에 맞춰 많은 레시피가 조정되었습니다. 이 작업을 통해 우리는 두 가지 약간 다른 자연스러운 설정을 만듭니다: (1) 대체 언어 데이터셋(중국어, 태국어, 키르기스어를 선택함)으로 훈련된 동일한 모델, 그리고 (2) 배치 크기의 4분의 1로 줄인 원래 설정입니다. 이 두 가지 모두 해당 설정에 맞춰 조정된 레시피가 없으므로, 우리는 반응형 옵티마이저를 새로운 학습을 한 번에(one-shot) 처리하려는 경우 가질 수 있는 자연스러운 선택지들과 비교합니다. 즉, 5개의 리더보드 레시피를 자연스럽게 조정하여 채택하는 방식입니다. 이 모든 새로운 설정에서, 반응형 옵티마이저는 최고의 기준선(baseline)을 능가하거나 일치합니다 (Figure 1).
우리는 또한 반응형 옵티마이저(Reactive Optimizer)를 Goldfish [17]에 직접 적용했습니다. Goldfish는 350개 언어의 데이터셋을 제안하고, 각 데이터셋마다 표준 레시피로 작은 GPT 모델을 학습시키지만, 이 방식은 특정 언어에 맞춰 조정되지는 않았습니다. 따라서 이는 옵티마이저의 적응성을 원샷(one-shot)으로 테스트하기에 이상적인 기반 환경입니다. Amharic 데이터셋에서 공식 레시피는 검증용(held-out) FLORES-200 텍스트 기준으로 문장당 83.48 비트에서 끝났습니다 (낮을수록 좋음). 이와 동일한 레시피로 초기화된 반응형 옵티마이저는 원샷으로 81.11에 도달했습니다 (스와힐리어의 경우 유사한 성능 향상을 얻었습니다). 표준 레시피의 학습률(learning rate)을 탐색하여, 정적 레시피가 반응형 옵티마이저보다 앞선 80.89를 기록하는 것을 발견했지만, 이 탐색된 레시피로 초기화된 반응형 옵티마이저는 다시 앞서 80.73에 도달했습니다.
또한 우리는 이미지 생성 분야에서도 반응형 옵티마이저를 시도해 보았습니다. 이 분야는 문헌상 확산 모델 설계(diffusion-model design)와 옵티마이저 설계가 주로 직교하는(orthogonal) 것으로 간주되기 때문에, 일반적으로 최적화 하이퍼파라미터(optimization hyperparameters) 탐색을 수행하지 않습니다. 반응형 옵티마이저는 AdamW로 학습하는 공식 레시피를 가진 LightningDiT-XL/1 [18]에서 앞섰으며 (FID-50K 기준 가이던스 사용 시 2.18 대비 1.83, 미사용 시 5.15 대비 4.33), Muon 변형으로 학습하는 레시피를 가진 RAEv2-B에서도 앞섰습니다 (3.39 대비 3.19).
그림 1은 이러한 실행 결과를 보여주며, 각 결과는 Claude Opus 5.5가 에이전트(agent) 역할을 수행했습니다. 우리는 우리의 반응형 옵티마이저를 OPAL (OPtimizer with an Agent in the Loop)이라는 이름으로 명명했습니다. 종합적으로 볼 때, 이 결과들은 원샷 환경에서 미리 결정된 레시피를 고수하는 것보다 반응형 옵티마이저가 더 나은 솔루션을 제공할 수 있음을 보여줍니다.
modded-nanogpt Track 3
| Task | 162M GPT의 modded-nanogpt 최적화 트랙(Track 3)은 모든 설정에서 동일한 실행을 수행합니다: 공식 NanoGPT, 트랙 자체 (공식 배치 524,288 토큰/업데이트에서의 FineWeb); NanoGPT, ¼ 배치 (FineWeb, 13,000 업데이트의 131,072 토큰); 중국어, 태국어 및 키르기스어에 대한 NanoGPT (1 GB Goldfish 중국어, 태국어 및 키르기스어 세트, 3,000 업데이트의 524,288 토큰) |
|---|
| Metric | 공식 NanoGPT: 검증 손실이 3.28에 도달한 첫 번째 업데이트, 트랙 목표치 (25 업데이트마다 검증); ¼ 배치: 13,000 업데이트 후의 검증 손실; 중국어, 태국어 및 키르기스어: 3,000 업데이트 후의 검증 손실. #46과 #36의 OPAL은 #46 자체 판독값(마지막 업데이트에 걸친 가중치의 이동 평균)을 사용하고, 다른 기록들은 마지막 가중치를 사용합니다. 중국어와 태국어에서는 OPAL 막대가 세 번의 실행 평균입니다. 손실 막대는 해당 설정의 최저값보다 적어도 0.05 낮은 가장 큰 0.1의 배수 c에서 시작합니다: 3.2 (¼ 배치), 4.0 (중국어), 4.5 (태국어) 및 3.7 (키르기스어) |
| Baseline | 리더보드의 다섯 가지 기록, #46, #30, #21 (Shampoo), #36 (조정된 Muon 기준선) 및 #2 (AdamW), 그리고 #36의 OPAL. 리더보드는 #46, #30, #21 및 #36에 대해 각각 2,690, 2,930, 4,100, 3,250 업데이트를 나열합니다; 공식 NanoGPT에서 #2의 막대는 리더보드의 5,625입니다: 우리의 #2 실행은 5,625 업데이트 후 3.28449에 끝났으며 3.28에 도달하지 못했습니다; ¼ 배치에서는 모든 기록이 학습률 전체 ×¼로 실행되었습니다. |
| Task | Amharic(골드피쉬)의 1GB 데이터로 처음부터 학습시킨 125M GPT, 10 에포크 (64,626 업데이트) |
|---|
| Metric | FLORES-200 (dev + devtest) 기준 문장당 비트 수 |
| Baseline | 공식 Goldfish 레시피(AdamW, 학습률 1e-4, 가중치 감쇠 0.01, 10% 워밍업 후 선형 감소; 그래디언트 클리핑 1.0); 추가적인 아홉 가지 정적 레시피들로, 각각 모든 10 에포크에 대해: 가중치 감쇠 0.1을 가진 학습률 6e-4, 1e-3, 2e-3 및 4e-3의 AdamW (네 값 스윕), 가중치 감쇠 0.1 또는 0.01을 가진 학습률 3e-4의 AdamW, 그리고 가중치 감쇠 0.1을 가진 학습률 1e-4의 AdamW, 그리고 공식 레시피의 숨겨진 블록 행렬에 적용된 Muon과 학습률 3e-4, 가중치 감쇠 0.1. OPAL은 두 번 실행됨: 공식 레시피(기본값 사용, 50개 결정; 기본 표시)와 스윕의 최적 지점인 AdamW (학습률 2e-3, 200개 결정; 스윕과 함께 표시)에서 |
LightningDiT-XL/1
| Task | VA-VAE 잠재 공간에서 ImageNet 256×256 이미지 생성; LightningDiT-XL/1, 64 에포크 |
|---|
| Metric | 최종 FID-50K (가이던스 유무에 따라); 곡선은 학습 중의 FID-10K를 보여줌 |
| Baseline | 우리가 재학습한 공식 레시피와 저자들이 공개한 체크포인트(가이던스 없음 시 FID-50K 5.23, 있음 시 2.16) |
RAEv2-B
| Task | 표현 표현 오토인코더(representation-autoencoder) 잠재 공간에서 ImageNet 256×256 이미지 생성; RAEv2-B, 20 에포크 |
|---|
| Metric | FID-50K |
| Baseline | 이미 Muon 변형인 GMuon으로 학습되는 공식 레시피 |
그림 1. OPAL (OPtimizer with an Agent in the Loop), 우리의 반응형 옵티마이저를 다양한 설정에서 보여줍니다. 최적화 하이퍼파라미터가 크게 조정되지 않은 대부분의 설정에서, 반응형 옵티마이저는 표준적인 선택에 비해 학습 손실과 테스트 성능을 향상시킵니다. modded-nanogpt에서는 #2가 AdamW이고, #21은 Shampoo이며, #36은 튜닝된 Muon이고, #30과 #46은 더 정교한 Muon 기반 레시피입니다.
실용적인 이점 외에도, 우리는 제로샷 에이전트가 스스로 관찰하고 분석한 학습 역학을 바탕으로 내리는 자발적인 결정에 흥미를 느낍니다. 다음 두 섹션에서는 반응형 옵티마이저(reactive optimizer)가 무엇을 감시하기로 선택하는지, 그리고 이에 대응하여 무엇을 변경하는지 탐구합니다.
예를 들어, DiT와 SiT의 표준 레시피[19]는 하나의 상수 학습률로 AdamW를 사용하여 훈련합니다. 이는 Muon과 그 변형 모델들이 확산 트랜스포머(diffusion transformers)에 적용되고([20] CMuon; [21] Li et al.; [22] RAEv2의 GMuon), 최근 확산 모델용 옵티마이저에 대한 통제된 벤치마크[23]가 나오면서 변화하기 시작했습니다.
2. 자발적 분석을 통한 업데이트 (Updating with spontaneous analysis)
그림 2. 에이전트가 훈련 중에 반복적으로 작성하고 재작성하는 두 개의 프로그램. 코드를 전체 인터페이스로 제공함으로써, 우리는 에이전트에게 무엇을 관찰할지 그리고 무엇을 변경할지를 선택할 수 있는 최대의 자유를 부여합니다.
그림 2는 에이전트가 옵티마이저 코드를 즉석에서 업데이트하는 방법을 보여줍니다. 구체적으로, 우리는 에이전트가 최적화의 모든 단계에 사용되는 updater.py 파일을 재작성하도록 허용합니다. 코드를 전체 인터페이스로 사용하는 이유는 이것이 에이전트가 최적화의 어느 부분이라도 쉽게 변경할 수 있게 해주기 때문입니다. 하지만 에이전트가 변경을 가하기 전에, 그것은 스스로 훈련이 어떻게 진행되고 있는지 관찰해야 하며, 이는 에이전트가 재작성할 수 있는 두 번째 프로그램인 analysis.py로 이어집니다.
에이전트는 의사 결정을 내리는 데 어떤 정보를 사용해야 할까요? 원시 훈련 상태(raw training states)는 비실용적일 것입니다. 가중치(weights), 활성화 값(activations), 기울기(gradients)가 컨텍스트를 압도할 것이기 때문입니다. 대신, 학습 역학의 통계 및 기하학(statistics and geometry)이 결정에 도움이 될 수 있습니다. 예를 들어, 각 파라미터 그룹의 기울기 노름(gradient norm), 변경하는 가중치 옆 업데이트 크기, 연속적인 기울기 간의 코사인 값, 그리고 비정상적으로 큰 기울기 스파이크를 보이는 텐서 등이 그것입니다. 따라서 옵티마이저와 마찬가지로, 우리는 에이전트에게 가장 많은 자유를 부여합니다: 즉, analysis.py라는 프로그램을 작성하게 합니다.
, 학습 역동성(training dynamics)에서 알고 싶은 모든 것을 추출합니다. 자체적으로 작성된 프로브(probes)는 훈련 전반에 걸쳐 실행되며, 훈련 상태로부터 에이전트가 요청하는 모든 것을 가져와 다음 결정을 위해 에이전트에게 전달합니다.
우리는 양방향 루프(bidirectional loop)를 완성합니다. analysis.py를 통해 훈련은 에이전트가 보고 싶은 어떤 형태로든 에이전트에게 도달하며, updater.py를 통해 에이전트의 결정은 훈련의 모든 단계에 전달됩니다. 이 두 프로그램은 에이전트가 과정 중에 최대의 유연성으로 새로운 최적화 알고리즘(optimization algorithms)을 고안할 수 있게 합니다.
3. 최적화 과정 중 동적으로 변화하는 것은 무엇인가?
이 섹션에서는 반응형 옵티마이저(reactive optimizer)가 훈련 전반에 걸쳐 무엇을 하는지 살펴보고, 이러한 결과 뒤에 숨겨진 결정들을 검토합니다. 요약하자면, 그림 3은 각 에이전트가 언제, 어떻게, 그리고 왜 코드를 수정했는지 보여줍니다.
프로그램을 변경한 모든 결정들
