.backward() and the .grad Buffer

How gradients land — and when they accumulate

tensor.backward() is the call that triggers the backward pass. It expects to be called on a scalar; if your tensor isn't scalar, you have to provide an "incoming gradient" argument that tells backward how to weight each output element. In normal supervised learning your loss is always a scalar, so this argument almost never appears in your code.

The .grad accumulator

After backward, the gradient lands in leaf_tensor.grad. Critical fact: repeated backward calls add to the existing .grad, not overwrite it. That's why every training loop has optimizer.zero_grad() at the top — to reset the accumulator.

The accumulation behavior is occasionally useful: it's what makes gradient accumulation (simulating a larger batch by accumulating gradients across several mini-batches before stepping the optimizer) trivial.

Non-scalar outputs

If you must call .backward() on a non-scalar tensor, you have to pass a tensor of the same shape as a "virtual incoming gradient" — this is the vector-Jacobian product (vJp). In practice, you'll see this in research code (computing per-sample gradients, custom losses) more than in normal training.

Code

Scalar backward — the standard case·python

import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = (x ** 2).sum()        # collapse to scalar
y.backward()
print(x.grad)             # tensor([2., 4., 6.])

Accumulation — why zero_grad exists·python

import torch

x = torch.tensor(2.0, requires_grad=True)

(x ** 2).backward()       # dy/dx = 4
print(x.grad)             # tensor(4.)

(x ** 3).backward()       # dy/dx = 12, but accumulates with the 4
print(x.grad)             # tensor(16.) — 4 + 12

# The fix in real training:
x.grad = None             # or .zero_() — both work; None is faster
(x ** 3).backward()
print(x.grad)             # tensor(12.)

Gradient accumulation — bigger batch on a smaller GPU·python

import torch
import torch.nn as nn

model = nn.Linear(10, 2)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()
accumulation_steps = 4

optimizer.zero_grad()
for i in range(8):  # 8 mini-batches → effective batch = 8 * micro-batch
    x = torch.randn(16, 10)
    y = torch.randn(16, 2)
    loss = loss_fn(model(x), y) / accumulation_steps   # scale!
    loss.backward()                                     # accumulates

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

.backward() and the .grad Buffer

How gradients land — and when they accumulate

The .grad accumulator

Non-scalar outputs

Code

External links

Exercise

Progress

Comments 0