# Matrix Calculus and Jacobians

Differentiating functions whose inputs and outputs are vectors or matrices.

**Rests on:** Linear Transformations, Tensors and Shapes

**Leads to:** Encoder-Decoder, Lagrange Multipliers

*Mathematical Foundations · Calculus*


## Why learn this

This is the machinery that lets a computer tune millions of knobs at once. When you train an AI model — whether it's a chatbot, a fraud detector, or a photo classifier — you're constantly asking: 'If I nudge this weight, how much does the final answer change?' Matrix calculus is how that question gets answered for every weight in a single step. Learning this unlocks the next step: backpropagation, the algorithm that actually updates all those weights. It also connects to the tools you'll use every day — PyTorch and TensorFlow compute gradients automatically, but knowing what they're doing underneath lets you debug when things go wrong, understand why certain tricks like learning rates matter, and eventually design your own custom layers. Without this, neural networks are a magic box; with it, they become a machine you can reason about and control.


## The intuition

You already know that a function takes an input and gives an output — like a vending machine: you put in a code, you get a snack. The derivative of a function tells you how sensitive the output is to small changes in the input. If you press the button a tiny bit harder, does the snack come out faster or slower? That's the derivative — a single number describing the slope of the vending machine's response. Now imagine a function that takes not one input, but three: income, debt, and job tenure. And its output isn't a single number, but a list of three numbers — say approval, interest rate, and risk score. That's a 'vector function'. The derivative of a vector function can't be a single number anymore, because each output depends on every input in a different way. You need a whole table of numbers: each entry tells you how one particular output changes when one particular input changes slightly. This table is called the Jacobian. It's like a sensitivity chart for the entire machine at once. And when you have hundreds of inputs and hundreds of outputs, that table isn't just a piece of paper — it's a matrix, and the rules of matrix multiplication let you combine these tables when you chain functions together, which is exactly what happens in a deep neural network.


## An analogy

Imagine you're the manager of a small tea factory, and you have a recipe that depends on three ingredients: tea leaves, milk, and sugar. The taste of the final tea depends on each ingredient, but not equally. Let's say the recipe gives you two taste qualities: sweetness and creaminess. So you have a machine with three inputs (grams of each ingredient) and two outputs (sweetness and creaminess scores on a scale of 0 to 10). Now, as the manager, you want to tweak the recipe to make the tea sweeter without changing creaminess too much. You need to know: if I add 1 gram of sugar, how much does sweetness go up? And how much does creaminess change? And if I add 1 gram of milk, how does sweetness change? And creaminess? The Jacobian is a table that answers all these questions at once. Each row is an output (sweetness, creaminess), each column is an input (leaves, milk, sugar). The entry in the row for sweetness and the column for sugar is a single number: the change in sweetness per unit change in sugar, holding everything else constant. That's the partial derivative. Now, the tea factory doesn't just make one type of tea — it has a pipeline. The taste scores feed into a second machine that outputs a quality score and a price. To know how changing sugar affects the final price, you can't just look at the taste table — you need to combine the two tables: the Jacobian of the taste machine and the Jacobian of the quality machine. This combining is exactly matrix multiplication. And in a neural network, each layer is such a machine, and training means adjusting the knobs (ingredients) of each layer to make the final output better. The Jacobian tells you which direction to turn each knob. The analogy stops at one point: in the real world, the tea recipe is fixed, but in AI, we start with random ingredients and learn the best ones from data. Also, a real network might have thousands of inputs per layer, so the Jacobian table is enormous — much bigger than any recipe chart. But the idea is the same.


## Definition

The Jacobian is a matrix that stores all the partial derivatives of a vector-valued function, showing how each output changes with respect to each input, and it's the key object that the chain rule uses to propagate gradients through a neural network.


## How it connects

You haven't started this path yet, so let's lay the groundwork. This concept sits at the meeting point of two older branches: calculus (the study of change) and linear algebra (the study of lists of numbers and the operations on them — like adding, scaling, and the dot product). The Jacobian is fundamentally a table of numbers, and combining Jacobians uses the rules of matrix multiplication. You have notes on 'Matrix Multiplication, Determinants and Inverses' — those will come in handy. Also, this directly feeds into 'Automatic Differentiation' and 'Computational Graphs', which are the practical tools used in PyTorch. The Jacobian itself is rarely computed in full — that's where those topics come in, computing vector-Jacobian products instead. So treat this as the conceptual root.


