Automatic Differentiation in Practice

Automatic differentiation is exact up to floating-point arithmetic. It is not symbolic algebra. It is not numerical perturbation.

Finite differences have a price you can write down. A gradient in n dimensions needs O(n) function evaluations. A Jacobian of f: R^n → R^m by central differences costs 2mn evaluations. And the estimate is ill-conditioned by construction: shrinking the step size h drives truncation error down while round-off error goes up. Baydin and three colleagues set that out in a 2018 survey in the Journal of Machine Learning Research, in a section titled “AD Is Not Numerical Differentiation”. It states the obstacle without hedging: “The O(n) complexity of numerical differentiation for a gradient in n dimensions is the main obstacle to its usefulness in machine learning, where n can be as large as millions or billions in state-of-the-art deep learning models”.

That is why finite differencing survives in a training stack as a checker rather than as a method. PyTorch's torch.autograd.gradcheck defaults to eps=1e-06, atol=1e-05 and rtol=0.001, and warns that those defaults assume float64. JAX's jax.test_util.check_grads is cheaper still. It probes one random direction rather than all n.

“Exact” carries a qualification that is easy to miss. The derivative is exact for the program, and the program is not the function. Bolte and Pauwels showed this in 2020, with three TensorFlow implementations of the same rectifier. For relu(t) = max(0, t) the framework returns a derivative of 0 at t = 0. For relu2(t) = relu(−t) + t it returns 1. For relu3(t) = (relu(t) + relu2(t))/2 it returns 1/2. Their statement of the reason is blunt: “AD does not act on functions, but on their representations, i.e., on programs. Different programs implementing the same function may provide different results, beyond numerical precision”. The sharpest case in the paper is zero(t) = relu2(t) − relu(t). Its output is constantly 0. Its automatic derivative at the origin is 1.

During tracked execution, the runtime associates values with operation history or stages a graph ahead of execution.

Keeping that history is the whole cost of reverse mode. The survey's “Reverse Mode” section states the shape of the bill: storage requirements grow, in the worst case, in proportion to the number of operations in the evaluated function. The tape is as long as the program.

The standard way to buy that memory back is to recompute instead of storing, and it has been measured twice. Chen and three colleagues stated the trade in 2016: “we design an algorithm that costs O(sqrt(n)) memory to train a n layer network, with only the computational cost of an extra forward pass per mini-batch”. Concretely, they cut the training memory of a 1,000-layer residual network from 48 GB to 7 GB on ImageNet problems, for roughly 30 percent extra running time. Korthikanti and colleagues at NVIDIA reported the same order of penalty at production scale in 2023 — “30 − 40% execution time overhead when full activation recomputation is used”. Sequence parallelism plus selective recomputation removed over 90 percent of that overhead, reaching 54.2% Model FLOPs Utilization on a 530B-parameter GPT-3-style model on 2240 A100 GPUs, against 42.1% with full recomputation.

What gets recorded is a decision with a price, not a framework detail. Differentiating through a computation that already contains a gradient step forces the tape to carry second-order context. Finn and colleagues measured what dropping it costs. Their 2017 first-order approximation of MAML omits those second derivatives, and on few-shot MiniImagenet classification “the performance of this method is nearly the same as that obtained with full second derivatives”. The saving was concrete. The approximation “removes the need for computing Hessian-vector products in an additional backward pass, which we found led to roughly 33% speed-up in network computation”.

The pass they dropped had a known, bounded price of its own. Pearlmutter's R-operator, published in 1994, computes an exact Hessian-vector product without ever forming the Hessian, at a cost his abstract puts as one “which takes about as much computation, and is about as local, as a gradient evaluation”. Martens built Hessian-free optimisation on exactly that result in 2010. He noted that the alternative finite-difference identity Hd = lim(ε→0)[∇f(θ+εd) − ∇f(θ)]/ε costs one extra gradient evaluation but is numerically fragile. So the second-order context is about one gradient's worth of work per pass. Removing it returned 33% of the network computation.

Detaching a tensor or using a stop-gradient operation treats its value as constant for the surrounding derivative, which is a useful thing to do for targets, teacher signals, and selected optimization designs.

