Chain Rule for Computation Graphs

Goal

Turn the chain rule from a formula about nested functions into an algorithm for differentiating a program.

The central idea is:

A large derivative is assembled from small local derivatives, following the same dependency structure as the original computation.

Why Composition Creates A New Problem

Lesson 36 described one function near one point with one Jacobian:

Real models are compositions. An input becomes a linear transformation, then an activation, then another transformation, and eventually a scalar loss. Writing the entire model as one expanded formula would hide the reusable structure and create a large bookkeeping problem.

A computation graph keeps that structure explicit:

input -> intermediate values -> prediction -> loss

Each node performs a small operation. The chain rule explains how their local derivatives combine.

The Invariant: Derivatives Must Connect Compatible Spaces

Suppose:

and:

The composition is:

Its derivative must map an $n$-dimensional input displacement to a $k$-dimensional output displacement. The Jacobian chain rule is:

The shapes force the order:

This is the same composition rule as matrix transformations. $F$ acts first, so its Jacobian is closest to the input displacement on the right.

Notation / Symbol Table

The bar in $\bar v$ does not mean an average. In this lesson it is compact notation for the derivative of the final scalar loss with respect to $v$.

First Example: A Scalar Chain

Let:

The graph is:

x --square--> u --scale and shift--> L

The local derivatives are:

The chain rule multiplies them:

At $x=2$, the forward values are:

The backward sensitivities are:

The forward pass computes values. The backward pass computes how the final loss depends on those values.

Why The Backward Pass Starts With One

Every reverse pass begins at the final scalar output with:

This is not an arbitrary initialization. A quantity changes one-for-one with itself. That unit sensitivity seeds the repeated chain-rule multiplications.

A Computation Graph Is A Factorization

Consider:

Factor it into named operations:

The graph is:

x --\
     multiply -> u -> add 1 -> v -> square -> L
y --/

At $(x,y)=(2,3)$, the forward pass gives:

Now record only the local derivatives:

Propagate backward:

Therefore:

The graph avoided expanding the whole expression. Each node needed only its own operation, stored forward values, and incoming sensitivity.

Branching Means Addition

Multiplication handles consecutive links along one path. Addition handles several paths from the same variable to the output.

Let:

The input $x$ influences $L$ along two paths:

       -> square -> u --\
x ---->              add -> L
       -> scale  -> v --/

The total derivative is the sum of both path contributions:

Since both addition-node derivatives are $1$:

This yields the operational rule:

Multiply local derivatives along a path. Add contributions from different paths that meet at the same value.

Failing to add at a branch is one of the most common computation-graph mistakes.

Vector Composition

The same rule works when intermediate values are vectors.

Let:

Then:

Their product is:

Because $G\circ F$ is scalar-valued, this result is the transpose of its column gradient:

This transposed form is important in machine learning. It says that an output gradient can be pulled backward through $F$ by multiplying with $J_F^T$.

A Concrete Vector Example

Define:

The local Jacobians are:

and:

At $(x,y)=(1,2)$, the intermediate value is:

Therefore:

Compose them:

Thus:

The shapes are not decoration. They tell us which maps can be composed and in what order.

Forward Mode And Reverse Mode

There are two natural ways to apply the chain rule without constructing every full derivative matrix.

Forward Accumulation

Start with an input direction $\Delta x$ and push its effect forward:

This computes a Jacobian-vector product:

It efficiently answers: How does this chosen input direction affect every output?

Reverse Accumulation

Start with sensitivity at a scalar loss and pull it backward:

This computes a vector-Jacobian product, usually represented with column gradients through the transpose.

It efficiently answers: How does the one final loss depend on every earlier input or parameter?

Machine learning commonly has millions of parameters but one scalar loss. That shape makes reverse accumulation especially valuable: one backward sweep can produce sensitivities for every parameter.

This Is The Foundation Of Backpropagation

Backpropagation is not a different derivative rule. It is reverse-mode chain-rule accumulation specialized to layered models and scalar losses.

This lesson focuses on the calculus and graph bookkeeping. Lesson 43 will apply the same mechanism to neural-network parameters, activations, and layers.

Connection To Programming

Here is the earlier branching example encoded as an explicit forward pass and backward pass:

type ForwardState = Readonly<{
  loss: number;
  scaled: number;
  squared: number;
  x: number;
}>;

function forward(x: number): ForwardState {
  const squared = x * x;
  const scaled = 3 * x;
  const loss = squared + scaled;

  return { loss, scaled, squared, x };
}

function backward(state: ForwardState): number {
  const lossSensitivity = 1;
  const squaredSensitivity = lossSensitivity * 1;
  const scaledSensitivity = lossSensitivity * 1;

  const fromSquared = squaredSensitivity * (2 * state.x);
  const fromScaled = scaledSensitivity * 3;

  return fromSquared + fromScaled;
}

const state = forward(2);
const derivative = backward(state);

console.assert(state.loss === 10);
console.assert(derivative === 7);

Notice the software structure:

1. The forward pass stores values needed by local derivative formulas.
2. The backward pass receives a sensitivity from each child.
3. Each edge multiplies by one local derivative.
4. Contributions are added when they reach a shared parent.

A full automatic-differentiation engine generalizes this pattern by storing operations and their local backward rules dynamically.

Common Mistakes

Reversing The Jacobian Product

For $G\circ F$, write the displacement flow:

Therefore $J_F$ is closest to $\Delta x$ on the right:

Evaluating A Local Derivative At The Wrong Value

$J_G$ must be evaluated at the intermediate value $F(x)$, not at the original input $x$ unless those happen to be the same type and value.

Multiplying Across Branches Instead Of Adding

