TACO가 옵티마이저 상태 메모리를 174배 줄여 32B 파인튜닝을 가능케 한다

무엇인가

LLM 전체 파라미터 파인튜닝은 옵티마이저 상태 메모리 때문에 발목이 잡힌다. 논문은 13B 모델을 예로 든다. BF16 파라미터가 약 26GB를 차지하는데, AdamW는 FP32 모멘트 버퍼 두 개로 약 104GB를 더 쓰고 그래디언트가 약 26GB를 더해 활성값이나 FP32 마스터 복사본을 계산하기 전에 이미 약 156GB가 된다. 즉 AdamW의 영구 옵티마이저 상태만 BF16 파라미터 메모리의 약 4배다. 기존 대응은 옵티마이저 상태를 양자화하거나(AdamW8bit, FlashAdamW), 2차 모멘트를 분해하거나(Adafactor), 그래디언트를 저랭크 부분공간에 사영하거나(GaLore, APOLLO-Mini), 아예 역전파를 포기하는 0차 방법(MeZO, ZO-Muon, HiZOO)이었다. Muon은 2차원 가중치를 행렬로 보고 스펙트럼 기하에서의 최급강하 방향을 근사해 메모리를 줄이지만 상태가 여전히 조밀하고, AdamW로 사전학습된 모델에 적용하면 기하 차이 때문에 성능이 떨어지는 옵티마이저 불일치가 보고된다.

어떻게 동작하나

TACO의 핵심은 업데이트 기하 자체를 바꾸는 것이다. Muon이 정규화된 2→2(스펙트럼) 연산자 노름을 쓴다면, TACO는 차원 정규화된 1→1 연산자 노름을 쓴다. 이 노름은 ‖D‖ = (n/m)·max_j‖d_j‖₁ 로 최대 열 노름이 되고, 제약이 열마다 분리되기 때문에 각 열은 ℓ1 공의 극점, 즉 그래디언트 절댓값이 가장 큰 좌표 하나에 예산을 몰아준다. 그 결과 닫힌 형태 해는 D* = -(m/n)·S(G)이고, S(G)는 각 열에서 최대 크기 원소의 부호만 남기고 나머지를 0으로 만든다. m×n 업데이트의 비영 원소는 최대 n개다. 이 희소성은 사후적 그래디언트 희소화가 아니라 국소 최적화 문제의 정확한 해에서 나온다는 점을 저자들은 강조한다. m/n 인수는 별도 계층별 스케일 하이퍼파라미터가 아니라 차원 정규화에서 따라 나온다. 이력 없는 vanilla TACO에 대해 저자들은 표준 매끄러움 가정 아래 ε-정상점에 O(ε⁻²) 반복으로 도달한다는 비볼록 수렴 정리(정리 1)를 제시한다.

무엇과 다른가

이론적 위치 설정도 흥미롭다. 고정 입력 z에 대해 목표가 층 출력에만 의존한다는 단순화 설정에서, TACO와 Adam 프록시는 같은 연속 집합 C(W₀, sign(z)) 안에 머물고 같은 극한 W* = W₀ + (u* - W₀z)sign(z)ᵀ/‖z‖₁ 에 도달한다(정리 2). 반면 Muon은 C(W₀, z)를 따라가며 W* = W₀ + rzᵀ/‖z‖₂² 라는 다른 해에 도달한다(정리 3). Muon에서 출발했지만 파인튜닝에서는 Adam 쪽에 가깝게 행동한다는 주장이다. 또한 열별 승자가 행에 고르게 분포한다는 가정(정리 4) 아래 ‖D‖₂ = (m/n)√(max_i c_i) = Θ(√(m/n)) 로 μP 스펙트럼 스케일링을 만족해, 모델 폭이 커져도 학습률을 그대로 옮길 수 있는 근거를 댄다.

어떻게 쓰나

실용 구현의 요점은 조밀한 그래디언트 이력을 버리는 것이다. 미니배치 그래디언트에 S를 바로 적용하면 선택 좌표가 노이즈에 흔들리는데, 조밀한 EMA 이력을 유지하면 메모리 이점이 사라진다. TACO는 각 열에서 절댓값이 큰 k개 원소만 남기는 동적 헤비히터 방식으로 M_t = H_k(μM_{t-1} + (1-μ)G_t) 를 유지한다(μ=0.95, k=16). 남긴 값은 FP8 E4M3 1바이트, 행 인덱스는 int32 4바이트로 저장해 m×n 행렬당 5kn 바이트만 쓴다(m ≤ 2¹⁵면 int16으로 3kn 바이트). 영구 상태 복잡도가 O(mn)에서 O(kn)으로 내려간다. 행렬 파라미터에만 TACO를 적용하고 스칼라·벡터 파라미터는 보조 AdamW가 맡는다. PyTorch autograd 훅으로 그래디언트가 계산되는 즉시 제자리에서 갱신하고 방출하며, 그래디언트 누적은 1로 두고 그래디언트 체크포인팅으로 활성값 메모리를 줄인다.

