Computational Graphs

A computational graph is a directed acyclic graph (DAG) $G = (V, E)$ that represents a mathematical expression as a composition of elementary operations. Each node $v \in V$ holds a value: either an input (a leaf node, such as a parameter or a data point), a constant, or the output of an operation applied to the values of its parent nodes. Each directed edge $(u, v) \in E$ means "the value at $u$ is an input to the operation that produces $v$." Because the graph is acyclic, every node's value can be computed once all of its predecessors' values are known. The key design choice is granularity: operations are decomposed into pieces small enough that each has a simple, known derivative (addition, multiplication, matrix multiplication, $\exp$, $\log$, ReLU, and so on), but large enough that the graph doesn't become unmanageable. A dense layer $y = Wx + b$ followed by an activation $a = \sigma(y)$ is not one atomic node — it is expanded into a multiply node, an add node, and an activation node, each with its own simple local derivative. This decomposition is exactly what makes the graph amenable to a generic differentiation algorithm rather than a bespoke derivation for every architecture. Two traversal orders matter. A forward pass processes nodes in topological order (each node after all its inputs), computing values bottom-up from leaves to the final output — this is ordinary Forward Propagation. A backward pass processes nodes in reverse topological order, propagating derivatives of the final output with respect to each node, from the output back to the leaves — this is the traversal that Backpropagation performs, and it is what lets The Chain Rule for Neural Networks be executed as an algorithm rather than written out symbolically term by term. The computational graph is the shared data structure that both passes operate on, and it is the object that general-purpose Automatic Differentiation systems build and traverse under the hood.

Building the graph. As the forward computation executes, each operation is intercepted and a node is created recording (a) the operation type, (b) references to its input nodes ("parents"), and (c) the computed output value. This can happen statically (the whole graph is built before any data flows, as in classic TensorFlow 1.x) or dynamically (the graph is built on the fly as Python executes, as in PyTorch's "define-by-run" — but the resulting structure is the same DAG either way).

Topological ordering. Because every edge points from a producer to a consumer, any valid evaluation order must place each node after all of its predecessors. Such an order — a topological sort — always exists for a DAG (and only for a DAG; a cycle would make no such order possible, which is precisely why the graph must be acyclic).

Forward evaluation. Walk the nodes in topological order; each operation node computes its value from its already-computed parents.

Backward evaluation (gradient computation). Walk the nodes in reverse topological order starting from the output. Each node accumulates, from every child that depends on it, the child's contribution to $\partial L/\partial(\text{this node})$ via the chain rule, using only the local derivative of the operation that produced the child. ```

# ---------- Node definition ---------- struct Node: op # operation type, e.g. "input", "add", "mul", "matmul", "relu" parents # list of Node references (empty for leaves) value # cached forward output grad # accumulated dL/d(this node), initialized to 0

# ---------- Building the graph (happens during forward prop) ---------- function make_node(op, parents): n = Node(op=op, parents=parents, value=None, grad=0) for p in parents: p.children.append(n) # optional: for eager reverse traversal return n

# ---------- Topological sort (DFS-based) ---------- function topo_sort(output_node): visited = empty set order = empty list function visit(n): if n in visited: return visited.add(n) for p in n.parents: visit(p) order.append(n) # post-order => parents appear before n visit(output_node) return order # order[0..k-1] is a valid topological order

# ---------- Forward pass ---------- function forward(order): for n in order: if n.op == "input": continue # value already set externally n.value = apply(n.op, [p.value for p in n.parents]) return order[-1].value # value of the final (output) node

# ---------- Backward pass ---------- function backward(order): output = order[-1] output.grad = 1.0 # dL/dL = 1, seed gradient for n in reversed(order): if n.op == "input": continue local_grads = local_derivative(n.op, [p.value for p in n.parents], n.value) for (p, dvalue_dparent) in zip(n.parents, local_grads): p.grad += n.grad dvalue_dparent # chain rule, accumulate at shared nodes return { n: n.grad for n in order if n.op == "input" } ``` The critical line is `p.grad += n.grad dvalue_dparent` (accumulation, not assignment). If a node is used by multiple downstream operations — the normal case in any network with parameter sharing or skip connections — its total gradient is the sum over all paths, exactly as the multivariable chain rule requires.

Notice that $x$ feeds into two downstream nodes ($a$ and $b$), so its gradient in the backward pass is the sum of two contributions — this is the accumulation step in the pseudocode, and it is the diagram's clearest illustration of why the chain rule on graphs requires summation over all paths, not just one.

Time complexity. Building the graph costs $O(1)$ additional work per operation beyond the operation's own cost (just allocating a node and recording parent pointers), so the forward pass with graph construction costs the same asymptotic time as forward propagation alone: $O(\sum_i c_i)$, where $c_i$ is the cost of operation $i$. The backward pass visits every node exactly once (in reverse topological order) and, at each node, does work proportional to the number of incoming edges to that node in the forward graph (equivalently, the number of parents), computing one local-derivative-times-incoming-gradient product per parent. Summed over the whole graph this is again $O(\sum_i c_i)$ up to a small constant factor — computing a local derivative is generally no more expensive than the operation itself (e.g., the derivative of a matrix multiply is another matrix multiply of the same shape). So:

This is the single most important efficiency guarantee of the computational-graph approach: gradients of a scalar loss with respect to all parameters simultaneously cost a small constant multiple of one forward pass, independent of the number of parameters. This is dramatically better than naively estimating each partial derivative separately (e.g., by finite differences), which would cost $O(P \cdot T_{\text{forward}})$ for $P$ parameters — utterly infeasible when $P$ is in the millions or billions.

Space complexity. The forward pass must cache every intermediate node's value, because the backward pass needs it to compute local derivatives (e.g., $\partial(xy)/\partial x = y$ requires knowing $y$). This gives $O(N)$ additional memory where $N$ is the number of nodes (intermediate activations) — a real cost, and the reason activation memory, not just parameter memory, dominates training large models. (Gradient checkpointing trades recomputation time for reduced cached memory, but does not change the fundamental $O(N)$ vs. $O(\sqrt{N})$-style tradeoff story here — the DAG's structure is what makes such tradeoffs even expressible.)

