Function as a DAG: every edge a derivative, every path a product

This article assumes what it means for a function to be represented as a directed acyclic graph (DAG).

E.g. the function can be represented as a DAG as shown here with the nodes defined like so:

Now let's say the goal is to compute the derivative of the output with respect to each input : . It has a nice visual interpretation in terms of the DAG:

If each edge carries the local derivative , then is the sum over paths from to of the product of edge derivatives along each path.

It also leads to the following insights:

• what JVP and VJP compute, and why they are equivalent
• when one is more efficient than the other

---

The gradient, one input at a time

First, let us assign to each edge the partial derivative .

Remember that is the sum over all paths products. To calculate it, we can follow these steps:

1. Find all paths from to .
2. For each path, multiply the edge derivatives along that path.
3. Sum the products from all paths.

Gradient w.r.t :

• There is only one path from to . The product of the edges along this path is .
• So

Gradient w.r.t :

• There are two paths: path one and path two from to .
• The product along the first path is .
• The product along the second path is .
• So

Gradient w.r.t :

• There is only one path from to . The product of the edges along this path is .
• So

---

Finding an efficient way to calculate sum of path products

Note that these sums can be computed in primarily two ways:

• Start from each input , follow all paths to , multiply edge derivatives along the way, and sum where the nodes meet. This is the JVP (forward mode) approach.
• Start from the output , follow all paths back to each input , multiplying edge derivatives along the way, and sum at the end. This is the VJP (reverse mode) approach.

Let's count the number of multiplications and additions needed for each approach.

Counting JVP vs VJP

Forward mode (JVP) needs one sweep per input (seed ). Reverse mode (VJP) needs one sweep per output (seed ).

	JVP	VJP
Sweep for	(3 muls, 0 adds)	—
Sweep for	(5 muls, 1 add)	—
Sweep for	(3 muls, 0 adds)	—
Backward sweep	—	(7 muls, 1 add)
Total	(11 muls, 1 add)	(7 muls, 1 add)

---

Put the tally on the canvas:

• Add the JVP and VJP counts

VJP needs one sweep, not three, so it wins as inputs grow. In general, VJP is more efficient for a function with many inputs and few outputs.

---

Inverting the function

If the edges of the DAG are reversed, one gets a function from to .

The edge derivatives are also switched for each edge. What was JVP is now VJP and vice versa. In this case, JVP wins. In general, VJP is more efficient for a function with few inputs and many outputs. In practice this rarely happens, if at all.