## The mathematics

We need to build up from nothing. First, what is a derivative? Imagine you're driving a car, and you have a speedometer that shows your speed in km/h. At a particular moment, the derivative of your distance with respect to time is exactly your speed — how many kilometres you cover per hour, at that instant. If you accelerate, the speed changes; the derivative tells you the rate at that exact moment. Now, a function can have more than one input. Suppose your tea preference scores depend on sugar and milk. If I ask 'how does sweetness change when I add a tiny bit of sugar, while I keep milk fixed?' — that's called the partial derivative of sweetness with respect to sugar. It's just the derivative, but we hold all other inputs constant. Now, a vector-valued function is a function with multiple outputs. Each output can be partially differentiated with respect to each input, producing a grid of numbers. That grid is the Jacobian. The rows are the outputs, the columns are the inputs. To compute it, you simply compute each partial derivative one by one. There's no new calculus here — just repeated application of the single-variable derivative rule, done for every pair of input and output. The Jacobian also has a geometric meaning: it's a linear approximation of the function near a particular point. If the function is the tea machine, the Jacobian tells you the 'slope' of each output with respect to each input, and those slopes together tell you the best linear guess for what the output would be if you change the inputs slightly. The chain rule for vectors says: if you compose two functions (put one after the other), the Jacobian of the composite is the matrix product of the two Jacobians. This is why matrix multiplication is so central — it's not just a bookkeeping trick, it's the correct way to combine sensitivities.

Given a function f: R^n -> R^m (n inputs, m outputs), the Jacobian matrix J is an m×n matrix with entries: J_{ij} = ∂f_i / ∂x_j, where f_i is the i-th output and x_j is the j-th input. Each row i corresponds to the gradient of the i-th output.

For a composite function g(f(x)), the chain rule in matrix form is: J_{g∘f} = J_g(f(x)) · J_f(x), where · is matrix multiplication.


## A worked example

```
Let's use a tiny function with 2 inputs and 2 outputs. Inputs: x1 (grams of sugar), x2 (grams of milk). Outputs: y1 (sweetness), y2 (creaminess). The function is defined as: y1 = 3*x1 + 2*x2, y2 = x1^2 + x2. We will evaluate the Jacobian at the point (x1=1, x2=2).

Step 1: ∂y1/∂x1 = 3   # We treat x2 as a constant. Derivative of 3*x1 is 3.
Step 2: ∂y1/∂x2 = 2   # Derivative of 2*x2 w.r.t x2 is 2, while 3*x1 is constant.
Step 3: ∂y2/∂x1 = 2*x1. At x1=1, this equals 2*1 = 2.   # Using the power rule: derivative of x^2 is 2x.
Step 4: ∂y2/∂x2 = 1   # Derivative of x2 is 1, while x1^2 is constant.

Result: The Jacobian matrix J is a 2x2 table. Rows: y1, y2. Columns: x1, x2. J = [[3, 2], [2, 1]]. At the point (1,2), this matrix tells you how each output changes for small changes in each input. For example, increasing sugar by 1 unit increases sweetness by 3 units and creaminess by 2 units (approximately, for very small changes).
```


## From scratch, in Python

```python
import math

# Define a tiny function with 2 inputs and 2 outputs
def f(x):
    # x is a list of two numbers: [x1, x2]
    y1 = 3 * x[0] + 2 * x[1]      # sweetness
    y2 = x[0] ** 2 + x[1]        # creaminess
    return [y1, y2]

# Compute the Jacobian numerically using finite differences (central difference)
def jacobian(f, x, h=1e-5):
    n = len(x)  # number of inputs
    # Call f to see how many outputs there are
    f0 = f(x)
    m = len(f0)
    # Initialize J as a list of m rows, each with n zeros
    J = [[0.0] * n for _ in range(m)]
    for j in range(n):
        x_plus = x.copy()
        x_plus[j] += h
        f_plus = f(x_plus)
        x_minus = x.copy()
        x_minus[j] -= h
        f_minus = f(x_minus)
        for i in range(m):
            J[i][j] = (f_plus[i] - f_minus[i]) / (2 * h)
    return J

# Test at point (1, 2)
x = [1.0, 2.0]
J_num = jacobian(f, x)
print("Numerical Jacobian:")
for row in J_num:
    print(row)

# Expected analytic Jacobian: [[3, 2], [2, 1]]
# The numerical result should be very close.
```


