6 ms·

> The Enzyme project is a tool for performing reverse-mode automatic differentiation (AD) of statically-analyzable LLVM IR. This allows developers to use Enzyme

by ksr 6y ago

> The Enzyme project is a tool for performing reverse-mode automatic differentiation (AD) of statically-analyzable LLVM IR. This allows developers to use Enzyme to automatically create gradients of their source code without much additional work. Can someone please explain applications of creating gradients of my source code?

6d65 6y ago
AFAIK, It's mainly used for implementing gradient descent, which is used for training neural networks. Frameworks like pytorch, tensorflow, probably used back propagation to calculate the gradient of a multidimensional function. But in involves tracing, and storing the network state during the forward pass. Static automatic differentiation should be faster and should look a lot like differentiation is done mathematically rather than numerically. Of course there are more applications to AD in scientific computing.
- oscargrouch 6y ago
  I think Swift is also going into this direction of backing in directly into the compiler and provide it as a higher level language construction. https://github.com/apple/swift/blob/main/docs/DifferentiableProgramming.md https://github.com/apple/swift/blob/main/docs/Differentiable... Which leads to "Swift for Tensorflow" that unlike other languages like Java, Go or Python is not just about bindings to the C++ tensorflow library.
  - 6d65 6y ago
    You're right about Swift for Tensorflow. It is an interesting development. I have installed swift on Linux just to play with it. Was dissapointed to see that Swift for Tensorflow was a fork of the Swift compiler, and haven't followed whether it's being refreshed, and whether there are any plans to merge it back. Ideally an LLVM tool should allow languages that compile to LLVM(I'm mostly interested in Rust and Julia) to leverage it without dramatically changing their compilers. It would be interesting to see something like JAX in Rust, exposing the AD functionality in the Standard Library, paired with a high performance SIMD/GPU array type. Things could get very interesting.
- cl3misch 6y ago
  I don't see how static AD removes the need to store the network state. Is this a fundamental property of static AD? Also, your statement sounds like pyTorch/TF are doing AD numerically, which is not the case. They build the analytical gradient from the traced computation graph.
  - wsmoses 6y ago
    Reverse mode AD can always get into situations where it needs to store original values (i.e. network state). One advantage, however, of doing a more whole-program approach to AD rather than individual operators is that one might be able to avoid caching values unnecessarily. For example if an input isn't modified (and still exists) by the time the value is needed in the reverse pass, you don't need to cache it but can simply use the original input without a copy. And yes PyTorch/TF tend to perform a (limited) form of AD as well, rather than do numerical differentiation (though I do think there may be an option for numerical?) I wouldn't really position a tool like Enzyme as a competitor to PyTorch/TF (they may have some better domain-specific knowledge after all), but rather a really nice complement. Enzyme can take derivatives of arbitrary functions, in any LLVM-based language rather than the DSL of operators supported by PyTorch/TF. In fact, we built a plugin for PyTorch/TF that uses Enzyme to import custom foreign code as a differentiable layer!
    - cl3misch 6y ago
      > One advantage, however, of doing a more whole-program approach to AD rather than individual operators I was under the impression that the big ML frameworks (and surely JAX with jit) are doing optimization on the complete compute graph, too. I didn't want to make this discussion too TF/pyTorch focused (I'm not even a ML researcher). But your optimization claims sound like the other AD frameworks are not doing any optimization at all, which is not the case. I was also thinking about derivatives of functions which are doing something iterative on the inside, like a matrix decomposition (combined with linear solve and/or matrix inversion). While a "high level" AD tracer can identify an efficient derivative of these operations, your LLVM introspection would only be able to compute the derivative through all the internal step of the matrix decomposition?
