머신러닝을 위한 수학적 최적화: Adam부터 볼록 최적화, ZeRO optimizer까지

목차

1. 최적화 기초: 볼록 최적화와 KKT 조건
2. 경사 하강법 계열 optimizer
3. 2차 최적화 방법
4. 학습률 스케줄링
5. 정규화 기법
6. 손실 함수 설계
7. LLM 학습 최적화
8. 퀴즈

최적화 기초

볼록 최적화 (Convex Optimization)

함수 $f: \mathbb{R}^n \to \mathbb{R}$이 **볼록(convex)**이면, 임의의 두 점 $x, y$와 $\lambda \in [0,1]$에 대해 다음이 성립합니다.

$f(\lambda x + (1-\lambda)y) \leq \lambda f(x) + (1-\lambda)f(y)$

볼록 함수의 핵심 성질:

• 모든 지역 최솟값(local minimum)이 전역 최솟값(global minimum)
• 경사 하강법이 수렴이 보장됨
• 딥러닝 손실 함수는 대부분 비볼록(non-convex)이지만, 볼록 분석 기법이 여전히 유용

강볼록(Strongly Convex): $m > 0$이 존재하여 $f(y) \geq f(x) + \nabla f(x)^T(y-x) + \frac{m}{2}\|y-x\|^2$가 성립하면, 수렴 속도가 선형(linear convergence)으로 빨라집니다.

라그랑주 승수법

등호 제약 최적화 문제를 다룹니다.

$\min_x f(x) \quad \text{subject to} \quad g_i(x) = 0, \; i = 1, \ldots, m$

라그랑주안(Lagrangian):

$\mathcal{L}(x, \lambda) = f(x) + \sum_{i=1}^{m} \lambda_i g_i(x)$

최적해에서 $\nabla_x \mathcal{L} = 0$, $\nabla_\lambda \mathcal{L} = 0$이 성립합니다.

KKT 조건

부등호 제약까지 포함한 일반적인 최적화:

$\min_x f(x) \quad \text{s.t.} \quad g_i(x) \leq 0, \; h_j(x) = 0$

KKT 조건 (필요 조건):

1. 정류성(Stationarity): $\nabla f(x^*) + \sum_i \mu_i \nabla g_i(x^*) + \sum_j \lambda_j \nabla h_j(x^*) = 0$
2. 원시 실현가능성(Primal feasibility): $g_i(x^*) \leq 0$, $h_j(x^*) = 0$
3. 쌍대 실현가능성(Dual feasibility): $\mu_i \geq 0$
4. 상보 여유(Complementary slackness): $\mu_i g_i(x^*) = 0$

볼록 문제에서 KKT 조건은 충분 조건도 됩니다.

안장점 (Saddle Point)

딥러닝 최적화에서 큰 문제는 지역 최솟값보다 안장점입니다. 안장점에서는 일부 방향으로는 함수값이 증가하고 다른 방향으로는 감소하여 경사가 0이 됩니다. SGD의 확률적 노이즈가 안장점 탈출에 도움이 됩니다.

경사 하강법 계열

SGD와 그 변형

기본 SGD:

$\theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L}(\theta_t; x_i, y_i)$

SGD with Momentum:

$v_{t+1} = \beta v_t + \nabla_\theta \mathcal{L}(\theta_t)$ $\theta_{t+1} = \theta_t - \eta v_{t+1}$

모멘텀 $\beta = 0.9$가 보통 사용되며, 이전 기울기 방향을 유지해 oscillation을 줄입니다.

Nesterov Accelerated Gradient (NAG):

$v_{t+1} = \beta v_t + \nabla_\theta \mathcal{L}(\theta_t - \beta v_t)$ $\theta_{t+1} = \theta_t - \eta v_{t+1}$

현재 위치가 아니라 "앞을 내다본" 위치에서 기울기를 계산합니다.

AdaGrad, RMSProp, Adam

AdaGrad: 파라미터별 학습률 적응

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

