# Looping and branching with Algorithmic Differentiation

Algorithmic (aka Automatic) Differentiation is a wonderful technique for numerical computation of derivatives. I understand how it relates to the fact that we know how to deal with every elementary operation in a computer program, but I am not sure to get how this applies to every computer program.

every computer program, no matter how complicated, executes a sequence of elementary arithmetic operations

.. which I totally agree with. However, it sounds then that the number $m$ of variables newly produced from any $n$ initial input is fixed and can be determined from static code analysis. But this is not straightforward to me since constructs like:

if x_1 > x_2:                         # branching
perform 4 elementary operations
else:
perform 84 elementary operations
endif


and:

while x_1 < x_2:                      # looping
perform 2 elementary operations
endwhile


do exist in so-called « complicated » computer programs. This make the number (and the type) of elementary operations not straightforward to compute in advance. And I even suspect it is impossible to gess that in general, right?

Does AD support such branching and looping programs?
Are there extensions of AD adapted to programs that are not just intricate closed-form expressions?
How does AD deal with Turing-completeness?

• Judging from the examples, automatic differentiation supports general control structures as long as they only depend on fixed hyperparameters. In other words, your program should be equivalent to a sequence of fixed arithmetic operations, and then automatic differentiation works. Feb 21, 2017 at 18:23

AD supports arbitrary computer programs, including branches and loops, but with one caveat: the control flow of the program must not depend on the contents of variables whose derivatives are to be calculated (or variables depending on them). Here is an example:

if x = 3 then 9 else x * x


At close inspection you will recognize that the above is really just an inefficient way of implementing $x*x$. If you evaluate this program at $x = 3$, then the result is the constant $9$. But the derivative of constants is zero, which is obviously not the right answer, which should be $6$.

The reason is that AD will typically only look at the executed branch of your program. It is perfectly ok to have branches on conditions that don't involve numbers, or conditions involving number variables not part of the computational graph for derivative calculations. It's also ok to look at fixed properties of derivative values (e.g. the dimension of a vector). But as soon as you "look" at the contents of a variable that depends on one of the inputs to determine what calculation to perform next, you will "break the chain". AD does basically just that: apply the chain rule.

Users of AD should view programs as "configuring" a fixed computational graph. If running the program with different inputs for which derivatives are requested gives you different computational graphs, AD may not always give the correct result.

• So, if I understand well, branches and loops should should not depend on the content of variables.. To me, this can be reformulated as: "the program should be closed form".. or "the program could be rewritten with no loops" or "the program terminates in a known number of steps", am I right? This is quite a restricted subset of "every computer program", is it not? Feb 22, 2017 at 9:34
• Obviously, one can only calculate derivatives of a program that terminates. You do not need to know the number of steps beforehand. No matter how long it runs, eventually you will have a computational trace of everything that happened. AD (reverse-mode) looks only at the trace of operations that affect derivatives, inferring a computational graph from it. You can then interpret the graph backwards using the chain rule to obtain gradients. The important requirement is that the graph must remain the same even if you changed the value of arguments for which you want to calculate gradients. Feb 23, 2017 at 14:54
• So, there is nothing to do for derivable functions for which I can write a program that both terminates in all cases but whose graph is still not the same when I change the arguments, right? For instance: take a smooth function $f : \mathbb{R}^2 \rightarrow \mathbb{R}$ for which each point is the result of a limit or an integration operation. Then if I write an adaptive iterative program P(a, b) that approximates $f(a,\ b)$ with a control flow which depends on the given value of (a, b) (for adaptive precision or efficiency, say), P cannot be analysed with AD, right? Feb 23, 2017 at 15:40
• Lets assume you have an iterative optimization algorithm that checks a value for convergence. That's a problem, because if you pass arguments that are already optimal, there will be only one iteration of the algorithm, i.e. a small computational graph. In order to "fix" the graph in a meaningful way, you could pass an iteration limit rather than a convergence criterion. In my experience most algorithms can be rewritten such that AD will do the right thing. It's certainly not the case that you can throw AD at any program that wasn't written with AD in mind and always expect sound results. Feb 24, 2017 at 16:22
• Okaay! :D So this is how I understand it: if your program (should it be adaptive) does terminate in all cases, then there exist a finite upper bound to the number of steps it'll perform. So you may refactor it so that it performs exactly this many steps and follow the same graph, no matter the data given as an input.. should you fill up the graph with dummy operations in this purpose. Too bad we cannot enjoy the adaptiveness computational gain then.. but I think I've got it. Cheers :) Feb 24, 2017 at 16:32

If you want the derivative everywhere, automatic differentiation can't handle branches and loops. If you are satisfied with getting the derivative "almost everywhere", automatic differentiation might be fine for some programs with (some kinds of) branches. For some kinds of optimization, this is sometimes good enough.

Automatic differentiation can support at least some branches and if-statements, but with the caveat that it might fail (or the derivative might fail to exist) at the "boundary" of the conditional statements.

Consider your case where we branch if x_1 > x_2, and suppose we are calculating the derivative with respect to x_1 at a particular point (x_1, x_2). Then if the branch is true at that point, we calculate the derivative of the true-branch with respect to x_1. If the branch is false at that point, we calculate the derivative of the false-branch with respect to x_1. (If x_1 == x_2, we might have a discontinuity and the derivative might not be defined.)

Why does this work? Consider the function

