AdamW 국룰 시대에 등장한 Lion, 써보니 어떤가?

TL;DR: Lion은 메모리를 33% 아끼면서 비슷한 성능을 냅니다. 하지만 하이퍼파라미터 튜닝 없이 쓰면 오히려 손해입니다.

1. Optimizer의 진화: SGD에서 Lion까지

딥러닝 optimizer의 역사는 "어떻게 하면 더 빠르고 안정적으로 수렴할 수 있을까?"에 대한 답을 찾는 여정입니다.

1.1 SGD의 시대 (1950s~2010s)

모든 것의 시작은 Stochastic Gradient Descent입니다:

$\theta_{t+1} = \theta_t - \eta \cdot g_t$

단순하지만 문제가 있었습니다:

Noisy gradient: mini-batch sampling으로 인한 variance
동일한 learning rate: 모든 파라미터에 같은 크기로 업데이트
Saddle point: 고차원에서 탈출이 어려움

1.2 Momentum의 등장 (1999)

Polyak의 momentum이 첫 번째 개선이었습니다:

$\begin{aligned}v_t &= \gamma v_{t-1} + \eta g_t \\\theta_{t+1} &= \theta_t - v_t\end{aligned}$

과거 gradient를 누적해서 "관성"을 부여합니다. 이는 noisy gradient를 평균화하고, narrow valley에서의 진동을 줄여줍니다.

1.3 Adaptive Learning Rate의 시대 (2011~)

AdaGrad (2011): 파라미터별로 다른 learning rate

$\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{G_t + \epsilon}} \cdot g_t$

여기서 $G_t$ 는 과거 gradient 제곱의 누적합입니다. 자주 업데이트되는 파라미터는 lr이 작아지고, 드물게 업데이트되는 파라미터는 lr이 유지됩니다.

RMSprop (2012): AdaGrad의 "망각" 버전

$v_t = \beta v_{t-1} + (1-\beta) g_t^2$

Exponential moving average를 사용해 오래된 gradient는 잊어버립니다.

1.4 Adam: 두 아이디어의 결합 (2015)

Kingma & Ba의 Adam은 momentum과 adaptive learning rate를 결합했습니다:

$\begin{aligned}m_t &= \beta_1 m_{t-1} + (1 - \beta_1) g_t \quad \text{(1st moment, momentum)} \\v_t &= \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \quad \text{(2nd moment, adaptive lr)} \\\theta_t &= \theta_{t-1} - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}\end{aligned}$

여기서 $\hat{m}_t, \hat{v}_t$ 는 bias-corrected 버전입니다.

Adam은 대부분의 태스크에서 "그냥 작동"했고, 빠르게 표준이 되었습니다.

1.5 AdamW: Weight Decay의 재발견 (2017)

Loshchilov & Hutter는 중요한 문제를 발견했습니다: Adam에서 L2 regularization이 제대로 작동하지 않는다.

문제: Adam은 gradient에 weight decay 항을 포함시킵니다.

$g_t = \nabla f(\theta) + \lambda \theta$

이러면 adaptive learning rate가 regularization 효과를 상쇄합니다. gradient가 작은 파라미터는 lr이 커지므로, weight decay의 효과도 커지는 비일관성이 발생합니다.

해결 (AdamW): Weight decay를 gradient와 분리(decouple)합니다.

$\theta_t = \theta_{t-1} - \eta \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_{t-1} \right)$

이 간단한 변경이 generalization을 크게 개선했고, AdamW가 새로운 표준이 되었습니다.

1.6 Lion: AutoML이 찾아낸 Optimizer (2023)

여기서 Google의 접근이 흥미롭습니다. 사람이 설계하지 않고, 프로그램이 발견하게 했습니다.

Lion 논문의 핵심 아이디어:

Optimizer를 "프로그램"으로 표현 (symbolic representation)
진화 알고리즘으로 수천 개의 optimizer 탐색
실제 태스크에서 평가하여 선별

