torch.func: 샘플별 기울기와 Vmap

함수형 autograd와 자동 배치 변환

torch.func는 예전의 독립 패키지 functorch를 PyTorch 안으로 통합한 기능으로, JAX와 비슷한 함수 변환을 제공해. grad, vmap, jacrev, hessian을 사용할 수 있고, 단일 샘플에 동작하는 함수를 배치 전체에 자동으로 벡터화할 수 있어.

실전에서 의외로 자주 만나는 용도는 두 가지야:

샘플별 기울기. 표준 .backward()는 배치에서 합친 손실의 기울기만 줘. 즉, 매개변수마다 기울기가 하나씩 나와. 차등 개인정보 보호, 영향 함수, GradSAM 같은 연구에서는 각 샘플의 손실에 대한 기울기가 필요해. torch.func.vmap(grad(...))를 쓰면 Python 반복문 없이 구할 수 있어.
고차 기울기. 헤시안-벡터 곱, 2차 최적화, 메타 학습은 모두 기울기를 다시 미분해야 해. torch.func.grad(grad(f))처럼 변환을 깔끔하게 조합할 수 있어.

관점을 바꿔 보기

표준 PyTorch에서는 텐서가 암묵적인 그래프 상태를 들고 있고 여기에 .backward()를 호출해. torch.func에서는 입력과 매개변수를 받는 순수 함수를 만들고, 변환이 또 다른 순수 함수를 만들어 내. JAX에 조금 더 가까운 관점이라 일반적인 학습에서는 다소 낯설 수 있지만, 앞의 두 경우에는 훨씬 강력해.

Code

함수 변환으로 기울기 계산하기·python

import torch
from torch.func import grad

def f(x):
    return torch.sin(x) * x

# grad(f) is a NEW function: x → df/dx
df_dx = grad(f)

x = torch.tensor(1.0)
print(df_dx(x))            # cos(1)*1 + sin(1) ≈ 0.541 + 0.841 = 1.381

# Higher order — grad(grad(f))
d2f_dx2 = grad(grad(f))
print(d2f_dx2(x))          # -sin(1)*1 + 2*cos(1) ≈ ...

vmap: 단일 샘플 함수 자동-배치·python

import torch
from torch.func import vmap, grad

def loss_per_example(w, x, y):
    pred = (w * x).sum()
    return (pred - y) ** 2

w = torch.randn(4)
batch_x = torch.randn(8, 4)         # batch of 8
batch_y = torch.randn(8)

# Single-sample gradient: grad w.r.t. w
single_grad = grad(loss_per_example)

# Vectorize over the batch dim of x and y (in_dims=(None, 0, 0))
per_sample_grads = vmap(single_grad, in_dims=(None, 0, 0))(w, batch_x, batch_y)
print(per_sample_grads.shape)  # torch.Size([8, 4]) — one gradient per sample

nn.Module의 샘플별 기울기를 구하는 현대적인 방법·python

import torch
import torch.nn as nn
from torch.func import functional_call, vmap, grad

model = nn.Linear(4, 2)
params = dict(model.named_parameters())

def compute_loss(params, x, y):
    pred = functional_call(model, params, (x.unsqueeze(0),)).squeeze(0)
    return ((pred - y) ** 2).mean()

batch_x = torch.randn(16, 4)
batch_y = torch.randn(16, 2)

# Per-sample gradient w.r.t. params
per_sample_grad = vmap(grad(compute_loss), in_dims=(None, 0, 0))(
    params, batch_x, batch_y
)
print({k: v.shape for k, v in per_sample_grad.items()})
# {'weight': torch.Size([16, 2, 4]), 'bias': torch.Size([16, 2])}

Exercise

2계층 MLP를 하나 골라 32개 샘플로 이뤄진 배치에서 torch.func.vmap(grad(...))로 샘플별 기울기를 계산해 봐. Python 반복문으로 샘플을 하나씩 처리한 결과와 일치하는지 확인하고 두 방법의 실행 시간도 재 봐. vmap 쪽이 몇 배 더 빠른지 기록해 둬.

함수형 autograd와 자동 배치 변환

관점을 바꿔 보기

Code

External links

Exercise

Progress

댓글 0