$$f(x) = \begin{cases} g(x) &\text{if x>c}\\ h(x) &\text{otherwise} \end{cases}$$

Then its derivative is

$$f'(x) = \begin{cases} g'(x) &\text{if x>c}\\ h'(x) &\text{otherwise} \end{cases}$$

with a possible discontinuity at $x=c$.

For conditionals with an equality comparison, things might fail at the value where equality holds (at the boundary value). For conditionals with $\le$ or $<$ or $\ge$ or $>$, things might fail at the boundary/threshold value, but it works elsewhere, I think.

Looping is harder, I think.

Markus Mottl's answer is better than mine; see his answer for explanation of why branches and loops are problematic for automatic differentiation.

• Okay for "almost everywhere". This limitation sure sounds natural: AD will obviously fail to compute derivative where there is no derivative, who would blame that ;) I'll have a look at Markus Mottl's answer. Feb 22, 2017 at 8:45
• @iago-lito, yes, but the limitation is a bit worse than that: it might also fail to compute the correct derivative at the boundary/threshold points, even if a derivative does exist (i.e., a derivative might exist, but it might output the wrong value). That is a bit less desirable.
– D.W.
Feb 22, 2017 at 16:45
• That is a bit less desirable indeed. :\ Feb 22, 2017 at 16:47

As explained in the other responses, AD alone cannot capture the effects of input-dependent control flow. Regarding your second and third questions, there are some recent tools that extend AD to still compute useful gradients (cf. [1] and [2]). For a comprehensive treatment of the matter, you can also have a look at a recent book (preprint) on "Differentiable Programming" [3], in particular chapters 5 and 13.

Slightly simplifying the first example from the question, we may have a program like:

if x >= 0:
return 1
else:
return 0


This program implements the Heaviside step function $$H(x)$$. AD wrt. $$x$$ will follow only one of the two possible control flow paths (here returning a constant) and return its derivative (here 0). This is the correct derivative apart from $$x = 0$$, where $$H(x)$$ jumps from 0 to 1, unbeknownst to AD. Mathematically speaking, the derivative of $$H(0)$$ is the Dirac delta, whose properties (0 outside of $$x = 0$$ and with integral 1) make it its own can of worms.

Clearly, the zero-derivatives returned by plain AD are not particularly informative, and we'd prefer derivatives that reflect more than just one control flow path. There are some efforts that try to "fill in" jumps or kinks using, e.g., subderivatives. However, when using AD for optimization (as is commonly done with gradient descent), this is still not informative. Take again the function $$H(x)$$, where the derivative is still zero almost everywhere. One basic idea is to instead interpolate between the paths using some smooth function.

If we pick the logistic function $$l(x) = \frac{1}{1 + e^-kx}$$ with some smoothing factor $$k$$, we get:

return 1 / (1 + exp(-kx))


This is quite a common way to deal with discontinuous functions in machine learning. For this program, AD can compute a useful derivative:

The downside of plain interpolation is that the degree of smoothing is somewhat arbitrary. Particularly if there are multiple consecutive branches, it's not clear how well the derivative still reflects the original program's behavior. This is a problem if we'd like to use the derivative for optimization since the optima can shift unpredictably.

A better way is to take a probabilistic view of the program. This works either by adding randomness to the inputs (for programs without randomness like our example above), or making use of the program's own randomness (for programs where the relevant branches depend on pseudo-random numbers). In both cases, averaging over many program runs then leads to a smooth function.

For example, if we add normally distributed noise from $$\mathcal{N}(0, \sigma)$$ to the inputs of our example program, we get a smoothed function $$\hat{H}(x)$$:

Here, the two possible paths are weighted according to the normal distribution at the branching point. This is the core idea of the methods discussed below. As there is only one branch in this case, it happens to look quite similar to the interpolation above. However, this type of smoothing has a clear relation to the original program: it's the convolution of the program output with a Gaussian kernel, adjustable by varying $$\sigma$$. Again, for stochastic programs, we don't even need this external perturbation (although from experience, it can sometimes be beneficial for gradient descent).

What remains is to calculate the corresponding derivative. Note that it is not enough to average across the derivatives obtained for each sample by AD (this is known as infinitesimal perturbation analysis). For instance, in our smoothed Heaviside function, we'd be averaging over a series of zero-derivatives. Practically speaking, there are currently at least two relatively generic solutions to this. What they have in common is to sample across multiple runs of the program and then combine plain AD derivatives with some method to capture the probabilities and effects of branches.

One is the Julia package StochasticAD [1], which works for programs that draw from discrete probability distributions. While that doesn't cover branches in a generic manner, many programs can be expressed that way. For each program run (sample), StochasticAD tracks one "primal" and one alternative control flow path, applying a custom chain rule along the way. The latter builds on the parametric forms of the involved (discrete) distributions, which are assumed to be known a priori in this case.

For C++ programs, there is DiscoGrad [2] (note: our own work), which handles generic "if-else" branching statements by computing density estimations over the conditions collected on the fly. This is necessary because in the completely generic case one cannot readily rely on parametric forms for the distributions. The smoothed derivative is then computed in a post-processing step.

Given the capability to differentiate across branches, loops can just be regarded as a special case, by moving the branch condition inside the loop body:

while x1 < x2:
some_operations()

while True:
if x1 < x2:
some_operations()
else
break


Note, however, that this still requires the program to terminate. This can for example be ensured by setting a limit for the number of iterations.