그 결과 발견된 것이 Lion입니다. 놀랍게도, 인간이 설계한 것보다 단순하면서도 효과적이었습니다.

2. Lion의 핵심 아이디어

2.1 수식으로 보는 차이

AdamW Update:

$\begin{aligned}m_t &= \beta_1 m_{t-1} + (1 - \beta_1) g_t \\v_t &= \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \\\hat{m}_t &= \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t} \\\theta_t &= \theta_{t-1} - \eta \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_{t-1} \right)\end{aligned}$

Lion Update:

$\begin{aligned}\theta_t &= \theta_{t-1} - \eta \lambda \theta_{t-1} \quad \text{(weight decay first)} \\c_t &= \beta_1 m_{t-1} + (1 - \beta_1) g_t \\\theta_t &= \theta_t - \eta \cdot \text{sign}(c_t) \\m_t &= \beta_2 m_{t-1} + (1 - \beta_2) g_t\end{aligned}$

2.2 핵심 차이점 분석

구분	AdamW	Lion
Momentum 저장	`m` + `v` (2x 메모리)	`m`만 (1x 메모리)
Update 크기	Adaptive ($\hat{m}/\sqrt{\hat{v}}$)	상수 ($\pm \eta$)
Weight Decay	Update와 동시	Update 전에 적용
Bias Correction	있음	없음
β 기본값	(0.9, 0.999)	(0.9, 0.99)

2.3 Lion의 독특한 구조: 두 개의 β

Lion에서 가장 특이한 점은 β₁과 β₂가 다른 역할을 한다는 것입니다.

β₁ (0.9): 업데이트 방향 계산에 사용
c_t = β₁ * m_{t-1} + (1-β₁) * g_t
→ sign(c_t)로 이번 스텝의 업데이트 방향 결정

β₂ (0.99): 다음 스텝을 위한 momentum 저장에 사용
m_t = β₂ * m_{t-1} + (1-β₂) * g_t
→ 더 smooth한 momentum을 유지

이 분리가 왜 효과적인지는 완전히 이해되지 않았습니다. AutoML이 발견했기 때문에 "왜 작동하는가"에 대한 이론적 설명이 부족합니다. 하지만 경험적으로는 잘 작동합니다.

2.4 Lion의 `sign()` 연산

Lion의 핵심은 sign() 연산입니다:

$\text{sign}(x) = \begin{cases} +1 & \text{if } x > 0 \\ 0 & \text{if } x = 0 \\ -1 & \text{if } x < 0 \end{cases}$

이는 gradient의 크기(magnitude)를 무시하고 방향(direction)만 사용한다는 의미입니다.

직관적으로 이상해 보이는 이유:

Gradient 크기는 "얼마나 움직여야 하는가"에 대한 정보
이걸 버리면 정보 손실 아닌가?

하지만 실제로는:

모든 파라미터가 동일한 크기로 업데이트됨
Learning rate η가 step size를 결정
Implicit regularization 효과 발생

3. 왜 sign()이 작동하는가? 이론적 분석

3.1 Implicit Regularization

sign() 연산은 gradient의 크기에 따라 다르게 작용합니다:

큰 gradient (|g| >> 0): sign(g) * η = η (크기 1로 "억제")
작은 gradient (|g| ≈ 0): sign(g) * η = η (크기 1로 "증폭")

이는 L∞ regularization과 유사한 효과를 냅니다. 모든 파라미터의 업데이트 크기가 균일해지므로, 특정 파라미터가 과도하게 커지는 것을 방지합니다.

3.2 Noise Robustness

Stochastic gradient는 "signal + noise"로 볼 수 있습니다:

$g_t = \nabla f(\theta) + \epsilon_t$

sign() 연산의 특성:

|signal| > |noise|이면: sign(g) ≈ sign(signal)
즉, noise의 크기가 방향을 바꿀 만큼 크지 않으면 무시됨

이는 noisy gradient 환경에서 더 안정적인 업데이트를 가능하게 합니다.

