10
$\begingroup$

Given a model $f$ as a programming routine, such that we are able to compute $y=f(x)$ for any $x \in \mathcal{D}$, I am interested in the sensitivity (or let us say derivative) of $y$ with respect to $x$. The important thing is that $f$ could be any model, for instance a machine learning model, a neural network or a model from a commercial software, and that this model should always have the ability to compute the output for any given input.

Of course, knowing more about the model will give more hints about methods to be used for sensitivity analysis. But talking generally and in an abstract fashion, one could for example perturb a point $x$ to $x+\epsilon$ and compute $\frac{dy}{dx}(x) = \frac{f(x+\epsilon) - f(x)}{\epsilon}$. One could also use a central difference scheme.

Are you aware of any other methods to perform this task ?

$\endgroup$
2
  • $\begingroup$ Automatic differentiation (autodiff) packages can accomplish this task without the finite difference approximation to the derivative. $\endgroup$ Aug 14 at 13:11
  • $\begingroup$ @sssssssssssss, I agree but isn't it that the autodiff restricted to the presence of a computational graph from the inputs to the outputs? Which is true for lets say TensorFlow models (neural networks). But if you go on a more abstract level and think of any model, will this still work? $\endgroup$
    – outlaw
    Aug 14 at 13:17
20
$\begingroup$

The act of using automatic differentiation for derivative calculations of general programs has become known as differentiable programming. There are many grades of differentiable programming across the spectrum, with some being confined to only models that could be represented statically, to systems like in Julia (and attempted in Swift) which allow for general differentiation of programs written within the language.

The mathematics for how this is accomplished is to understand automatic differentiation as simply the composition of Jacobian-vector or Jacobian-transpose-vector products. The MIT 18.337 Scientific Machine Learning course notes go into detail on this, specifically sections on forward mode, reverse mode, and finally differentiable programming. The crux of the idea is that you want to calculate $Jv$ where $J$ is the Jacobian $f'(x)$ and $v$ is some vector, since $Jv$ is equivalent to a directional derivative. If your compute program is two function calls, $g(f(x))$, then simply by linear algebra you have $Jv = J_g (J_f v)$, and so you can calculate the Jacobian-vector product of a whole program by doing it along the various steps. Forward-mode automatic differentiation can be done by defining a dual number arithmetic $x = a + b \epsilon$ (think about it as a two dimensional number, like a complex number), where $f(x) = f(a) + f'(a)b \epsilon$. If you define the arithmetic on a set of primitives, say $+$,$*$,$\sin$, etc., and then your program is composed of said primitives, this will differentiate entire generic programs. This is the basis of systems like ForwardDiff.jl. The Julia programming language then for example has the library ChainRules.jl which allows for registering primitive derivative definitions of any function in the language (those in Julia Base or in any package) for AD systems to in its derivative program construction. See Lyndon White's video on defining primitives for ChainRules.jl as a nice introduction.

While reverse mode automatic differentiation for handling general dynamic programs is much more complex, it can be done. Reverse mode AD is equivalent to the calculation of $J^Tv$, or $v^T J$, and so by linear algebra you can see the composition is reversed: $v^T J = (v^T J_g)J_f$ which means you have to figure out how to take the forward path of a program and then pull a vector transposed reverse through each of the operations. Notice that if $f(g(x))$ returns a scalar and $v$ is one, then $v^T J = \nabla f$ the gradient, which is the relationship to "backpropagation" and the use in machine learning. In reverse mode AD, the primitives are these $v^T J$ operations, and then any program composed of primitives is differentiable. A purely interpreted form of differentiable programming exists, known as tracing-based AD, where a single run of a program $f(x)$ builds out a tape of the operations that are used during that run, and that static description of the forward program at $x$ is then differentiated. Because this requires a new code for every reverse pass, you cannot generally compile the pass (unless the runtime is known to be much longer than the compile time), you lose optimizations in the forward pass (since you have to dynamically allocate a call tree), and many constructs cannot be well supported because lacking general analyses for escaping and aliasing (for example, mutation). For these reasons, machine learning frameworks which generally focus performance in large kernel calls (like matrix multiplication) use this technique and a notable software here is PyTorch. This handles some forms of dynamism but with major limitations for smaller cost function calls (e.g. nonlinear functions defined by scalar operations like is commonly seen in differential equation definitions), but if used in the correct domain that can be ignored.

