🎓 A1Learner
🔍

PyTorch Per-Sample Gradients with vmap

A per-sample gradient is the gradient of one batch row's own loss. vmap lifts a single-row function over a whole batch in one call.

The loop gives each row a gradient

The loop clones the weights and runs backward once per row. Four rows mean four backward calls, and the prints stack the four gradients.

import torch

w = torch.tensor([0.5, -1.0, 2.0])
X = torch.tensor([[1.0, 0.0, 2.0],
                  [0.0, 1.0, 1.0],
                  [2.0, 1.0, 0.0],
                  [1.0, 1.0, 1.0]])
y = torch.tensor([1.0, 0.0, 2.0, 1.0])


def one_loss(w, x_row, target):
    pred = (x_row * w).sum()
    return (pred - target) ** 2


rows = []
for i in range(len(X)):
    w_loop = w.clone().requires_grad_(True)
    loss_i = one_loss(w_loop, X[i], y[i])
    loss_i.backward()
    rows.append(w_loop.grad)
loop_grads = torch.stack(rows)
print(loop_grads)

grad makes a gradient function

torch.func.grad wraps one_loss and returns a new function. Calling it on the first row prints that row's gradient alone.

from torch.func import grad

one_grad = grad(one_loss)
print("First row gradient:", one_grad(w, X[0], y[0]).tolist())

vmap runs it over the batch

in_dims (None, 0, 0) keeps the weights fixed and walks both data arguments. The result matches the loop stack exactly.

from torch.func import vmap

all_grads = vmap(grad(one_loss), in_dims=(None, 0, 0))(w, X, y)
print("Shape:", tuple(all_grads.shape))
print("Matches the loop:", torch.equal(all_grads, loop_grads))

Example

import torch
from torch.func import grad, vmap

torch.manual_seed(0)
w = torch.tensor([0.5, -1.0, 2.0])
X = torch.tensor([[1.0, 0.0, 2.0],
                  [0.0, 1.0, 1.0],
                  [2.0, 1.0, 0.0],
                  [1.0, 1.0, 1.0]])
y = torch.tensor([1.0, 0.0, 2.0, 1.0])


def one_loss(w, x_row, target):
    pred = (x_row * w).sum()
    return (pred - target) ** 2


per_sample = vmap(grad(one_loss), in_dims=(None, 0, 0))(w, X, y)
print("Per-sample gradients:")
print(per_sample)
mean_grad = per_sample.mean(dim=0)

w2 = w.clone().requires_grad_(True)
batch_loss = ((X @ w2 - y) ** 2).mean()
batch_loss.backward()
print("Mean of per-sample gradients:", mean_grad.tolist())
print("Batch gradient:", w2.grad.tolist())
print("They match:", torch.allclose(mean_grad, w2.grad))

Key points

  • Sample: one row of a batch, with its input and its target
  • torch.func.grad turns a loss function into a gradient function
  • vmap maps that gradient function over every row at once
  • in_dims names which arguments carry a batch axis and which stay fixed
  • The mean of the per-sample gradients equals the gradient of the mean loss
  • In short, grad plus vmap gives every row its own gradient without a hand loop

In this example, one_loss scores a single row against the shared weight row w. torch.func.grad turns that loss into a function that returns its gradient. vmap maps the gradient function over all four rows in one call. The in_dims argument keeps w fixed and walks X and y row by row. The per-sample table prints one gradient row per batch row. The mean of those rows prints as 0.0, minus 0.25, and 4.25. A plain batch backward on the mean loss prints the same three values. The closing check prints True, so the two routes agree exactly.

Note

torch.func is version-sensitive, and these lessons pin torch 2.14.1+cpu

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does in_dims (None, 0, 0) tell vmap in this lesson?