Skip to content

Higher-Order Gradients & Jacobians

Overview

First-order training (loss.backward()) is the 99% case. But meta-learning, adversarial robustness, scientific ML (PDE solvers), and deep equilibrium models all need Jacobians, Hessians, and gradient-of-gradient. PyTorch gives you two tools: autograd with create_graph=True, and the torch.func (functorch) functional API.

  • torch.autograd.grad(loss, params, create_graph=True) — first forward, second derivative via chains.
  • torch.func.jacrev/jacfwd/jvp/vjp/hessian — explicit Jacobians for any callable.
  • Hessian-vector products (HVP) are cheap (jvp(vjp(...))) — never build the full Hessian.
  • Double backward requires the graph to be kept: retain_graph=True or recompute.

💡 Rule: for anything requiring a full Jacobian, use torch.func/vmap; for HVP/GN updates, compose jvp and vjp symbolically.

autograd.grad with create_graph

import torch

x = torch.tensor([2.0, 3.0], requires_grad=True)
loss = (x ** 3).sum()

g1 = torch.autograd.grad(loss, x, create_graph=True)[0]   # 3x^2
g2 = torch.autograd.grad(g1.sum(), x)[0]                   # 6x
print("first:", g1.tolist(), " second:", g2.tolist())

retain_graph vs create_graph

  • retain_graph=True → reuse the same graph for another backward (memory heavy).
  • create_graph=True → build another graph that will itself be backwarded later.
  • Prefer autograd.grad for selective params; backward(create_graph=True) for everything.

Vector-Jacobian Products (VJP) — the natural primitive

vjp(f, inputs, v) computes v^T J — exactly what autograd does, exposed directly.

from torch.func import vjp, jvp

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

x = torch.tensor([1.0, 2.0], requires_grad=True)
v = torch.tensor([1.0, 1.0])
out, vjp_fn = vjp(f, x)
grad = vjp_fn(v)[0]
print("vjp (grad):", grad)   # sum of cos

Full Jacobians — jacrev / jacfwd

from torch.func import jacrev, jacfwd, vmap

def vec_fun(x):
    return torch.stack([x[0] ** 2, x[0] * x[1]])

x = torch.tensor([2.0, 3.0])
J = jacrev(vec_fun)(x)          # 2x2
print(J)
# [[4., 0.], [3., 2.]]
API Cost Use when
jacrev output-dim backwards output small, input large
jacfwd input-dim forwards input small, output large
jacrev(jacrev(...)) Hessian dims tiny

Hessian-Vector Products Without the Hessian

from torch.func import jvp, vjp

def hessian_vector_product(f, x, v):
    _, vjp_fn = vjp(f, x)
    _, jvp_fn = jvp(lambda z: vjp_fn(z)[0], (x,), (v,))
    return jvp_fn(x)[0]  # careful: use correct composition

def g(z):
    return (z ** 3).sum()

x = torch.tensor([2.0])
Hv = hessian_vector_product(g, x, torch.tensor([1.0]))
print("Hv (should be 6x*1 = 12):", Hv)

💡 Cost of HVP ≈ 2 forwards — this is why second-order optimizers (Newton/L-BFGS/K-FAC-ish) are feasible at scale.

vmap — Batch Without a For-Loop

from torch.func import vmap

def single(x):
    return (x @ torch.ones(3)).sum()

xs = torch.randn(10, 3)
outs = vmap(single)(xs)             # implicit batch dim
print(outs.shape)

Use case: per-sample gradients, per-sample losses, Jacobians-of-batches.

params = [torch.randn(3, 3, requires_grad=True)]
def loss_fn(p, x):
    return (x @ p).square().mean()

# per-sample grads in one call:
batch_grads = vmap(lambda x_: torch.autograd.grad(loss_fn(params[0], x_), params[0])[0])(xs)
print("per-sample grads shape:", batch_grads.shape)  # (10, 3, 3)

Practical Guards

import torch

def safe_second_order(f, x, v=None):
    """Stable Hessian-vector product with NaNs guarded."""
    g1 = torch.autograd.grad(f(x).sum(), x, create_graph=True)[0]
    if v is None:
        v = torch.ones_like(x)
    g2 = torch.autograd.grad((g1 * v).sum(), x)[0]
    assert torch.isfinite(g2).all()
    return g2

Decision Table

Need Tool
Gradient of one scalar wrt params .backward() / autograd.grad
Second derivative for meta-loss create_graph=True double-backward
Full Jacobian (small dims) jacrev/jacfwd
Batch of small problems vmap
HVP for implicit/Hessian-based optim jvp ∘ vjp
Per-sample grads (fisher, robust) vmap over autograd.grad

Key Takeaways

  • create_graph=True extends the graph for a second backward.
  • jvp/vjp are the atoms; Jacobians and Hessians are compositions of them.
  • HVP ≈ 2 forwards; never materialize a full Hessian.
  • vmap replaces batch loops and unlocks per-sample methods.
  • Verify with torch.autograd.gradcheck on custom fns before trusting results.