AI Glossary / Probabilistic machine learning / Intermediate

Vector-Jacobian product

Definition

A vector-Jacobian product, or VJP, combines sensitivities at a function's outputs with its derivatives to produce sensitivities at its inputs. With column-vector notation, it is the Jacobian transpose multiplied by an output-sized vector.

What it means

A function may take thousands of inputs or parameters and return several outputs. The Jacobian records how every output changes with every input. Often, a calculation does not need that entire matrix; it needs the effect of a particular weighted combination of outputs.

A VJP carries those output weights backward through the computation. This is the central operation in reverse-mode automatic differentiation and backpropagation. For a scalar loss, it allows a gradient with respect to many parameters to be obtained without explicitly storing the full Jacobian.

The word “vector” does not imply a list of model predictions. Its role is an incoming sensitivity, sometimes called a cotangent. Keeping that role distinct from an input perturbation prevents confusion with a Jacobian-vector product.

How it works

Let f map n real inputs to m real outputs, and let J be its m-by-n Jacobian at a specified input x. For an m-dimensional output sensitivity u, calculate J transpose times u, producing n input sensitivities. Written with row vectors, the equivalent operation is u transpose times J.

For fixed u, this equals the gradient of the scalar function u transpose f(x). If the weights themselves depend on x, differentiating their full product also includes derivatives of those weights; a VJP alone treats the supplied sensitivity as the incoming coefficient.

Automatic differentiation composes local VJPs in reverse order using the chain rule. It can avoid materializing the full Jacobian, although it often needs intermediate values from the forward computation. Memory-management strategies such as checkpointing trade recomputation for storage.

Three practical examples

  • Neural-network training propagates a scalar loss's output sensitivity back to millions of model parameters.
  • A sensitivity-analysis calculation asks how a weighted combination of several outputs changes when inputs vary.
  • An optimization system differentiates through a sequence of numerical operations without constructing every intermediate Jacobian matrix.

Think of it this way

Think of tracing a total bill back through the items that contributed to it. Each output's contribution is weighted, and those weights are passed backward to identify how changing earlier quantities would affect the total.

The analogy explains reverse accumulation. Derivatives are local mathematical sensitivities, not necessarily causal effects, monetary values or reliable predictions for large changes.

Worked example

Pulling three output sensitivities back to two inputs

Define f(x1,x2) = (x1 squared + x2, x1x2, x1 − x2). At x = (2,3), the output is (7,6,−1). The Jacobian has rows (4,1), (3,2) and (1,−1).

Choose output sensitivity u = (2,−1,0.5). The VJP is J transpose u = (4 times 2 + 3 times −1 + 1 times 0.5, 1 times 2 + 2 times −1 − 1 times 0.5), which is (5.5,−0.5).

This is the gradient of L(x) = 2f1(x) − f2(x) + 0.5f3(x), with the weights held fixed. At the chosen input, L = 7.5.

For input direction v = (1,−2), the predicted directional derivative is (5.5,−0.5) dot (1,−2) = 6.5. A step of 0.01 gives the first-order prediction L approximately 7.565. Direct substitution at (2.01,2.98) gives 7.5654; the small difference is due to the nonlinear second-order term.

Equivalently, Jv = (2,−1,3), and u dot Jv also equals 6.5. This equality checks consistency between forward and reverse derivative products without treating the two products as the same vector.

Illustrative case study

A fictional simulation team optimizes a design using a weighted score made from energy consumption, vibration and mass. Its first derivative implementation constructs a full output-by-parameter Jacobian even though the optimizer only needs the score gradient.

The team replaces that calculation with a VJP seeded by the score's output sensitivities. It checks a small model against explicit derivatives and directional finite differences, then measures memory and runtime for the larger simulation.

The team separately checks that the score weights and units express the intended design trade-off. Faster differentiation cannot correct a poorly chosen objective or establish that simulated sensitivities match physical behavior.

Common misconceptions and mistakes

  • A VJP and JVP have interchangeable input vectors: the VJP takes output-sized sensitivities; the JVP takes input-sized perturbations.
  • Reverse mode must construct the full Jacobian first: it can compose local products directly.
  • Automatic differentiation removes all numerical issues: unstable operations, nondifferentiable points and floating-point effects still matter.
  • A derivative measures a causal intervention: it describes the specified mathematical computation, whose relationship to the world requires separate justification.

Business use

VJPs help make gradient-based training and optimization practical when the number of parameters is large and the final objective has few outputs. They also support targeted sensitivity questions without producing a huge derivative table.

When reviewing an implementation, verify dimensions, array reductions and the meaning of the incoming sensitivity. If a loss already depends nonlinearly on the outputs, use its actual output gradient as the VJP seed. Compare derivative checks at representative inputs and near boundaries. The choice between forward and reverse mode should consider both dimensions and memory cost.

People also ask

Why is it called vector-Jacobian when the formula uses a transpose?

Row-vector and column-vector conventions write the same operation differently. State the shapes to remove ambiguity.

Does one VJP give the full Jacobian?

Generally no. It gives one weighted combination of Jacobian rows; multiple suitable seeds can recover the full matrix.

What happens for a scalar output?

Seeding with one gives the ordinary gradient with respect to the inputs.

Check your understanding

A function has four inputs and seven outputs. What are the sizes of the VJP seed and result?

Answer: The seed has seven components, matching the outputs. The result has four components, matching the inputs. The Jacobian is seven-by-four.

Related terms

  • Jacobian
  • Jacobian-vector product
  • Backpropagation
  • Chain rule
  • Gradient
  • Checkpointing
  • Sensitivity analysis
  • reverse-mode automatic differentiation
  • cotangent

Learn next

Write a small Jacobian explicitly and calculate both Jv and J transpose u. Check their dimensions and verify u transpose Jv equals the VJP dotted with v. Then follow the same calculation through a composition of two functions to understand reverse accumulation.

  • Jacobian

More in Probabilistic machine learning

  • Variational filtering sequential Monte Carlo - Variational filtering sequential Monte Carlo refers here to combining particle-based sequential inference with variational optimization of proposals or...
  • Variational message passing - Variational message passing, or VMP, is a way to organize variational inference as local computations and messages in a probabilistic graphical model.
  • Variational recurrent neural network - A variational recurrent neural network, or VRNN, is a sequence model that combines recurrent neural-network memory with latent random variables at each time...
  • Vector-quantized generative adversarial network - A vector-quantized generative adversarial network, usually called VQGAN, learns a discrete codebook for image representations while using reconstruction...
  • Vector-quantized variational autoencoder - A vector-quantized variational autoencoder, or VQ-VAE, is a generative representation model that maps encoder outputs to entries in a learned discrete...
  • Von Mises-Fisher distribution - The von Mises-Fisher distribution is a probability model for directions represented by unit-length vectors.

Sources

JAX documentation, Forward- and reverse-mode autodiff, VJP definitions and composition: https://docs.jax.dev/en/latest/jacobian-vector-products.html

JAX autodiff cookbook, practical derivative-product operations: https://docs.jax.dev/en/latest/301/cookbook.html

The three-output calculation and fictional simulation case are original teaching illustrations. The entry uses real-valued functions; complex differentiation requires additional conventions.

Previous: Vector field reconstruction Next: Vector length

  • Home
  • Daily AI Updates
  • Blog
  • AIFBL Academy
  • AI Glossary
  • Products and Tools
  • Store
  • Free CFO AI Prompt Pack
  • About Shihan Sheriff
  • Privacy Policy
  • Terms of Use
  • Refund Policy
  • Cookie Policy
  • Earnings Disclaimer