자주 나타나는 특성은 학습률이 줄어들고, 드물게 나타나는 특성은 학습률이 큽니다. 단점: 학습률이 단조 감소하여 학습이 멈출 수 있습니다.

RMSProp: AdaGrad의 누적 문제 해결

$v_t = \beta v_{t-1} + (1-\beta) g_t^2$ $\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{v_t + \epsilon}} g_t$

Adam (Adaptive Moment Estimation):

$m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t \quad \text{(1차 모멘트)}$ $v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \quad \text{(2차 모멘트)}$

Bias correction:

$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}$

$\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \hat{m}_t$

기본 하이퍼파라미터: $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$

import torch
import torch.optim as optim

model = ...  # 모델 정의

# Adam optimizer
optimizer_adam = optim.Adam(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    eps=1e-8
)

# AdamW optimizer (weight decay 분리)
optimizer_adamw = optim.AdamW(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01  # L2와 독립적으로 적용
)

AdamW와 Lion

AdamW: Adam에서 weight decay를 L2 패널티가 아닌 파라미터 업데이트에 직접 적용합니다.

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

일반 Adam의 L2 정규화와 수학적으로 동등하지 않습니다 (자세한 설명은 퀴즈 참고).

Lion (EvoLved Sign Momentum):

$u_t = \beta_1 m_{t-1} + (1-\beta_1) g_t$ $\theta_{t+1} = \theta_t - \eta \cdot \text{sign}(u_t)$ $m_t = \beta_2 m_{t-1} + (1-\beta_2) g_t$

Lion은 부호(sign)만 사용하므로 메모리 효율적이며, 업데이트 크기가 균일합니다.

2차 최적화

Newton 방법

2차 미분(Hessian)을 활용합니다.

$\theta_{t+1} = \theta_t - H_t^{-1} \nabla f(\theta_t)$

여기서 $H_t = \nabla^2 f(\theta_t)$는 Hessian 행렬입니다. 이차 수렴(quadratic convergence)하지만, $n \times n$ Hessian의 역행렬 계산이 $O(n^3)$으로 딥러닝에서 비현실적입니다.

L-BFGS (Limited-memory BFGS)

Hessian을 직접 저장하지 않고, 최근 $m$개의 기울기 차이로 근사합니다.

$H_t^{-1} \approx \text{두 벡터 시퀀스 } \{s_k\}, \{y_k\}\text{로 근사}$

여기서 $s_k = \theta_{k+1} - \theta_k$, $y_k = \nabla f_{k+1} - \nabla f_k$입니다.

import torch
import torch.optim as optim

# L-BFGS는 클로저(closure) 함수 필요
optimizer = optim.LBFGS(
    model.parameters(),
    lr=1.0,
    max_iter=20,
    history_size=10,
    line_search_fn='strong_wolfe'
)

def closure():
    optimizer.zero_grad()
    output = model(input_data)
    loss = criterion(output, target)
    loss.backward()
    return loss

optimizer.step(closure)

자연 경사 하강법 (Natural Gradient)

Fisher Information Matrix를 사용하여 파라미터 공간의 곡률을 고려합니다.

$\theta_{t+1} = \theta_t - \eta F(\theta_t)^{-1} \nabla \mathcal{L}(\theta_t)$

Fisher Matrix: $F(\theta) = \mathbb{E}\left[\nabla \log p(y|x;\theta) \nabla \log p(y|x;\theta)^T\right]$

K-FAC(Kronecker-factored Approximate Curvature)은 자연 경사법을 실용적으로 구현합니다.

학습률 스케줄링

Warmup

초기에 학습률을 서서히 증가시켜 학습을 안정화합니다.

$\eta_t = \eta_{\max} \cdot \frac{t}{T_{\text{warmup}}} \quad (t \leq T_{\text{warmup}})$

Cosine Annealing

$\eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\frac{\pi t}{T}\right)$

from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR, ReduceLROnPlateau

# Cosine Annealing
scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)