Correctness argument. The argument has two parts. 1. A topological order exists and is valid. Since $G$ is a DAG, a topological order exists (standard graph theory: repeatedly remove a node with no unprocessed predecessors; a DAG always has one, otherwise a cycle would be forced). Processing nodes in this order for the forward pass guarantees every operation's inputs are ready before it executes — this is precisely the correctness condition for the forward pass. 2. The backward pass computes the correct total derivative. This follows from the multivariable chain rule. If node $n$'s value feeds into children $c_1, \dots, c_k$, then by definition of total derivative,

Processing nodes in reverse topological order guarantees that when node $n$ is processed, every child $c_j$ of $n$ has already had its own gradient $\partial L / \partial c_j$ fully accumulated (all of $c_j$'s own children were processed earlier in the reverse order, since they appear later in the forward topological order). So the accumulation `p.grad += n.grad * local_derivative` performed once per edge $(p, n)$, summed over all children of $p$, exactly reconstructs the right-hand side of the chain-rule identity above. By induction from the output node (seeded with $\partial L/\partial L = 1$) down to the leaves, every node's `grad` field equals its true partial derivative of the loss when the backward pass completes.

Honest tradeoffs. The graph abstraction is not free: it requires storing the entire computation history (memory-hungry for very deep or very wide networks, motivating checkpointing and reversible architectures), and dynamic (define-by-run) graph construction adds Python-level overhead per operation compared to a graph compiled once and reused (as in static-graph or JIT-compiled systems). The abstraction also assumes the computation is expressible as a DAG — genuinely cyclic dependencies (e.g., certain fixed-point iterations) need different tools (implicit differentiation) rather than plain graph backward-passes.

Trace $f(x, y) = (x \cdot y) + \sin(x)$ at $x = 2.0$, $y = 3.0$, matching the diagram above.

Building the graph. Nodes, in creation/topological order: `x` (leaf), `y` (leaf), `a = mul(x, y)`, `b = sin(x)`, `f = add(a, b)`.

Forward pass (topological order `x, y, a, b, f`):

Backward pass (reverse order `f, b, a, y, x`), seeding `f.grad = 1`:

Final gradients: $\partial f/\partial x = 2.5839$, $\partial f/\partial y = 2.0$.

Sanity check by direct calculus. $f(x,y) = xy + \sin x \Rightarrow \partial f/\partial x = y + \cos x = 3.0 + \cos(2.0) = 3.0 - 0.4161 = 2.5839$, and $\partial f/\partial y = x = 2.0$. Both match exactly — confirming the reverse-topological accumulation reconstructed the true total derivative, including correctly summing the two paths from $x$ (through $a$ and through $b$) rather than overwriting one with the other.

The idea of decomposing a computation into elementary steps with known local derivatives, then combining them mechanically, has roots in numerical analysis long before deep learning. Automatic differentiation was formalized in the 1960s–1980s, and its reverse mode — the mathematical procedure that computational graphs make executable as a graph algorithm — was independently rediscovered several times, most famously in the context of training multilayer neural networks in the 1980s.

Every mainstream deep learning framework — PyTorch, TensorFlow, JAX, MXNet — represents a model's forward computation internally as a computational graph, whether built eagerly (dynamic) or compiled ahead of time (static/traced via `jit`/`XLA`). This graph is exactly what enables `loss.backward()` or `jax.grad(f)` to compute gradients for arbitrary user-defined architectures without any hand-derived calculus. Compilers for deep learning (XLA, TorchScript, ONNX Runtime, TensorRT) operate directly on this graph representation to perform operator fusion, memory planning, and hardware-specific code generation, since the DAG exposes exactly which computations are independent and can be reordered, parallelized, or fused. Beyond neural networks, computational graphs underlie probabilistic programming systems, differentiable physics simulators, and differentiable rendering pipelines — anywhere a complex forward model needs gradients with respect to its inputs or parameters. The same graph abstraction is also what makes distributed training frameworks able to automatically insert communication operations (e.g., gradient all-reduce) at the correct points in the backward traversal.

Ask the tutor: How does the computational graph change when a node's output is reused in a loop, like in a recurrent neural network?