Severing the path by accident is not a mysterious event either. The frameworks enumerate the ways. TensorFlow's guide “Introduction to gradients and automatic differentiation” lists four cases where a gradient comes back None. A tf.Variable accidentally replaced by a tf.Tensor (x = x + 1 instead of x.assign_add(1)). A step computed outside TensorFlow, for example np.mean. An integer or string dtype. And a read from a stateful object, such as a Variable after assign_add, a tf.data iterator or a tf.queue. The fourth case states the rule in two sentences: “State stops gradients. When you read from a stateful object, the tape can only observe the current state, not the history that lead to it.” PyTorch states the complementary rule on the recording side. An operation enters the backward graph only if at least one of its inputs requires grad, and only leaf tensors with requires_grad=True receive gradients in .grad.

A stop-gradient can also be the load-bearing part of a method rather than a convenience. SimSiam networks learn useful representations without negative pairs, without large batches and without a momentum encoder. Chen and He reported that in 2021, and named the reason: “collapsing solutions do exist for the loss and structure, but a stop-gradient operation plays an essential role in preventing collapsing”. Remove that one call and nothing raises an error. The training loop runs. What the network learns changes completely.

Many frameworks add new gradients to existing parameter buffers instead of replacing them.

Accumulating the gradients correctly is still not the same as running the full batch, and the gap produced a real, dated defect. On 16 October 2024 Hugging Face published “Fixing Gradient Accumulation”, which opens with the symptom: “Gradient accumulation is supposed to be mathematically equivalent to full batch training; however, losses did not match between training runs where the setting was toggled on and off.” The cause was the denominator rather than the sum. A token-level cross-entropy averaged inside each micro-batch and then summed across micro-batches divides by the wrong count. The fix sums with reduction="sum" and divides by the total non-padding token count, num_items_in_batch. It shipped in transformers PR #34191, opened 16 October and merged 17 October 2024, after Unsloth surfaced the report. On 19 October 2024 the same defect was raised against another trainer, as PyTorch Lightning issue #20350. In neither case did anything raise. The run was simply not the full-batch run its author believed it was.

In eager systems, loops and conditionals can build different graphs for different inputs, and autodiff differentiates the executed path rather than branches that were skipped. That holds only while the program stays eager.

Under JAX's jit, the branch is resolved at trace time. The documentation page “Control flow and logical operators with JIT” says so directly: “Python control flow and logical operators are evaluated at JIT compile time, such that the compiled function represents a single path through the control flow graph”. Branching on a traced value fails outright, because the abstract value cannot be coerced to True or False. Data-dependent branching has to be written with the structured primitives lax.cond and lax.scan. PyTorch's answer is the same shape. torch.cond is a structured control-flow operator, carrying restrictions on its branches and operands precisely so that torch.compile and torch.export can capture it.

Compilation and tracing tools therefore impose restrictions on data-dependent Python behavior. Confirm whether control flow is captured, specialized, or frozen during export. Do not assume the exported program branches the way the eager one did.

A parameter gradient can go missing or look suspicious. Inspect the graph first. Changing the optimizer comes later.

Work the documented causes before the exotic ones. Confirm the parameter took part in the loss-producing path, and that it is a leaf tensor with requires_grad=True: PyTorch admits an operation into the backward graph only when at least one of its inputs requires grad, and only leaf tensors receive gradients in .grad. Then walk TensorFlow's four None-gradient cases as a checklist that transfers between frameworks. A variable accidentally replaced by a tensor (x = x + 1 instead of x.assign_add(1)). A step computed outside the framework, such as np.mean. An integer or string dtype. A read from a stateful object like a tf.data iterator or a tf.queue. Then search for detachments and stop-gradient calls. Then remove in-place operations around saved values. In PyTorch the version counter checked at saved_tensors will tell you when a saved tensor was overwritten; in JAX the assignment would have raised TypeError before it ever ran.

Only then compare numerically, and compare inside the checker's stated limits. torch.autograd.gradcheck defaults to eps=1e-06, atol=1e-05 and rtol=0.001, and those defaults assume float64. jax.test_util.check_grads probes one random direction, not all n. A finite-difference check on a deterministic tiny case is a verdict on that case, at that precision, in that direction.