3.3 Loss Landscape와 Flat Minima

Sharp minima vs Flat minima 가설:

Sharp minima: 좁은 계곡, generalization이 안 좋음
Flat minima: 넓은 분지, generalization이 좋음

sign() 업데이트의 효과:

상수 크기의 step으로 이동
Sharp minima에서는 "튕겨나감" (step이 너무 커서)
Flat minima에서만 안정적으로 머무름

이 가설은 Lion이 더 좋은 generalization을 보이는 이유를 설명할 수 있습니다.

3.4 Preconditioner 관점에서의 분석

Optimizer를 일반화하면:

$\theta_{t+1} = \theta_t - \eta \cdot P_t \cdot g_t$

여기서 $P_t$ 는 preconditioner입니다.

Optimizer	Preconditioner $P_t$
SGD	$I$ (identity)
Adam	$\text{diag}(1/\sqrt{\hat{v}_t})$
Lion	$\text{diag}(\text{sign}(\cdot)/g_t)$ (effective)

기하학적 해석:

SGD: Loss surface를 있는 그대로 탐색
Adam: Loss surface를 "늘리거나 줄여서" 탐색 (coordinate-wise scaling)
Lion: 모든 방향으로 동일한 크기로 이동 (isotropic in update magnitude)

4. Lion 논문의 주요 실험 결과

4.1 Image Classification (ImageNet)

논문에서 ViT-B/16, ViT-L/16 등 다양한 모델로 실험했습니다.

주요 발견:

Lion은 AdamW와 비슷하거나 약간 더 좋은 성능
메모리 사용량은 확실히 적음
Large batch (4K+)에서 더 안정적

4.2 Language Modeling

GPT-2 스타일 모델로 실험:

주요 발견:

Perplexity 측면에서 AdamW와 동등
학습 안정성이 좋음
LLM scale에서 메모리 이득이 크게 작용

4.3 Diffusion Models

Image generation 태스크에서:

주요 발견:

FID score에서 AdamW와 비슷
학습 속도는 약간 느릴 수 있음
하지만 최종 품질은 동등

4.4 논문의 핵심 권장사항

Lion 논문에서 강조하는 하이퍼파라미터 가이드:

하이퍼파라미터	AdamW 대비 Lion
Learning rate	3-10배 작게
Weight decay	3-10배 크게
β₁	0.9 (동일)
β₂	0.99 (AdamW의 0.999보다 작음)

중요: 이 비율은 태스크에 따라 다릅니다. Vision은 3배, NLP는 10배에 가까운 경향이 있습니다.

5. 실험: CIFAR-10에서의 비교

5.1 실험 설계

python

# AdamW 설정 (기본값)
adamw_config = {
    "lr": 1e-3,
    "weight_decay": 0.01,
    "betas": (0.9, 0.999)
}

# Lion 설정 (논문 권장)
lion_config = {
    "lr": 1e-4,       # AdamW의 1/10
    "weight_decay": 0.1,  # AdamW의 10배
    "betas": (0.9, 0.99)
}

5.2 결과

Metric	AdamW	Lion
Final Val Acc	84.2%	83.8%
Best Val Acc	85.1%	84.5%
Optimizer Memory	2.4 MB	1.2 MB
Generalization Gap	4.2%	3.8%

5.3 시각화

6. 하이퍼파라미터 튜닝 완전 가이드

6.1 Learning Rate 찾기

Lion의 lr은 AdamW보다 훨씬 작아야 합니다. 이유:

sign()이 magnitude를 1로 고정
AdamW의 $1/\sqrt{v}$ 가 보통 < 1이므로 effective lr이 작아짐
Lion은 이 scaling이 없어서 lr 자체를 줄여야 함

LR Range Test 방법:

python

# 1. AdamW 최적 lr 찾기 (기존 방법대로)
# 2. 그 값의 1/3 ~ 1/10로 Lion lr 설정
# 3. 작은 데이터셋에서 빠르게 검증

