고차 그래디언트와 Jacobians

그래디언트 자체를 다시 미분하고 싶을 때도 JAX 변환을 자유롭게 합성할 수 있어. grad(grad(f))가 그냥 작동해.

import jax
import jax.numpy as jnp

def f(x):
    return x ** 4

# 1차 미분: 4x³
print(jax.grad(f)(2.0))  # 32.0

# 2차 미분: 12x²
print(jax.grad(jax.grad(f))(2.0))  # 48.0

# 3차 미분: 24x
print(jax.grad(jax.grad(jax.grad(f)))(2.0))  # 48.0

실용 예, Newton's 메서드 (convergence 빠름, hessian 사용):

def f(x):
    return x ** 3 - 5 * x + 2

f_prime = jax.grad(f)
f_double_prime = jax.grad(f_prime)

x = 0.5
for _ in range(10):
    x = x - f_prime(x) / f_double_prime(x)
print(f"근사근: {x}")

Jacobian, 벡터 입력/출력의 미분

스칼라 함수에 grad. 벡터 함수엔 Jacobian:

def g(x):  # R^3 → R^2
    return jnp.array([x[0] * x[1], x[1] ** 2 + x[2]])

x = jnp.array([1.0, 2.0, 3.0])

# Jacobian: 2x3 행렬, J[i,j] = ∂g_i / ∂x_j
J = jax.jacrev(g)(x)
print(J)
# [[2., 1., 0.],   ∂(x0*x1)/∂x = [x1, x0, 0]
#  [0., 4., 1.]]   ∂(x1²+x2)/∂x = [0, 2*x1, 1]

jacrev vs jacfwd, 역방향 모드와 순방향 모드.

jacrev: 입력 차원 ≫ 출력 차원일 때 효율적이야. (예: 학습에서 매개변수 1M, 손실 1개)
jacfwd: 출력 차원 ≫ 입력 차원일 때 효율적이야. (예: 입력 3차원에서 출력 100차원)

스칼라 손실의 그래디언트는 사실 jacrev(loss)와 같아. jax.grad는 스칼라 출력 일 때 jacrev를 호출하면서 squeeze.

Hessian, 2차 미분 행렬

def f(x):
    return x[0] ** 2 + x[1] ** 2 + x[0] * x[1]

x = jnp.array([1.0, 2.0])

# Hessian = grad of grad, 또는 jacobian of grad
H = jax.hessian(f)(x)
print(H)
# [[2., 1.],
#  [1., 2.]]

# 동등 표현
H_alt = jax.jacrev(jax.grad(f))(x)
H_alt2 = jax.jacfwd(jax.grad(f))(x)

🧮 순방향 모드와 역방향 모드의 차이

자동 미분에는 순방향 모드와 역방향 모드가 있어. 순방향 모드는 입력 방향 하나마다 한 번씩 계산하므로 입력 차원 N에 비례하고, 역방향 모드는 출력 방향 하나마다 한 번씩 계산하므로 출력 차원 M에 비례해. 그래서 출력이 하나인 학습 손실에는 역방향 모드가, 입력은 적고 출력이 많은 문제에는 jacfwd가 유리해. 이 차이를 알면 큰 모델의 역문제와 민감도 분석에서 계산 방식을 고르기 쉬워.

Hessian은 N×N 행렬이므로 큰 모델에서는 직접 계산하지 않아. 대신 jax.jvp나 jax.vjp로 Hessian-벡터 곱을 효율적으로 계산해. 2차 최적화 (LBFGS, K-FAC, 자연 그래디언트)에 등장해.

Code

import jax
import jax.numpy as jnp

def f(x):
    return jnp.sin(x)

df = jax.grad(f)         # cos(x)
d2f = jax.grad(df)       # -sin(x)
d3f = jax.grad(d2f)      # -cos(x)
d4f = jax.grad(d3f)      # sin(x)

x = jnp.array(jnp.pi / 4)
print(f"f(x)   = {f(x):.4f}")     # 0.7071 (sin)
print(f"f'(x)  = {df(x):.4f}")    # 0.7071 (cos)
print(f"f''(x) = {d2f(x):.4f}")   # -0.7071 (-sin)
print(f"f'''(x)= {d3f(x):.4f}")   # -0.7071 (-cos)

import jax
import jax.numpy as jnp

def f(x):
    return x[0] ** 2 * x[1] + x[1] ** 3

# Hessian: matrix of second partial derivatives
hessian_fn = jax.hessian(f)
x = jnp.array([1.0, 2.0])
H = hessian_fn(x)
print(H)
# [[ 4.  2.]   d^2f/dx0^2 = 2*x1 = 4,  d^2f/dx0dx1 = 2*x0 = 2
#  [ 2. 12.]]  d^2f/dx1dx0 = 2*x0 = 2,  d^2f/dx1^2  = 6*x1 = 12

import jax
import jax.numpy as jnp

def vector_fn(x):
    """Maps R^3 -> R^2."""
    return jnp.array([x[0] * x[1], x[1] ** 2 + x[2]])

x = jnp.array([1.0, 2.0, 3.0])

# jacrev: Reverse-mode Jacobian (efficient when output dim < input dim)
J_rev = jax.jacrev(vector_fn)(x)
print(J_rev)
# [[2. 1. 0.]    dy0/dx0=x1, dy0/dx1=x0, dy0/dx2=0
#  [0. 4. 1.]]   dy1/dx0=0,  dy1/dx1=2*x1, dy1/dx2=1
print(J_rev.shape)  # (2, 3)

# jacfwd: Forward-mode Jacobian (efficient when input dim < output dim)
J_fwd = jax.jacfwd(vector_fn)(x)
print(jnp.allclose(J_rev, J_fwd))  # True — same result, different algorithm

🧮 순방향 모드와 역방향 모드의 차이

Code

External links

Exercise

Progress

댓글 0