ant6n 6y ago
Optimization. Then again, one could probably calculate gradients numerically.
- cameronperot 6y ago
  One could, but automatic differentiation is much more efficient than numerical differentiation, thus for high performance applications it is preferable to use automatic differentiation.
  - wsmoses 6y ago
    Adding onto this, numerical derivatives have two potential problems which is why they tend not to be used in big scientific/ML frameworks. First of all they suffer from accuracy decay. For example if you were to do the standard f'(x) \approx [f(x+h)-f(x)]/h, you'd subtract two similar numbers and waste many bits of precision. In contrast if you were to generate the derivative function directly like below, you'd end up far more accurate. double square(double x) { return x * x; } double d_square(double x) { return __enzyme_autodiff(square, x); } becomes double d_square(double x) { return 2 * x; } Secondly, from a performance perspective numerical differentiation is really slow -- especially for gradient computation. For example you would need to evaluate the function at once per argument in numeric differentiation to get the whole gradient. In contrast, reverse mode AD lets you generate the entire gradient in one call. In addition to these generic issues, we illustrate in our paper how doing this at a compiler level allows for significant additional optimization (by removing unnecessary code from the forward pass, finding common expressions, etc). These issues also are amplified as you make higher-order derivatives and so on.
PartiallyTyped 6y ago
It constructs an analytical gradient from the code. The reason is that you can compute the gradient directly. This can enable optimizations such as avoiding caching big matrices because you don't need to keep track of states/trace the graph, or you can compute the 2nd, 3rd, 4th... and so on derivatives because you have an analytical gradient. For example in an affine function, the gradient of the bias/intercept is the gradient of the loss wrt the activation function and for the weights, it's the product of loss wrt activation function and the input to the layer. With automatic graph construction e.g. eager Tensorflow/Pytorch, the layer needs to cache the input of the layer, so that it can compute the gradient of the weights. If the layer receives inputs multiple times within the computation graph, you end up caching it multiple times. With analytical gradients, you may be able to save memory by finding optimizations because you have the analytical gradient, e.g. above you can sum the inputs ie (dL/dz)input1 + (dL/dz)input2 = (dL/dz)(input1+input2).
- mhh__ 6y ago
  That last line doesn't make much sense to my eye at least.
  - wsmoses 6y ago
    I think in essence what PartiallyTyped is trying to say is that one potential optimization opportunity in whole-program AD is that you can avoid having to cache the original inputs of the program if you know that derivative computation won't need it (e.g. its only used in a sum and not a product or something whose derivative depends on the value). Some ML frameworks must cache all of the inputs to an operation since they don't know whether it will be necessary for the reverse pass of an operation. You could go even further and decide to cache a different & smaller set of intermediate values that still lets you compute the gradient. Beyond cache reduction, in our paper we demonstrate a lot of interesting ways that combining AD with a compiler can create potential speed-up. For example, we are often able to dead-code eliminate part of the original forward-pass code since it's not needed to compute the gradient. We also have a cool example in the paper showing an asymptotic [O(N^2) => O(N)] speedup on a code for normalizing a vector because doing AD in the compiler lets Enzyme run after optimization (and in that example benefit from loop invariant code motion in a way that tools that aren't in the compiler cannot do).
deleted 6y ago
[deleted]
wsmoses 6y ago
Say you have some existing virus simulation codebase that you want to use ML on to derive an effective policy on. Without an AD tool like Enzyme, you'd have to spend significant time and effort understanding and rewriting that obnoxious 100K lines of fortran into TensorFlow, when you could've been spending it solving your problem. The reason you need to do this rewriting is because many ML algorithms require the derivatives of functions to be able to use them and Enzyme provides an easy way to generate derivatives of existing code. This is also useful in the scientific world where derivatives of functions are commonplace. You could also use it in more performance-engineering/computer systems ways as well by using the derivatives to perform uncertainty quantification and perhaps decide to use 32-bit floats rather than 64-bit doubles.
- shoyer 6y ago
  I'm a big believer in auto-diff, but I'm skeptical that any autodiff tool would differentiate a 100k line simulation code correctly and efficiently without manual intervention. I'd certainly love to be proven wrong, though and absolutely AD can be a big time saver :)
  - wsmoses 6y ago
    Whoops added one too many zero’s there, agreed that would be really nice :P

