To illustrate automatic differentiation with dynamic control flow, consider a contrived function where a while loop scales an input vector $$\mathbf{a}$$ until its norm exceeds a threshold, followed by an if statement that returns different multiples of the vector based on the sum of its elements. Even though the number of loop iterations and the final scaling factor depend entirely on the specific values of the input, the deep learning framework constructs a concrete computational graph during the forward pass. Because the resulting function $$f(\mathbf{a})$$ is linear with a piecewise defined scale, its gradient with respect to $$\mathbf{a}$$ can be logically verified to match the element-wise division $$f(\mathbf{a}) / \mathbf{a}$$.

Example: Differentiating a Function with Dynamic Control Flow

Dynamic control flow is pervasive in modern deep learning applications. For instance, in natural language processing, the sequence of operations and the resulting computational graph inherently depend on the varying lengths of input text sequences. Because it is impossible to compute the exact gradient formula for such models a priori, the ability of automatic differentiation to handle dynamically realized computational graphs is vital for training complex statistical models.

Dynamic Control Flow in Deep Learning

One of the major benefits of automatic differentiation is its ability to compute gradients even when the sequence of operations depends on dynamic control flow, such as Python if statements, while loops, and arbitrary function calls. Since the exact path of computation is not known a priori, the computational graph is realized dynamically only during the forward evaluation on a specific input. Once this specific graph is constructed, deep learning frameworks can execute the backward pass to accurately calculate the gradient of the resulting variable.

Claude

A computational graph is a directed dependency graph where nodes represent mathematical operations or variables, mapping the functional dependency from inputs to outputs. When evaluating a vector-valued function, this graph is traversed in a forward direction. Conversely, to compute gradients automatically, the graph is traced in a backwards direction from output to input, which involves multiplying matrices of partial derivatives along the paths.

Computational Graph

Dive into Deep Learning

To unfold a computational graph is to unfold a recursive or recurrent computation into repetitive structure, typically corresponding to a chain of events.

Unfolding Operation

During automatic differentiation, it is sometimes necessary to exclude specific calculations from the recorded computational graph. This is often required when inputs generate auxiliary intermediate variables that should not contribute to the gradient computation. To achieve this, the computational graph associated with these intermediate terms is detached from the final result. A new variable is instantiated with the same numerical value, but its provenance—the record of operations that created it—is completely erased. Consequently, this new variable acts as a constant with no ancestors in the graph, preventing gradients from flowing back through it to earlier variables.

Detaching Computation

While automatic differentiation is used to optimize deep learning models in a statistical sense (e.g., through gradient descent), the computational optimization of the autograd libraries themselves is a vital area of framework design. Framework developers leverage techniques from compilers and graph manipulation to execute gradient computations in the most expedient and memory-efficient manner possible.

Optimization of Automatic Differentiation Libraries

Backpropagation is a systematic computational procedure for applying the chain rule to calculate gradients automatically. It operates by traversing a computational graph in a backwards direction—from the output loss back to the input parameters—multiplying matrices of partial derivatives at each step to determine how parameters affect the final output.

Backpropagation

When evaluating a vector-valued function, its underlying computational dependency graph is traversed in a forward direction, from the input variables to the output. This forward pass calculates the function's intermediate and final values before any gradients are computed.

Forward Graph Traversal

To compute the gradients of a vector-valued function, its computational dependency graph is traversed in a backwards direction, tracing from the output back to the inputs. During this backward pass, matrices of partial derivatives are multiplied along the paths to determine how changes in the input variables affect the final output.

Backward Graph Traversal

Automatic Differentiation with Dynamic Control Flow

Plotting a computational graph helps visualize the mathematical dependencies between variables and operators during the forward propagation of a neural network. In these visual representations, variables are typically denoted by squares and operators by circles. Arrows are used to indicate the direction of data flow; in a forward pass, this flow typically proceeds from the network inputs (often positioned at the lower-left) toward the outputs (at the upper-right).

Learn Before

Related

Learn After