Tracking gradients

Automatic Differentiation

Hand-deriving gradients for a 100-million-parameter network is a non-starter. Every modern framework ships an automatic differentiation engine that:

Records each operation onto a computational graph.
Walks the graph in reverse to apply the chain rule.
Returns the gradient with respect to every input you asked about — typically the model parameters.

This chapter teaches the API; the rest of the book leans on it.

A worked example

We’ll differentiate

y = 2\,\mathbf{x}^\top \mathbf{x}

with respect to the column vector \mathbf{x}. The analytic gradient is \nabla_\mathbf{x} y = 4\mathbf{x} — a useful sanity-check target.

from mxnet import autograd, np, npx
npx.set_np()

x = np.arange(4.0)
x

We tell the framework to track operations on x and reserve a slot for its gradient:

# We allocate memory for a tensor's gradient by invoking `attach_grad`
x.attach_grad()
# After we calculate a gradient taken with respect to `x`, we will be able to
# access it via the `grad` attribute, whose values are initialized with 0s
x.grad

Then run the forward pass — y is built from x, so the engine records the dependency:

# Our code is inside an `autograd.record` scope to build the computational
# graph
with autograd.record():
    y = 2 * np.dot(x, x)
y

Backward pass

A single call walks the recorded graph backwards:

y.backward()
x.grad

The result lands in x.grad. Compare with the analytic answer, 4\mathbf{x}:

x.grad == 4 * x

Resetting & re-using

Gradients accumulate by default — call .zero_() (or its equivalent) before computing a fresh gradient:

with autograd.record():
    y = x.sum()
y.backward()
x.grad  # Overwritten by the newly calculated gradient

For non-scalar y, the engine sums up gradients computed for each output element (or you supply weights):

with autograd.record():
    y = x * x  
y.backward()
x.grad  # Equals the gradient of y = sum(x * x)

Detaching from the graph

Sometimes we want a value treated as a constant in the backward pass — e.g., the auxiliary u below should not propagate gradients into x:

with autograd.record():
    y = x * x
    u = y.detach()
    z = u * x
z.backward()
x.grad == u

After detach() (or stop_gradient / lax.stop_gradient), the gradient flows around the detached tensor, not through it:

y.backward()
x.grad == 2 * x

Gradients through control flow

Autograd doesn’t care about Python ifs and whiles — it records whichever ops actually executed. Here’s a function whose behavior depends on its input:

def f(a):
    b = a * 2
    while np.linalg.norm(b) < 1000:
        b = b * 2
    if b.sum() > 0:
        c = b
    else:
        c = 100 * b
    return c

The number of while iterations and the branch taken both depend on the value of a.

…it just works

Run the function on a random scalar and ask for the gradient:

a = np.random.normal()
a.attach_grad()
with autograd.record():
    d = f(a)
d.backward()

The gradient is correct even though the path through the function is data-dependent. Here f(a) ends up linear in a along whichever branch ran, so f'(a) = f(a) / a:

a.grad == d / a

Recap

Mark inputs as needing gradients.
Run the forward pass — the engine records ops.
backward() (or grad()) walks the graph in reverse via the chain rule.
Gradients accumulate; reset between iterations.
detach / stop_gradient to break the graph.
Works through arbitrary Python control flow.