vmap + grad: 예제별 그래디언트

vmap과 grad를 합성하는 가장 흔한 패턴은 예제별 그래디언트야. 표준 학습은 배치 평균 그래디언트를 쓰지만, 가끔 예제마다 그래디언트를 따로 보고 싶을 때가 있어 (그래디언트 노이즈 규모, 차등 프라이버시, 영향 함수).

import jax
import jax.numpy as jnp

# 단일 example 의 loss
def loss_one(params, x, y):
    '''x: (D,), y: (), params: (D,)'''
    pred = jnp.dot(x, params)
    return (pred - y) ** 2

# grad — 단일 example 의 gradient
g_one = jax.grad(loss_one)
grad_for_one_example = g_one(params, x_single, y_single)  # shape: (D,)

# vmap — batch 의 모든 example 의 gradient
g_each = jax.vmap(g_one, in_axes=(None, 0, 0))
per_grads = g_each(params, batch_x, batch_y)  # shape: (B, D)

# 비교: 표준 batch loss 의 gradient
def batch_loss(params, X, Y):
    pred = X @ params
    return jnp.mean((pred - Y) ** 2)

batch_grad = jax.grad(batch_loss)(params, batch_x, batch_y)  # shape: (D,)

# 검증: per_grads.mean(0) ≈ batch_grad
print(jnp.allclose(per_grads.mean(0), batch_grad))  # True (거의)

이게 PyTorch에서 훅이나 Functorch 없이는 어려웠던 일. JAX에서는 한 줄이면 돼.

활용 1: 그래디언트 분포 분석

@jax.jit
def per_example_grads(params, X, Y):
    return jax.vmap(jax.grad(loss_one), in_axes=(None, 0, 0))(params, X, Y)

g_each = per_example_grads(params, X_batch, Y_batch)  # (B, D)

# gradient norm 분포
norms = jnp.linalg.norm(g_each, axis=1)  # (B,)
print(f"평균 norm: {norms.mean():.4f}")
print(f"최대 norm: {norms.max():.4f} (outlier 가능성)")
print(f"norm std: {norms.std():.4f}")

# outlier example 식별
outlier_idx = jnp.argmax(norms)
print(f"가장 큰 gradient 의 example: {outlier_idx}, norm={norms[outlier_idx]}")

실전에서 학습이 불안정할 때는 이상치 예제를 찾는 게 첫 디버깅 단계야.

활용 2: Differential Privacy (DP-SGD)

def dp_sgd_step(params, X, Y, lr, clip_norm, noise_scale, key):
    # per-example gradients
    g_each = jax.vmap(jax.grad(loss_one), in_axes=(None, 0, 0))(params, X, Y)

    # 각 example 의 norm clipping
    g_each = jax.vmap(
        lambda g: g * jnp.minimum(1, clip_norm / jnp.linalg.norm(g))
    )(g_each)

    # 평균 + noise
    g_clipped_mean = g_each.mean(0)
    noise = jax.random.normal(key, g_clipped_mean.shape) * noise_scale
    return params - lr * (g_clipped_mean + noise)

학술 알고리즘인 DP-SGD를 30줄 이하로 구현할 수 있어. 다른 프레임워크에선 별도 라이브러리가 필요해.

활용 3: 영향 함수

"이 한 예제를 빼면 모델이 어떻게 바뀔까?"에 대한 근사야. 예제별 그래디언트가 기본 재료야.

# 매우 간단화된 influence
def influence_score(params, X, Y, x_test, y_test):
    g_each = jax.vmap(jax.grad(loss_one), in_axes=(None, 0, 0))(params, X, Y)
    g_test = jax.grad(loss_one)(params, x_test, y_test)
    return -jnp.dot(g_each, g_test)  # train example 마다 한 score