## The same thing, with a framework

```python
import numpy as np

# Using a numerical finite difference is not how frameworks do it.
# They use reverse-mode automatic differentiation, which computes the
# vector-Jacobian product efficiently.

# In PyTorch, you would define a tensor with requires_grad=True and call backward().
# Example:
import torch

x = torch.tensor([1.0, 2.0], requires_grad=True)
y1 = 3*x[0] + 2*x[1]
y2 = x[0]**2 + x[1]
# To get Jacobian, we would need to call backward separately for each output.
# But for vector-Jacobian product, we can do:
loss = y1 + y2
loss.backward()
print(x.grad)  # This is the gradient of (y1+y2) w.r.t x, not the full Jacobian.

# To get the full Jacobian rows, you would do:
J = []
for y in [y1, y2]:
    x.grad = None
    y.backward(retain_graph=True)
    J.append(x.grad.numpy())
print(np.array(J))

# Underneath, PyTorch computes vector-Jacobian products analytically using the
# chain rule, without ever forming the full matrix when not needed.
```


## What the code does

We start by defining a function `f` that takes a list `x` of two numbers. Inside, we compute two outputs: `y1` and `y2`.\n\nThen we define `jacobian(f, x)`. This function uses the central difference method to approximate derivatives: for each input `j`, we make a tiny positive step `h` and a tiny negative step, evaluate `f` at both points, and take the difference divided by `2h`. This gives an approximation of the partial derivative.\n\nWe build a list of lists `J` with `m` rows (one per output) and `n` columns (one per input), initialised to zeros. Then we fill it with the computed derivatives.\n\nFinally, we call it with `x = [1, 2]` and print the result. The output should be close to the exact Jacobian [[3, 2], [2, 1]].


## Common misconceptions

**The Jacobian is just a matrix of all derivatives — you can always compute it directly.** — To train a network, I can just write out the full Jacobian of the whole network and multiply it out.

For a network with millions of parameters, the full Jacobian would be millions times the number of outputs — astronomically large. Storing or multiplying it is impossible. In fact: Frameworks never build the full Jacobian. They use reverse-mode automatic differentiation to calculate the product of the Jacobian with a vector (the incoming gradient) — that's enough to update parameters.

**Bias is a bad thing in neural networks.** — The 'bias' term in a neuron introduces unfairness, so I should remove it.

In neural networks, 'bias' is a learned parameter that shifts the activation, allowing the model to fit data that isn't centered at the origin. It has nothing to do with social bias or fairness. In fact: The bias parameter is just a constant added to the weighted sum before the activation. It's a tool, not an error.

**The Jacobian is only for square matrices.** — A Jacobian must be a square matrix because I've only seen it with equal inputs and outputs.

Most neural network layers map a vector of one size to a vector of a different size. The Jacobian has a row for each output and a column for each input — it can be rectangular. In fact: The Jacobian's shape is (number of outputs) × (number of inputs) — usually not square.

**The chain rule for Jacobians is just multiplying numbers.** — When I chain two functions, I just multiply the two derivative numbers.

That works only for one input and one output. For vectors, each function's derivative is a matrix, and the chain rule uses matrix multiplication, not scalar multiplication. In fact: For vector-valued functions, the chain rule is: J(g∘f) = J(g) · J(f), where · is matrix multiplication.

**The Jacobian gives the exact change in output.** — If I increase input x1 by 0.1, the output y1 will change by exactly J[0][0] * 0.1.