# OneCycleLR: Warmup + Cosine Decay
scheduler = OneCycleLR(
    optimizer,
    max_lr=1e-3,
    total_steps=1000,
    pct_start=0.3,         # 30% warmup
    anneal_strategy='cos'
)

# ReduceLROnPlateau: 검증 손실이 개선되지 않으면 감소
scheduler = ReduceLROnPlateau(
    optimizer,
    mode='min',
    factor=0.5,
    patience=10,
    min_lr=1e-6
)

Cyclical Learning Rate (CLR)

학습률을 주기적으로 변동시켜 안장점 탈출을 돕습니다.

$\eta_t = \eta_{\min} + (\eta_{\max} - \eta_{\min}) \cdot \max\left(0, 1 - \left|\frac{t}{\text{step\_size}} - 2k + 1\right|\right)$

정규화 기법

L1 / L2 정규화

L2 정규화 (Ridge):

$\mathcal{L}_{\text{reg}} = \mathcal{L} + \frac{\lambda}{2} \|\theta\|_2^2$

기울기: $\nabla_\theta \mathcal{L}_{\text{reg}} = \nabla_\theta \mathcal{L} + \lambda \theta$

L1 정규화 (Lasso):

$\mathcal{L}_{\text{reg}} = \mathcal{L} + \lambda \|\theta\|_1$

L1은 희소(sparse) 솔루션을 유도합니다.

Batch Normalization vs Layer Normalization

Batch Normalization (BN):

$\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}} \cdot \gamma + \beta$

여기서 $\mu_B$, $\sigma_B^2$는 미니배치 내 통계량입니다. 배치 방향으로 정규화합니다.

Layer Normalization (LN):

$\hat{x} = \frac{x - \mu_L}{\sqrt{\sigma_L^2 + \epsilon}} \cdot \gamma + \beta$

여기서 통계량은 각 샘플의 피처 차원을 따라 계산됩니다.

Weight Decay vs L2 정규화

SGD에서:

$\theta_{t+1} = \theta_t - \eta(\nabla \mathcal{L} + \lambda \theta_t) = (1 - \eta\lambda)\theta_t - \eta \nabla \mathcal{L}$

이 경우 weight decay와 L2 정규화는 동일합니다. 하지만 Adam에서는:

• L2 Adam: 기울기에 $\lambda \theta$를 더한 후 적응형 학습률 적용 → 적응 계수로 나누어져 정규화 효과 약화
• AdamW: 파라미터 업데이트 후 $\lambda \theta$를 직접 빼냄 → 모든 파라미터에 균등한 weight decay

# Dropout
import torch.nn as nn

class RegularizedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(512, 256)
        self.bn1 = nn.BatchNorm1d(256)
        self.ln1 = nn.LayerNorm(256)
        self.dropout = nn.Dropout(p=0.3)

    def forward(self, x):
        x = self.fc1(x)
        x = self.bn1(x)   # 또는 self.ln1(x)
        x = torch.relu(x)
        x = self.dropout(x)
        return x

손실 함수 설계

Cross-Entropy Loss

$\mathcal{L}_{CE} = -\sum_{c=1}^{C} y_c \log \hat{p}_c$

이진 분류: $\mathcal{L}_{BCE} = -[y \log p + (1-y)\log(1-p)]$

Focal Loss

클래스 불균형 문제를 해결합니다. 쉬운 샘플의 기여를 줄입니다.

$\mathcal{L}_{FL} = -(1-p_t)^\gamma \log(p_t)$

여기서 $p_t$는 정답 클래스의 예측 확률, $\gamma \geq 0$는 focusing parameter입니다. $\gamma = 0$이면 일반 Cross-Entropy와 동일합니다.

import torch
import torch.nn.functional as F

class FocalLoss(torch.nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, logits, targets):
        bce_loss = F.binary_cross_entropy_with_logits(
            logits, targets.float(), reduction='none'
        )
        p = torch.sigmoid(logits)
        p_t = p * targets + (1 - p) * (1 - targets)
        focal_weight = (1 - p_t) ** self.gamma
        alpha_t = self.alpha * targets + (1 - self.alpha) * (1 - targets)
        loss = alpha_t * focal_weight * bce_loss
        return loss.mean()