A more general form of automatic differentiation can be done via source code transformations on the lowered compiler representation known as SSA (single static assignment). Mike Innes' writeup of SSA-level differentiation is a good introduction, and Keno Fischer's talk on Optics from Category Theory being used in Diffractor is a nice example of doing it deeper within the compiler, while Billy and Valentin's discussion of Enzyme discuss the requirements of composability with compiler transformations. From these pieces together, you can notice that to do this both correctly and efficiently you need to mix compiler analysis idea like escape analysis, dead code elimination for thunks, etc. directly into the AD code generation mechanism, which is why this is done in a style that interacts with the compiler.

The interesting thing about the SSA form is that, like Mike's paper alludes to in its title "Don't Unroll Adjoint", SSA already has an O(1) representation of dynamic constructs like loops which it can keep intact. For example, if you think about the while loop while x < 1, the derivative program can replace this with the same check proceeded by an act to add a boolean to a stack. The reverse mode pass is then a for loop for i in 1:length(stack), when then runs the internals of the while loop in reverse, replacing each operation with its vector-transposed-Jacobian-product. And you can similarly analyze general programming constructs to define their primitives. What's the reverse mode derivative of locking a thread? It turns out that it's to lock a thread, because when you run a "lock f(x) unlock" program in reverse, you want to perform "lock $v^T f'(x)$ unlock" program in reverse.

So with this mixture of things like ChainRules.jl for defining language-wide primitives, AD systems which act directly on the language's lowered intermediate representation (Zygote.jl, Diffractor.jl, or even lower on the LLVM IR Enzyme.jl), you can design a language system to be fully differentiable on programs without ever requiring that they opt into the system. Of course, each AD system has its caveats due to its design constraints (something I discuss in much more detail in a blog post on gluing AD systems and another on limitations of quasi-static graph representations), so it's still a topic of ongoing research to make it so literally every program can just call derivative(f,x) and expect it to work. But it's both not impossible and it's something that entire communities are working on.

Let me end by mentioning that thinking about mathematical transformations as compiler passes can have many uses beyond generating codes for derivatives. It can detect sparsity patterns, make numerical models more stable, reduce floating point operations, generate code for "intrusive" uncertainty quantification methods, and more. The general concept here is abstract interpretation of functions for mathematical purposes, and that's really the key principle that ties together all of these compiler techniques.

[And just as an aside because I know someone will ask, Why Swift For TensorFlow gives a nice insight into why Google was looking at Swift and Julia for differentiable programming but not Python. Though IMO that discussion is somewhat indirect. A more direct reason is that for this to work, you need to be working on a program representation that (a) has an SSA-form IR that is amenable to inference and optimization (since that's what you will be doing code generation from and to) and (b) has a minimal number of primitives to support. So something where most of the libraries are written in the language and where its representations are static (Swift, Fortran, C++) or almost static with a language designs made to ensure compiler analyses (Julia). Python is a fantastic language for many things because its dynamism allows many natural structures, but its intermediate representation is not made for compilers to fully analyze (for example, objects can add fields at any time, so you need a non-local analysis to know which fields to include when constructing an object in reverse) and it generally uses a lot of calls to C libraries, each of which would need AD primitives to get support (which can decrease numerical stability, but that's a topic for another time). That's why machine learning frameworks like PyTorch have focused on subsets of the language and have provided their own C libraries (torch.numpy is not numpy) to minimize the primitive support required.]

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.