by ksr 6y ago

6d65 6y ago

AFAIK, It's mainly used for implementing gradient descent, which is used for training neural networks. Frameworks like pytorch, tensorflow, probably used back propagation to calculate the gradient of a multidimensional function. But in involves tracing, and storing the network state during the forward pass. Static automatic differentiation should be faster and should look a lot like differentiation is done mathematically rather than numerically. Of course there are more applications to AD in scientific computing.

oscargrouch 6y ago
I think Swift is also going into this direction of backing in directly into the compiler and provide it as a higher level language construction. https://github.com/apple/swift/blob/main/docs/DifferentiableProgramming.md https://github.com/apple/swift/blob/main/docs/Differentiable... Which leads to "Swift for Tensorflow" that unlike other languages like Java, Go or Python is not just about bindings to the C++ tensorflow library.
- 6d65 6y ago
  You're right about Swift for Tensorflow. It is an interesting development. I have installed swift on Linux just to play with it. Was dissapointed to see that Swift for Tensorflow was a fork of the Swift compiler, and haven't followed whether it's being refreshed, and whether there are any plans to merge it back. Ideally an LLVM tool should allow languages that compile to LLVM(I'm mostly interested in Rust and Julia) to leverage it without dramatically changing their compilers. It would be interesting to see something like JAX in Rust, exposing the AD functionality in the Standard Library, paired with a high performance SIMD/GPU array type. Things could get very interesting.
cl3misch 6y ago
I don't see how static AD removes the need to store the network state. Is this a fundamental property of static AD? Also, your statement sounds like pyTorch/TF are doing AD numerically, which is not the case. They build the analytical gradient from the traced computation graph.
- wsmoses 6y ago
  Reverse mode AD can always get into situations where it needs to store original values (i.e. network state). One advantage, however, of doing a more whole-program approach to AD rather than individual operators is that one might be able to avoid caching values unnecessarily. For example if an input isn't modified (and still exists) by the time the value is needed in the reverse pass, you don't need to cache it but can simply use the original input without a copy. And yes PyTorch/TF tend to perform a (limited) form of AD as well, rather than do numerical differentiation (though I do think there may be an option for numerical?) I wouldn't really position a tool like Enzyme as a competitor to PyTorch/TF (they may have some better domain-specific knowledge after all), but rather a really nice complement. Enzyme can take derivatives of arbitrary functions, in any LLVM-based language rather than the DSL of operators supported by PyTorch/TF. In fact, we built a plugin for PyTorch/TF that uses Enzyme to import custom foreign code as a differentiable layer!
  - cl3misch 6y ago
    > One advantage, however, of doing a more whole-program approach to AD rather than individual operators I was under the impression that the big ML frameworks (and surely JAX with jit) are doing optimization on the complete compute graph, too. I didn't want to make this discussion too TF/pyTorch focused (I'm not even a ML researcher). But your optimization claims sound like the other AD frameworks are not doing any optimization at all, which is not the case. I was also thinking about derivatives of functions which are doing something iterative on the inside, like a matrix decomposition (combined with linear solve and/or matrix inversion). While a "high level" AD tracer can identify an efficient derivative of these operations, your LLVM introspection would only be able to compute the derivative through all the internal step of the matrix decomposition?

ant6n 6y ago

Optimization. Then again, one could probably calculate gradients numerically.