Contrastive Loss와 Triplet Loss

Contrastive Loss (Siamese Network):

$\mathcal{L} = (1-y)\frac{d^2}{2} + y \cdot \max(0, m - d)^2$

여기서 $d = \|f(x_1) - f(x_2)\|_2$, $y=0$이면 유사 쌍, $y=1$이면 비유사 쌍입니다.

Triplet Loss:

$\mathcal{L}_{trip} = \max(0, \|f(a) - f(p)\|_2^2 - \|f(a) - f(n)\|_2^2 + m)$

앵커(a), 포지티브(p), 네거티브(n) 샘플을 사용합니다.

InfoNCE Loss (NT-Xent)

대조 학습(Contrastive Learning)의 핵심 손실 함수입니다.

$\mathcal{L}_{InfoNCE} = -\log \frac{\exp(\text{sim}(z_i, z_j)/\tau)}{\sum_{k=1}^{2N} \mathbf{1}_{k \neq i} \exp(\text{sim}(z_i, z_k)/\tau)}$

여기서 $\tau$는 temperature parameter, $\text{sim}$은 코사인 유사도입니다.

import torch
import torch.nn.functional as F

def info_nce_loss(features, temperature=0.07):
    """
    features: (2N, D) - 각 이미지의 두 augmentation view
    """
    N = features.shape[0] // 2
    features = F.normalize(features, dim=1)

    # 유사도 행렬 계산
    similarity = torch.matmul(features, features.T) / temperature

    # 자기 자신 제거 (대각선을 -inf로)
    mask = torch.eye(2 * N, dtype=torch.bool, device=features.device)
    similarity.masked_fill_(mask, float('-inf'))

    # 포지티브 쌍: i와 i+N, i+N과 i
    labels = torch.cat([
        torch.arange(N, 2 * N),
        torch.arange(N)
    ]).to(features.device)

    loss = F.cross_entropy(similarity, labels)
    return loss

LLM 학습 최적화

Gradient Clipping

기울기 폭발(exploding gradient)을 방지합니다.

$g \leftarrow g \cdot \min\left(1, \frac{\text{clip\_norm}}{\|g\|_2}\right)$

import torch

def train_with_clipping(model, optimizer, loss, max_norm=1.0):
    optimizer.zero_grad()
    loss.backward()

    # 기울기 노름 모니터링
    total_norm = 0
    for p in model.parameters():
        if p.grad is not None:
            param_norm = p.grad.data.norm(2)
            total_norm += param_norm.item() ** 2
    total_norm = total_norm ** 0.5

    # 클리핑 적용
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_norm)
    optimizer.step()
    return total_norm

ZeRO Optimizer (Zero Redundancy Optimizer)

모델 학습 시 메모리를 3단계로 최적화합니다.

혼합 정밀도(mixed precision) + ZeRO-3로 수십억 파라미터 모델을 단일 노드에서 학습 가능합니다.

8-bit Adam

양자화(quantization)를 통해 optimizer 상태를 FP32 대신 INT8로 저장합니다.

• Optimizer 상태 메모리를 75% 절감 (FP32 대비)
• Block-wise quantization으로 정밀도 손실 최소화
• bitsandbytes 라이브러리로 구현 가능

# bitsandbytes 8-bit Adam
import bitsandbytes as bnb

optimizer = bnb.optim.Adam8bit(
    model.parameters(),
    lr=1e-4,
    betas=(0.9, 0.999)
)

Adafactor

Adam의 2차 모멘트를 행렬 분해로 근사합니다.

$V_t \approx \hat{r}_t \hat{v}_t^T \quad (\text{rank-1 근사})$

파라미터 크기에 비례한 메모리만 사용 (행 + 열 벡터). T5, PaLM 등 초대형 모델 학습에 사용됩니다.

퀴즈