A VJP pulls an output sensitivity back to the inputs
Calculate cotangent shapes, seed a non-scalar output, accumulate fan-out contributions, and distinguish a reverse product from a full Jacobian.
On this page
The question for today
What you’ll understand
- Compute Jᵀu with an output-shaped real cotangent seed.
- Explain why contributions add at a reused intermediate value.
- Choose reverse mode for required output sensitivities without claiming a full Jacobian is free.
A scalar training loss can depend on millions of parameters. Reverse mode starts from a sensitivity at the output and propagates it backward through local derivative rules. For a vector output, the caller must say which output sensitivity matters by supplying a cotangent seed. Following that seed makes both the matrix orientation and the runtime's need for saved forward values much clearer than treating backward as an unexplained gradient button.
Choose an output seed before asking for a gradient
For a real map f:R^n→R^m, a cotangent u has m components and the reverse product J_f(x)ᵀu has n. This is often called a vector-Jacobian product when written with a row seed uᵀJ; here column notation makes the transpose explicit. For a scalar loss, seed one yields its input gradient. For a vector output, choosing all ones gives the gradient of the sum of outputs, not every Jacobian row separately. The seed's meaning belongs in the API contract and the mathematical question.
The reverse seed is held fixed for this product. A loss L=uᵀf(x) has gradient Jᵀu when u is constant; if u depends on x, a total derivative includes the additional dependency unless it is intentionally detached. This lesson uses real arrays and Euclidean inner products. Complex differentiation requires conventions about conjugation and real objectives that cannot be inferred by swapping an ordinary transpose into these formulas. State dtype, shape, and seed structure before comparing two framework results that may implement different derivative conventions.
Reverse traversal adds every use's contribution
If an intermediate feeds two downstream operations, both paths contribute to its cotangent. Reverse accumulation sums them; overwriting one with the other silently loses a chain-rule term. For z=a+b, the output cotangent passes to both inputs. For z=ab, each input receives the output cotangent multiplied by the other input's saved primal. This local reverse rule explains why forward values or equivalent recomputation are needed later. A graph's fan-out is mathematical reuse even when the two uses share storage or appear in different lines of source code.
One reverse product can be efficient for a scalar objective with many inputs, commonly costing a small multiple of the underlying program under suitable primitive assumptions. It still needs storage or recomputation and is not a constant-memory guarantee. Getting a full m×n Jacobian generally needs multiple seeds or batched transformations; batching can change memory and kernel costs. PyTorch 2.9's autograd.grad accepts output-gradient seeds and returns the requested input products, while its accumulation-oriented backward interface has different destination behavior. Do not confuse where gradients are stored with which derivative product was computed.
Use a linear map to check the transpose orientation
On a small example, write the Jacobian dimensions and compute one JVP and VJP by hand. Their inner products should satisfy uᵀ(Jv)=vᵀ(Jᵀu) for compatible real vectors. This adjoint check can catch orientation and seed-shape mistakes, but mutually consistent wrong rules can pass, so it does not replace independent derivative checks. Test accumulation through repeated uses and verify that a zero output seed produces zero input sensitivity within the valid domain. Numerical tolerance should reflect magnitude and floating-point behavior rather than hiding shape errors.
The matrix lab shows one fixed linear family. Reset, set BOTH diagonal scale a=1.5 and shear b=0.5, inspect A=[[1.5,0.5],[0,1]], then Play or Step. On paper, v=(1,−1) gives Av=(1,−1), while u=(2,1) gives Aᵀu=(3,2); both inner products equal one. Edits retain t; playback interpolates geometry, while readouts describe the selected final A. The instrument does not compute VJPs, transpose an arbitrary program, or turn the intermediate animation into a changing derivative rule.
Work it through
Pull three output sensitivities back to two inputs
Use f(x,y)=(x²,xy,y²) at (2,3) with constant output cotangent u=(1,2,−1). All values are real and dimensionless.
The Jacobian is [[4,0],[3,2],[0,6]], a 3×2 matrix. The cotangent seed has three entries to match the output.
Accumulate into x: 4×1+3×2+0×(−1)=10. Accumulate into y: 0×1+2×2+6×(−1)=−2.
The VJP is (10,−2). It is also the gradient of x²+2xy−y² at this input with u fixed, not a three-component gradient or the full six-entry Jacobian.
Reverse mode maps the selected output sensitivity to an input-shaped two-vector and accumulates every output path's contribution.
The living notebook
Explore the model
Interactive lab · matrix
Inspect the linear map before its adjoint
Reset, set BOTH a=1.5 and b=0.5, inspect selected A=[[1.5,0.5],[0,1]], then Play or Step. Compute Av and Aᵀu on paper. Edits retain t; animation interpolates geometry while readouts describe final A. This instrument does not calculate arbitrary VJPs or display a program's reverse tape.
- Area scale · det A
- 1.4
- Singular values · σmax, σmin
- 1.597, 0.876
- Condition number · κ₂(A)
- 1.823
The selected matrix is A = [[a,b],[0,1]]. Playback interpolates geometry between the identity and A; readouts always describe A. Its determinant is a and its singular values come from AᵀA. Positive a makes every supported matrix invertible. This two-dimensional family is not a general eigensolver or a learned network. Parameter changes pause playback and recompute the model at the current elapsed time. Reset restores both parameters and returns time to zero. Playback starts paused. Adjust either parameter, then choose Play or Step.
Look forA VJP uses the adjoint of the selected local linear map and returns an input-shaped sensitivity.
Pause & reason
Check your understanding
Your turn
Try it yourself
For f(x,y)=(x+y,x²y) at (1,2), compute the VJP for constant output seed u=(3,−1). Compare its shape with the JVP input seed shape.
Take a moment to make a prediction or work it out on paper. Open a hint when you need one.
Give me a hint
Transpose Jacobian [[1,1],[4,1]] and multiply by the output seed.
Compare with the solution
Jᵀu=(3−4,3−1)=(−1,2), an input-shaped two-vector. Here both input and output happen to have dimension two, but their seed roles differ; matching sizes do not make Jv and Jᵀu the same operation.
Keep the idea
Take it with you
- A VJP maps a fixed output cotangent to an input sensitivity.
- Reverse accumulation adds contributions from every downstream use.
- Efficient scalar-loss gradients do not make full Jacobians or saved-state costs disappear.
Read further
Original sources and deeper explanations to follow at your own pace.
Optional progress, saved in this browser when available. Revisit any lesson at any time.
Lesson not marked complete.