전제와 한계

OPT-13B 실험에서 TACO는 AdamW8bit 대비 영구 옵티마이저 상태를 174배 줄이고(27.7GB → 0.16GB), 최대 학습 메모리를 2.9배 줄인다(80.6GB → 27.5GB). SST-2 효율 비교에서는 가중치 25.7GB 위에 추가 학습 메모리가 1.8GB뿐이고, 그중 영구 옵티마이저 상태는 0.16GB다. 같은 조건에서 Adafactor·GaLore·FlashAdamW는 추가 메모리가 각각 26.9GB, 31.0GB, 42.0GB로 최대 메모리가 52.6GB, 56.7GB, 67.7GB가 된다. 정확도와 속도는 SST-2에서 TACO 94.2%(276 tokens/s), Adafactor 95.5%(350 tokens/s), GaLore 95.7%(31 tokens/s), FlashAdamW 95.5%(547 tokens/s)로, TACO는 정확도-메모리 트레이드오프의 저메모리 끝을 차지하면서 실용적 처리량을 유지한다. 규모 실험에서는 OPT 1.3B~30B, Qwen3 8B~32B가 단일 H100 80GB 안에 들어가고, OPT-30B의 최대 메모리는 SST-2 62.5GB, RTE 64.7GB, BoolQ 68.1GB다(Adam은 이 모델을 단일 H100에 올리지 못한다). 모델 크기별 하이퍼파라미터 튜닝 없이 같은 설정을 쓰며, OPT 네 크기 모두 학습률 3×10⁻⁵에서 최적 검증 정확도를 보인다. 아키텍처 전이에서는 Llama-3.1-8B와 Qwen3-32B가 세 과제 모두 강했고, Pythia-12B와 Mistral-24B는 특히 RTE와 BoolQ에서 과제 의존 편차가 컸다. 절제 실험에서 k=16은 옵티마이저 상태 38.3MB로 조밀 이력 대비 34.4배 줄이면서 BoolQ와 RTE 최고 성능을 냈고, μ=0.95는 8개 과제 중 6개를 개선했으며(MultiRC +8.9 F1A, CB +7.1%) SQuAD와 DROP이 예외였다. 스텝 예산을 2K에서 16K로 늘리면 BoolQ +6.5%, RTE +4.3%, SQuAD +2.1 F1이 더 오른다. 열별 Top-1보다 조밀한 Top-k 변형은 성능이 떨어지고 Top-8 이상에서 특히 나빠진다.

개발자 입장에서 이 논문은 단일 80GB GPU에서 30B급 모델의 전체 파라미터 파인튜닝을 시도할 때의 실질적 선택지를 제시한다. 다만 적용 전에 확인할 것이 있다. TACO는 행렬 파라미터에만 쓰이고 나머지는 AdamW가 담당하는 혼합 구성이므로, 학습 스크립트에서 어떤 파라미터를 어느 옵티마이저에 배정할지 정해야 한다. autograd 훅 기반 제자리 갱신과 그래디언트 누적 1 설정은 기존 학습 루프와 충돌할 수 있어 통합 비용이 든다. 또 저자들이 공개한 기본 설정(k=16, μ=0.95, 2K 스텝)은 그대로 쓰되, 2K 스텝이 성능 상한이 아니라는 점과 과제별 편차(Pythia·Mistral의 RTE/BoolQ)를 감안해 스텝 예산과 검증을 늘려 잡는 편이 안전하다.

저자들이 밝힌 전제와 한계는 분명하다. Adam과의 연속 기하 동등성과 Muon과의 분리는 고정 입력에 목표가 의존한다는 단순화된 설정(가정 2, 3)에서 증명된 것이지 일반 파인튜닝 전체에 대한 결과가 아니다. μP 스케일링 정리도 열별 승자가 행에 고르게 분포한다는 가정(정리 4)에 의존한다. 실용 구현의 희소 그래디언트 이력은 조밀 이력의 근사이며, 이를 정당화하는 충분 조건은 별도 정리(정리 5)와 부록에 실려 있다. 실험 범위는 1.3B~32B 모델, 단일 80GB H100, 8개 과제로 한정되고, 기본 2K 스텝 예산은 성능 상한이 아니다.