cameronperot 6y ago
One could, but automatic differentiation is much more efficient than numerical differentiation, thus for high performance applications it is preferable to use automatic differentiation.
- wsmoses 6y ago
  Adding onto this, numerical derivatives have two potential problems which is why they tend not to be used in big scientific/ML frameworks. First of all they suffer from accuracy decay. For example if you were to do the standard f'(x) \approx [f(x+h)-f(x)]/h, you'd subtract two similar numbers and waste many bits of precision. In contrast if you were to generate the derivative function directly like below, you'd end up far more accurate. double square(double x) { return x * x; } double d_square(double x) { return __enzyme_autodiff(square, x); } becomes double d_square(double x) { return 2 * x; } Secondly, from a performance perspective numerical differentiation is really slow -- especially for gradient computation. For example you would need to evaluate the function at once per argument in numeric differentiation to get the whole gradient. In contrast, reverse mode AD lets you generate the entire gradient in one call. In addition to these generic issues, we illustrate in our paper how doing this at a compiler level allows for significant additional optimization (by removing unnecessary code from the forward pass, finding common expressions, etc). These issues also are amplified as you make higher-order derivatives and so on.

PartiallyTyped 6y ago

It constructs an analytical gradient from the code. The reason is that you can compute the gradient directly. This can enable optimizations such as avoiding caching big matrices because you don't need to keep track of states/trace the graph, or you can compute the 2nd, 3rd, 4th... and so on derivatives because you have an analytical gradient. For example in an affine function, the gradient of the bias/intercept is the gradient of the loss wrt the activation function and for the weights, it's the product of loss wrt activation function and the input to the layer. With automatic graph construction e.g. eager Tensorflow/Pytorch, the layer needs to cache the input of the layer, so that it can compute the gradient of the weights. If the layer receives inputs multiple times within the computation graph, you end up caching it multiple times. With analytical gradients, you may be able to save memory by finding optimizations because you have the analytical gradient, e.g. above you can sum the inputs ie (dL/dz)input1 + (dL/dz)input2 = (dL/dz)(input1+input2).

mhh__ 6y ago
That last line doesn't make much sense to my eye at least.
- wsmoses 6y ago
  I think in essence what PartiallyTyped is trying to say is that one potential optimization opportunity in whole-program AD is that you can avoid having to cache the original inputs of the program if you know that derivative computation won't need it (e.g. its only used in a sum and not a product or something whose derivative depends on the value). Some ML frameworks must cache all of the inputs to an operation since they don't know whether it will be necessary for the reverse pass of an operation. You could go even further and decide to cache a different & smaller set of intermediate values that still lets you compute the gradient. Beyond cache reduction, in our paper we demonstrate a lot of interesting ways that combining AD with a compiler can create potential speed-up. For example, we are often able to dead-code eliminate part of the original forward-pass code since it's not needed to compute the gradient. We also have a cool example in the paper showing an asymptotic [O(N^2) => O(N)] speedup on a code for normalizing a vector because doing AD in the compiler lets Enzyme run after optimization (and in that example benefit from loop invariant code motion in a way that tools that aren't in the compiler cannot do).

deleted 6y ago

[deleted]

wsmoses 6y ago

Say you have some existing virus simulation codebase that you want to use ML on to derive an effective policy on. Without an AD tool like Enzyme, you'd have to spend significant time and effort understanding and rewriting that obnoxious 100K lines of fortran into TensorFlow, when you could've been spending it solving your problem. The reason you need to do this rewriting is because many ML algorithms require the derivatives of functions to be able to use them and Enzyme provides an easy way to generate derivatives of existing code. This is also useful in the scientific world where derivatives of functions are commonplace. You could also use it in more performance-engineering/computer systems ways as well by using the derivatives to perform uncertainty quantification and perhaps decide to use 32-bit floats rather than 64-bit doubles.

shoyer 6y ago
I'm a big believer in auto-diff, but I'm skeptical that any autodiff tool would differentiate a 100k line simulation code correctly and efficiently without manual intervention. I'd certainly love to be proven wrong, though and absolutely AD can be a big time saver :)
- wsmoses 6y ago
  Whoops added one too many zero’s there, agreed that would be really nice :P