Consecutive edges multiply. Alternative paths add. A reused value can receive several valid sensitivity contributions.

Mixing Forward Values With Backward Sensitivities

$v$ is an intermediate value. $\bar v=\partial L/\partial v$ is a sensitivity. They may have compatible shapes, but they represent different things.

Materializing A Full Jacobian Unnecessarily

If the goal is only a directional effect or scalar-loss gradient, Jacobian-vector and vector-Jacobian products can apply the local map without storing every entry of a large Jacobian.

Exercises

1. In one sentence, explain why the chain rule for $G\circ F$ has the order $J_G(F(x))J_F(x)$ rather than $J_F(x)J_G(F(x))$.

2. Let:

   $$u=2x-1, \qquad L=u^3.$$

   Draw the two-operation computation graph, compute the forward values at $x=2$, and use local derivatives to find $dL/dx$.

3. For:

   $$u=xy, \qquad v=u+2, \qquad L=v^2,$$

   run a forward and backward pass at $(x,y)=(3,-1)$ and report $\partial L/\partial x$ and $\partial L/\partial y$.

4. Let $F:\mathbb R^4\to\mathbb R^3$ and $G:\mathbb R^3\to\mathbb R^2$. State the shapes of $J_F$, $J_G$, and $J_{G\circ F}$. Explain how the dimensions type-check.

5. A scalar $x$ is used in both:

   $$u=x^2, \qquad v=\sin x, \qquad L=u+v.$$

   Explain why the backward contributions to $x$ must be added, then compute $dL/dx$.

6. For the vector composition:

   $$F(x,y)= \begin{bmatrix} x-y\\ x+y \end{bmatrix}, \qquad G(u,v)=u^2+3v,$$

   compute $J_F$, $J_G$, and $\nabla(G\circ F)$ at $(x,y)=(2,1)$ using Jacobian multiplication.

7. Suppose an intermediate scalar $v$ feeds two children $a(v)$ and $b(v)$, and $L=a+b$. Write $\bar v$ in terms of $\bar a$, $\bar b$, $da/dv$, and $db/dv$. State the general branch-accumulation rule in words.

8. Compare forward and reverse accumulation for a function $F:\mathbb R^{1000}\to\mathbb R$. Which direction is naturally efficient for computing the full gradient of the scalar output, and why?

9. Write TypeScript forward and backward functions for:

   $$L(x,y)=(xy+1)^2.$$

   Store only the intermediate values needed by the backward pass. Check the result at $(2,3)$ against the analytic gradient $[42,28]$.

10. In your own words, explain why a computation graph does not change the function being differentiated. What does the graph add that the single expanded expression does not?

Practice 37.1: Local Derivatives And Branch Accumulation

This is a short correction checkpoint, not a new lesson. Work it on paper without expanding the full functions first. For every computation, write the forward values separately from the backward sensitivities.

1. Local Values Stay Local

Let:

At $x=-1$:

1. Compute the forward values $u$ and $L$.
2. Write the local derivatives $\partial L/\partial u$ and $\partial u/\partial x$.
3. Compute $dL/dx$ using the stored value of $u$.
4. In one sentence, explain why replacing $u^2$ with $x^2$ would change the computation.

2. Follow The Actual Operations

For:

write $J_F(x,y)$ by differentiating one output component at a time. Label rows by outputs and columns by inputs before filling in the entries.

Then compute:

and explain what the resulting vector means.

3. Addition At A Branch

Let:

1. Draw the two branches from $x$ to $L$.
2. Write the local derivatives at the final addition node.
3. Compute $dL/dx$ symbolically.
4. Evaluate it at $x=1$.
5. Explain why neither branch’s forward value multiplies the other branch’s derivative.

4. Accumulate Child Contributions

An intermediate value $z$ feeds two children:

At $z=3$, run the backward pass from $\bar L=1$. Show $\bar a$, $\bar b$, the contribution from each child to $\bar z$, and their accumulated total.

Then write the general rule:

using $\bar a$, $\bar b$, $da/dz$, and $db/dz$.

5. One Complete Composition

Define:

At $(x,y)=(2,1)$:

1. Compute the forward value $F(2,1)$.
2. Write $J_F(x,y)$ and $J_G(u,v)$ with shapes.
3. Evaluate both Jacobians at the correct local values.
4. Compute $J_{G\circ F}(2,1)=J_G(F(2,1))J_F(2,1)$.
5. Convert the resulting scalar-output Jacobian into a column gradient.
6. Check the result by expanding $G(F(x,y))$ only after completing the graph-based calculation.

Readiness Check

Lesson 37 is ready to close when the work consistently shows all four habits:

• Local derivatives are evaluated using the stored value at their own node.
• The derivative matches the operation actually written in each node or component.
• Derivatives multiply along one path.
• Contributions add when several paths return to the same parent.

Summary

• A computation graph factors a function into small operations and explicit dependencies.
• The chain rule composes local derivative maps.
• For $G\circ F$, the Jacobian is $J_G(F(x))J_F(x)$.
• The forward pass computes and stores values.
• The backward pass propagates final-output sensitivities through local derivatives.
• Local derivatives multiply along one path and contributions add across multiple paths.
• Forward accumulation pushes a chosen input direction toward outputs.
• Reverse accumulation pulls a scalar loss sensitivity back toward every input and parameter.
• Backpropagation is reverse chain-rule accumulation applied to neural networks.

Next Lesson

Lesson 38 makes matrix-calculus layout conventions explicit. We will use shape checking to distinguish gradients with respect to vectors and matrices, translate between common notation conventions, and derive reusable identities without relying on memorized formula tables.

Completed Exercises

Handwritten work submitted after completing this lesson.