adamw_lr = 1e-3  # 기존에 찾은 값
lion_lr_candidates = [adamw_lr / 3, adamw_lr / 5, adamw_lr / 10]

6.2 Weight Decay 설정

Lion에서 weight decay는 더 중요합니다:

Update magnitude가 고정이므로 WD의 상대적 영향이 큼
WD가 너무 작으면 overfitting
WD가 너무 크면 underfitting

경험적 규칙:

Task	AdamW WD	Lion WD
Vision (ViT)	0.05	0.5
Vision (ResNet)	1e-4	1e-3
NLP (BERT)	0.01	0.1
NLP (GPT)	0.1	1.0
Diffusion	0.01	0.05

6.3 β 값 조정

대부분의 경우 기본값 (0.9, 0.99)가 작동합니다. 하지만:

β₁ (momentum for update):

더 작게 (0.8): 현재 gradient에 더 민감
더 크게 (0.95): 과거 방향을 더 유지

β₂ (momentum for storage):

더 작게 (0.95): 빠른 적응, 덜 smooth
더 크게 (0.999): 느린 적응, 더 smooth

일반적으로:

Short training: β₂를 약간 낮춤 (0.95~0.99)
Long training: β₂를 높임 (0.99~0.999)

6.4 Batch Size와의 관계

Lion은 large batch에서 더 안정적입니다:

Batch Size	AdamW 안정성	Lion 안정성
32-128	좋음	보통
256-1024	좋음	좋음
2048-4096	보통	좋음
8192+	주의 필요	좋음

이유:

sign()이 gradient magnitude를 정규화
Large batch의 작은 variance가 문제되지 않음

6.5 AdamW에서 Lion으로 전환 체크리스트

python

def adamw_to_lion_config(adamw_config, task_type="general"):
    """
    AdamW 설정을 Lion으로 변환

    Args:
        adamw_config: dict with lr, weight_decay, betas
        task_type: "vision", "nlp", "diffusion", "general"
    """
    factors = {
        "vision": {"lr_div": 3, "wd_mul": 10},
        "nlp": {"lr_div": 10, "wd_mul": 10},
        "diffusion": {"lr_div": 5, "wd_mul": 5},
        "general": {"lr_div": 7, "wd_mul": 7},
    }

    f = factors[task_type]

    return {
        "lr": adamw_config["lr"] / f["lr_div"],
        "weight_decay": adamw_config["weight_decay"] * f["wd_mul"],
        "betas": (0.9, 0.99),  # Lion 기본값
    }

# 사용 예시
adamw = {"lr": 1e-3, "weight_decay": 0.01, "betas": (0.9, 0.999)}
lion = adamw_to_lion_config(adamw, task_type="nlp")
# lion = {"lr": 1e-4, "weight_decay": 0.1, "betas": (0.9, 0.99)}

6.6 튜닝 실패 시 디버깅

증상 1: Loss가 발산함

원인: lr이 너무 큼
해결: lr을 2-3배 줄임

증상 2: 학습이 너무 느림

원인: lr이 너무 작거나 WD가 너무 큼
해결: lr을 1.5-2배 올리거나 WD를 줄임

증상 3: Train은 잘 되는데 Val이 안 좋음

원인: WD가 너무 작음 (overfitting)
해결: WD를 2-3배 올림

증상 4: Train/Val 모두 안 좋음

원인: WD가 너무 큼 (underfitting)
해결: WD를 2-3배 줄임

7. 실전 적용 가이드

7.1 Lion을 써야 할 때

메모리가 부족할 때: LLM/Large Diffusion 모델에서 optimizer state가 병목
대규모 batch 학습: Lion은 큰 batch에서 더 안정적
충분한 튜닝 시간이 있을 때: hyperparameter 민감도가 높음

7.2 Lion을 피해야 할 때

