advanced25 min studyLesson 502 in Engineering

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

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.

u∈Rm,VJP⁡f(x;u)=Jf(x)Tu∈Rnu\in\mathbb R^m,\qquad \operatorname{VJP}_f(x;u)=J_f(x)^\mathsf T u\in\mathbb R^n
f maps n real dimensionless inputs to m real dimensionless outputs, and u is fixed during the product. With physical units, cotangent units depend on the selected scalar objective. Complex-valued conventions are outside this real-transpose statement.↔ Scroll wide equations horizontally.

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.

aˉ+=bzˉ,bˉ+=azˉ(z=ab)\bar a\mathrel{+}=b\bar z,\qquad \bar b\mathrel{+}=a\bar z\quad(z=ab)
Bars denote real cotangents and += denotes accumulation of all downstream uses, not mutation of the primal values. The rule assumes an ordinary differentiable product; required a and b must describe the same forward evaluation.↔ Scroll wide equations horizontally.

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.

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.

  1. The Jacobian is [[4,0],[3,2],[0,6]], a 3×2 matrix. The cotangent seed has three entries to match the output.

  2. Accumulate into x: 4×1+3×2+0×(−1)=10. Accumulate into y: 0×1+2×2+6×(−1)=−2.

  3. 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.

The result

Reverse mode maps the selected output sensitivity to an input-shaped two-vector and accumulates every output path's contribution.

Explore the model

Interactive lab

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.

Inspect the linear map before its adjointThe 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. Current readings: Area scale · det A: 1.4; Singular values · σmax, σmin: 1.597, 0.876; Condition number · κ₂(A): 1.823. All parameters are adjustable with the labeled controls below.Matrix A with first row 1.4, 0.6 and second row zero, one; display interpolates between identity and A while final determinant remains 1.4.−4−2024−2−1012x · equal geometric scaleyA = [ 1.4 0.6 ; 0 1 ]Display blend: 0% · final values below
Interpolated geometryOriginal unit squareSelected final transform
Area scale · det A
1.4
Singular values · σmax, σmin
1.597, 0.876
Condition number · κ₂(A)
1.823
t = 0 s

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.

Check your understanding

For a three-component output, what does an all-ones reverse seed compute at a differentiable real input?

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.

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.