You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Copy file name to clipboardExpand all lines: usage/automatic-differentiation/index.qmd
+84-44Lines changed: 84 additions & 44 deletions
Original file line number
Diff line number
Diff line change
@@ -12,68 +12,82 @@ using Pkg;
12
12
Pkg.instantiate();
13
13
```
14
14
15
-
## Switching AD Modes
15
+
## What is Automatic Differentiation?
16
16
17
-
Turing currently supports four automatic differentiation (AD) backends for sampling: [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) for forward-mode AD; and [Mooncake](https://github.com/compintell/Mooncake.jl) and [ReverseDiff](https://github.com/JuliaDiff/ReverseDiff.jl) for reverse-mode AD.
18
-
`ForwardDiff` is automatically imported by Turing. To utilize `Mooncake` or `ReverseDiff` for AD, users must explicitly import them with `import Mooncake` or `import ReverseDiff`, alongside the usual `using Turing`.
17
+
Automatic differentiation (AD) is a technique used in Turing.jl to evaluate the gradient of a function at a given set of arguments.
18
+
In the context of Turing.jl, the function being differentiated is the log probability density of a model, and the arguments are the parameters of the model (i.e. the values of the random variables).
19
+
The gradient of the log probability density is used by various algorithms in Turing.jl, such as HMC (including NUTS), mode estimation (which uses gradient-based optimization), and variational inference.
19
20
20
-
As of Turing version v0.30, the global configuration flag for the AD backend has been removed in favour of [`AdTypes.jl`](https://github.com/SciML/ADTypes.jl), allowing users to specify the AD backend for individual samplers independently.
21
-
Users can pass the `adtype` keyword argument to the sampler constructor to select the desired AD backend, with the default being `AutoForwardDiff(; chunksize=0)`.
21
+
The Julia ecosystem has a number of AD libraries.
22
+
Behind the scenes, Turing.jl uses DifferentiationInterface.jl as a unified interface to many of these libraries: the list of available backends can be seen on the [DifferentiationInterface documentation](https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/).
23
+
You can switch between these using the unified [ADTypes.jl](https://github.com/SciML/ADTypes.jl/) interface, which for a given AD backend, provides types such as `AutoBackend` (see [the documentation](https://docs.sciml.ai/ADTypes/stable/) for more details).
24
+
For example, to use the [Mooncake.jl](https://github.com/compintell/Mooncake.jl) package for AD, you can run the following:
22
25
23
-
For `ForwardDiff`, pass `adtype=AutoForwardDiff(; chunksize)` to the sampler constructor. A `chunksize` of `nothing` permits the chunk size to be automatically determined. For more information regarding the selection of `chunksize`, please refer to [related section of `ForwardDiff`'s documentation](https://juliadiff.org/ForwardDiff.jl/dev/user/advanced/#Configuring-Chunk-Size).
26
+
```{julia}
27
+
# Turing re-exports AutoForwardDiff, AutoReverseDiff, and AutoMooncake.
28
+
# Other ADTypes must be explicitly imported from ADTypes.jl or
29
+
# DifferentiationInterface.jl.
30
+
using Turing
31
+
setprogress!(false)
24
32
25
-
For `ReverseDiff`, pass `adtype=AutoReverseDiff()` to the sampler constructor. An additional keyword argument called `compile` can be provided to `AutoReverseDiff`. It specifies whether to pre-record the tape only once and reuse it later (`compile` is set to `false` by default, which means no pre-recording). This can substantially improve performance, but risks silently incorrect results if not used with care.
33
+
# Note that if you specify a custom AD backend, you must also import it.
34
+
import Mooncake
26
35
27
-
Pre-recorded tapes should only be used if you are absolutely certain that the sequence of operations performed in your code does not change between different executions of your model.
36
+
@model function f()
37
+
x ~ Normal()
38
+
# Rest of your model here
39
+
end
28
40
29
-
Thus, e.g., in the model definition and all implicitly and explicitly called functions in the model, all loops should be of fixed size, and `if`-statements should consistently execute the same branches.
30
-
For instance, `if`-statements with conditions that can be determined at compile time or conditions that depend only on fixed properties of the model, e.g. fixed data.
31
-
However, `if`-statements that depend on the model parameters can take different branches during sampling; hence, the compiled tape might be incorrect.
32
-
Thus you must not use compiled tapes when your model makes decisions based on the model parameters, and you should be careful if you compute functions of parameters that those functions do not have branching which might cause them to execute different code for different values of the parameter.
The previously used interface functions including `ADBackend`, `setadbackend`, `setsafe`, `setchunksize`, and `setrdcache` have been removed.
44
+
By default, if you do not specify a backend, Turing will default to [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl).
45
+
In this case, you do not need to import ForwardDiff, as it is already a dependency of Turing.
35
46
36
-
For `Mooncake`, pass `adtype=AutoMooncake(; config=nothing)` to the sampler constructor.
47
+
## Choosing an AD Backend
37
48
38
-
## Compositional Sampling with Differing AD Modes
49
+
There are two aspects to choosing an AD backend: firstly, what backends are available; and secondly, which backend is best for your model.
39
50
40
-
Turing supports intermixed automatic differentiation methods for different variable spaces. The snippet below shows using `ForwardDiff` to sample the mean (`m`) parameter, and using `ReverseDiff` for the variance (`s`) parameter:
51
+
### Usable AD Backends
41
52
42
-
```{julia}
43
-
using Turing
44
-
using ReverseDiff
53
+
Turing.jl uses the functionality in [DifferentiationInterface.jl](https://github.com/JuliaDiff/DifferentiationInterface.jl) ('DI') to interface with AD libraries in a unified way.
54
+
Thus, in principle, any AD library that has integrations with DI can be used with Turing; you should consult the [DI documentation](https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/) for an up-to-date list of compatible AD libraries.
45
55
46
-
# Define a simple Normal model with unknown mean and variance.
47
-
@model function gdemo(x, y)
48
-
s² ~ InverseGamma(2, 3)
49
-
m ~ Normal(0, sqrt(s²))
50
-
x ~ Normal(m, sqrt(s²))
51
-
return y ~ Normal(m, sqrt(s²))
52
-
end
56
+
Note, however, that not all AD libraries in there are thoroughly tested on Turing models.
57
+
Thus, it is possible that some of them will either error (because they don't know how to differentiate through Turing's code), or maybe even silently give incorrect results (if you are very unlucky).
58
+
Turing is most extensively tested with **ForwardDiff.jl** (the default), **ReverseDiff.jl**, and **Mooncake.jl**.
59
+
We also run a smaller set of tests with Enzyme.jl.
53
60
54
-
# Sample using Gibbs and varying autodiff backends.
Before describing how to choose the best AD backend for your model, we should mention that we also publish a table of benchmarks for various models and AD backends in [the ADTests website](https://turinglang.org/ADTests/).
64
+
These models aim to capture a variety of different features of Turing.jl and Julia in general, so that you can see which AD backends may be compatible with your model.
65
+
Benchmarks are also included, although it should be noted that many of the models in ADTests are small and thus the timings may not be representative of larger, real-life models.
66
+
67
+
If you have suggestions for other models to include, please do let us know by [creating an issue on GitHub](https://github.com/TuringLang/ADTests/issues/new)!
65
68
66
-
Generally, reverse-mode AD, for instance `ReverseDiff`, is faster when sampling from variables of high dimensionality (greater than 20), while forward-mode AD, for instance `ForwardDiff`, is more efficient for lower-dimension variables. This functionality allows those who are performance sensitive to fine tune their automatic differentiation for their specific models.
69
+
### The Best AD Backend for Your Model
67
70
68
-
If the differentiation method is not specified in this way, Turing will default to using whatever the global AD backend is.
69
-
Currently, this defaults to `ForwardDiff`.
71
+
Given the number of possible backends, how do you choose the best one for your model?
70
72
71
-
The most reliable way to ensure you are using the fastest AD that works for your problem is to benchmark them using the functionality in DynamicPPL (see [the API documentation](https://turinglang.org/DynamicPPL.jl/stable/api/#AD-testing-and-benchmarking-utilities)):
73
+
A simple heuristic is to look at the number of parameters in your model.
74
+
The log density of the model, i.e. the function being differentiated, is a function that goes from $\mathbb{R}^n \to \mathbb{R}$, where $n$ is the number of parameters in your model.
75
+
For models with a small number of parameters (say up to ~ 20), forward-mode AD (e.g. ForwardDiff) is generally faster due to a smaller overhead.
76
+
On the other hand, for models with a large number of parameters, reverse-mode AD (e.g. ReverseDiff or Mooncake) is generally faster as it computes the gradients with respect to all parameters in a single pass.
77
+
78
+
The most exact way to ensure you are using the fastest AD that works for your problem is to benchmark them using the functionality in DynamicPPL (see [the API documentation](https://turinglang.org/DynamicPPL.jl/stable/api/#AD-testing-and-benchmarking-utilities)):
72
79
73
80
```{julia}
81
+
using ADTypes
74
82
using DynamicPPL.TestUtils.AD: run_ad, ADResult
75
83
using ForwardDiff, ReverseDiff
76
84
85
+
@model function gdemo(x, y)
86
+
s² ~ InverseGamma(2, 3)
87
+
m ~ Normal(0, sqrt(s²))
88
+
x ~ Normal(m, sqrt(s²))
89
+
return y ~ Normal(m, sqrt(s²))
90
+
end
77
91
model = gdemo(1.5, 2)
78
92
79
93
for adtype in [AutoForwardDiff(), AutoReverseDiff()]
@@ -84,6 +98,32 @@ end
84
98
85
99
In this specific instance, ForwardDiff is clearly faster (due to the small size of the model).
86
100
87
-
We also have a table of benchmarks for various models and AD backends in [the ADTests website](https://turinglang.org/ADTests/).
88
-
These models aim to capture a variety of different Turing.jl features.
89
-
If you have suggestions for things to include, please do let us know by [creating an issue on GitHub](https://github.com/TuringLang/ADTests/issues/new)!
101
+
::: {.callout-note}
102
+
## A note about ReverseDiff's `compile` argument
103
+
104
+
The additional keyword argument `compile=true` for `AutoReverseDiff` specifies whether to pre-record the tape only once and reuse it later.
105
+
By default, this is set to `false`, which means no pre-recording.
106
+
Setting `compile=true` can substantially improve performance, but risks silently incorrect results if not used with care.
107
+
Pre-recorded tapes should only be used if you are absolutely certain that the sequence of operations performed in your code does not change between different executions of your model.
108
+
:::
109
+
110
+
## Compositional Sampling with Differing AD Modes
111
+
112
+
When using Gibbs sampling, Turing also supports mixed automatic differentiation methods for different variable spaces.
113
+
The following snippet shows how one can use `ForwardDiff` to sample the mean (`m`) parameter, and `ReverseDiff` for the variance (`s`) parameter:
114
+
115
+
```{julia}
116
+
using Turing
117
+
using ReverseDiff
118
+
119
+
# Sample using Gibbs and varying autodiff backends.
0 commit comments