빠른 프로토타이핑: AdamW 기본값이 대부분 작동
작은 모델/데이터셋: 메모리 이득이 미미
기존 AdamW 레시피가 있을 때: 검증된 설정을 버릴 이유 없음

7.3 Gradient Clipping과 함께 사용

Lion 논문에서는 gradient clipping을 권장합니다:

python

# Lion + gradient clipping 조합
optimizer = Lion(model.parameters(), lr=1e-4, weight_decay=0.1)

for batch in dataloader:
    loss = model(batch)
    loss.backward()

    # Gradient clipping (Lion 논문 권장)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    optimizer.step()
    optimizer.zero_grad()

7.4 Learning Rate Scheduler와의 조합

Lion도 일반적인 scheduler와 잘 작동합니다:

python

optimizer = Lion(model.parameters(), lr=1e-4, weight_decay=0.1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=num_epochs,
    eta_min=1e-6  # Lion의 경우 더 작은 min_lr
)

Warmup 참고:

Lion은 warmup이 짧아도 됨 (AdamW 대비)
논문에서는 전체 step의 1-3% warmup 사용

8. 메모리 절감의 실제 임팩트

8.1 이론적 계산

Optimizer state 메모리 = 파라미터 수 × 저장하는 moment 수 × dtype 크기

Optimizer	Moment 수	7B 모델 (bf16)
SGD	0	0 GB
SGD+momentum	1	14 GB
Adam/AdamW	2	28 GB
Lion	1	14 GB

8.2 7B LLM 학습 시나리오

Component	AdamW	Lion
Model (bf16)	14 GB	14 GB
Gradients	14 GB	14 GB
Optimizer m	14 GB	14 GB
Optimizer v	14 GB	0 GB
Total	56 GB	42 GB

25% 메모리 절감으로 가능해지는 것들:

Batch size 2배 증가
Sequence length 4K → 8K
A100 80GB에서 더 큰 모델 학습 가능

8.3 ZeRO와의 시너지

Lion + ZeRO-2 조합은 매우 효율적입니다:

ZeRO-2 with AdamW (4 GPU):
- 각 GPU의 optimizer state: 28GB / 4 = 7GB

ZeRO-2 with Lion (4 GPU):
- 각 GPU의 optimizer state: 14GB / 4 = 3.5GB

추가 절감: 3.5GB per GPU

9. 다른 Optimizer들과의 비교

9.1 LAMB (Layer-wise Adaptive Moments)

BERT의 large batch 학습을 위해 개발:

Layer별로 다른 lr 적용
매우 큰 batch (32K+)에서 효과적
Lion보다 복잡한 구현

9.2 Shampoo

2nd order 정보를 효율적으로 근사:

이론적으로 더 빠른 수렴
하지만 계산/메모리 비용이 큼
실무에서는 잘 안 쓰임

9.3 Adafactor

Transformer를 위해 설계:

2nd moment를 factorize하여 메모리 절감
Lion보다 복잡하지만 비슷한 메모리 이득
T5 학습에 사용됨

9.4 비교 요약

Optimizer	메모리	수렴 속도	튜닝 난이도	추천 상황
SGD+M	1x	느림	낮음	Vision, 최고 성능
AdamW	2x	빠름	낮음	기본 선택
Lion	1x	보통	높음	메모리 제한
LAMB	2x	빠름	중간	Very large batch
Adafactor	1.5x	보통	중간	Transformer

10. 결론: 국룰은 국룰인 이유가 있다

Lion은 분명히 흥미로운 optimizer입니다. AutoML이 발견했다는 점, 단순하면서도 효과적이라는 점에서 주목할 가치가 있습니다.

하지만 현실에서는:

AdamW가 여전히 안전한 선택: 수년간 검증된 설정과 레시피
Lion은 메모리가 critical할 때: 대규모 모델 학습에서 진가 발휘
무조건적인 교체는 위험: 튜닝 없이 Lion으로 바꾸면 성능 하락 확률 높음

10.1 언제 Lion을 고려할까?