The Jacobian is a linear approximation at a specific point. For small changes it's accurate, but for larger changes, the function's curvature means the actual change can differ. In fact: The Jacobian gives the local slope; the exact change is given by the function itself, which may be nonlinear.

**Gradient descent requires computing the full Jacobian every iteration.** — To update weights, I have to compute a table of all partial derivatives of all outputs with respect to all weights.

You only need the gradient of the loss (a scalar) with respect to each weight. That can be computed as a vector-Jacobian product, which is more efficient. In fact: Backpropagation computes the gradient of the loss with respect to all parameters without ever forming the full Jacobian matrix.


## In short

The Jacobian is a matrix that records how every output of a function changes when you nudge each input. It's the multi-dimensional generalization of the derivative. For neural networks, it's the building block of backpropagation: the chain rule for Jacobians uses matrix multiplication, so you can propagate gradients backward through layers. In practice, you never compute the full Jacobian — you use reverse-mode automatic differentiation to compute vector-Jacobian products, which is what PyTorch and TensorFlow do under the hood. Understanding the Jacobian helps you debug custom layers, design efficient training, and grasp why gradient-based optimisation works.


## Where this is used

The Jacobian is the engine behind gradient-based training in deep learning, and it appears everywhere in modern AI.\n\n**Large language models (LLMs):** An LLM like GPT has billions of parameters (weights and biases). Training it involves computing the gradient of the loss (a single number) with respect to every parameter. The Jacobian of the entire network with respect to all parameters would be impossibly huge — billions of columns. So frameworks use reverse-mode automatic differentiation, which computes the vector-Jacobian product — a way to get the gradient of the loss without ever storing the full Jacobian. The chain rule is applied layer by layer, and that's the essence of backpropagation. A developer doesn't hand-write Jacobians; they define the model with PyTorch, call `loss.backward()`, and the framework computes the gradients efficiently. If a custom layer's backward pass is wrong, the model fails to learn or diverges, and that's where a hand-derived Jacobian might be needed.\n\n**Small language models on devices:** When a model must run on a phone or laptop, it needs to be small and efficient. Quantization (reducing the precision of numbers) and pruning (removing some weights) are common. Both rely on knowing how much each weight contributes to the output — which is the gradient information from the Jacobian. For example, pruning uses the magnitude of gradients to decide which weights are least important. So even on a tiny device, the Jacobian (or its proxy) guides the compression.\n\n**Retrieval-augmented generation (RAG):** RAG combines a language model with a retrieval system that fetches relevant documents. The retrieval step often uses embeddings (vectors representing text) and a similarity measure like dot product. The model needs to learn good embeddings, and the training loss depends on the dot product between query and document embeddings. The gradient of that dot product with respect to the embedding parameters is a vector (a one-row Jacobian), and that's computed via backpropagation. So the Jacobian enables the embeddings to be tuned so that relevant documents are close to queries.\n\n**Agentic systems that plan and call tools:** These systems use a language model as a controller that decides which tools to call. The model's parameters are trained similarly to an LLM, so the Jacobian is central to training. At inference time, the Jacobian isn't used directly, but the trained model's decisions are emergent from the learned weights. Some advanced reinforcement learning approaches use policy gradients, which again rely on Jacobians of the policy outputs with respect to parameters.\n\n**Everyday applied ML:** Regression, classification, forecasting, recommendation — all rely on gradient descent to minimize a loss function. For a linear regression, the Jacobian of the loss with respect to weights is a simple vector, computed exactly. For neural networks, it's the chain of layer Jacobians. A developer often uses libraries like scikit-learn or PyTorch. When you do gradient descent manually, you compute the Jacobian (or use a library's autograd). If you implement your own custom loss, you need to ensure its derivative is correct; that's where you might hand-compute a Jacobian and use gradient checking to verify.\n\nIn practice, a developer rarely writes the full Jacobian by hand. They define the forward pass and let the framework compute gradients. But when a custom layer or loss is needed, a correct derivative is crucial — a wrong Jacobian can cause training to diverge, and debugging it requires understanding this concept.


---

Read this on the web: https://basavapurushottam.com/learn-ai/matrix-calculus-and-jacobians