체크리스트:
□ GPU 메모리가 부족해서 batch size를 줄여야 하는가?
□ Optimizer state가 전체 메모리의 상당 부분을 차지하는가?
□ 하이퍼파라미터 튜닝에 시간을 투자할 수 있는가?
□ Large batch (1K+) 학습을 하는가?

3개 이상 체크되면 Lion을 시도해볼 가치가 있습니다.

10.2 실무자의 시선

"새로운 optimizer가 나올 때마다 '이번엔 다르다'고 하지만, 결국 AdamW가 살아남는다. Lion도 예외는 아닐 것이다. 다만, 메모리 효율이 필요한 특정 상황에서는 충분히 고려할 가치가 있다."

Lion의 진정한 가치는 "AdamW를 대체하는 것"이 아니라, 메모리와 성능의 trade-off가 필요한 상황에서 새로운 선택지를 제공하는 것입니다.

References

Chen, X., et al. "Symbolic Discovery of Optimization Algorithms." arXiv:2302.06675 (2023)
Loshchilov, I., & Hutter, F. "Decoupled Weight Decay Regularization." ICLR 2019
Kingma, D. P., & Ba, J. "Adam: A Method for Stochastic Optimization." ICLR 2015
You, Y., et al. "Large Batch Optimization for Deep Learning: Training BERT in 76 minutes." ICLR 2020
Shazeer, N., & Stern, M. "Adafactor: Adaptive Learning Rates with Sublinear Memory Cost." ICML 2018

이 글의 실험 코드는 [GitHub](https://github.com/sotaaz/adamw_vs_lion_optimizer)에서 확인할 수 있습니다.

AdamW 국룰 시대에 등장한 Lion, 써보니 어떤가?

1. Optimizer의 진화: SGD에서 Lion까지

1.1 SGD의 시대 (1950s~2010s)

1.2 Momentum의 등장 (1999)

1.3 Adaptive Learning Rate의 시대 (2011~)

1.4 Adam: 두 아이디어의 결합 (2015)

1.5 AdamW: Weight Decay의 재발견 (2017)

1.6 Lion: AutoML이 찾아낸 Optimizer (2023)

2. Lion의 핵심 아이디어

2.1 수식으로 보는 차이

2.2 핵심 차이점 분석

2.3 Lion의 독특한 구조: 두 개의 β

2.4 Lion의 sign() 연산

3. 왜 sign()이 작동하는가? 이론적 분석

3.1 Implicit Regularization

3.2 Noise Robustness

3.3 Loss Landscape와 Flat Minima

3.4 Preconditioner 관점에서의 분석

4. Lion 논문의 주요 실험 결과

4.1 Image Classification (ImageNet)

4.2 Language Modeling

4.3 Diffusion Models

4.4 논문의 핵심 권장사항

5. 실험: CIFAR-10에서의 비교

5.1 실험 설계

5.2 결과

5.3 시각화

6. 하이퍼파라미터 튜닝 완전 가이드

6.1 Learning Rate 찾기

6.2 Weight Decay 설정

6.3 β 값 조정

6.4 Batch Size와의 관계

6.5 AdamW에서 Lion으로 전환 체크리스트

6.6 튜닝 실패 시 디버깅

7. 실전 적용 가이드

7.1 Lion을 써야 할 때

7.2 Lion을 피해야 할 때

7.3 Gradient Clipping과 함께 사용

7.4 Learning Rate Scheduler와의 조합

8. 메모리 절감의 실제 임팩트

8.1 이론적 계산

8.2 7B LLM 학습 시나리오

8.3 ZeRO와의 시너지

9. 다른 Optimizer들과의 비교

9.1 LAMB (Layer-wise Adaptive Moments)

9.2 Shampoo

9.3 Adafactor

9.4 비교 요약

10. 결론: 국룰은 국룰인 이유가 있다

10.1 언제 Lion을 고려할까?

10.2 실무자의 시선

References

2.4 Lion의 `sign()` 연산