EnzymeAD / EnzymeAD/Enzyme.jl

Supporting covariant derivatives

Open
#1,334 2 comments 4 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

## The problem:

Currently, when you compute the reverse mode derivative of some function of a `struct MyStruct` in enzyme, i.e. `autodiff(Reverse, f, Duplicated(A::MyStruct, dA::MyStruct))`, the object `dA` must be an identical type to `A`, but it should be interpreted in a very different way from `A`. What `dA` actually is, is the object you get by treating `MyStruct` as a Cartesian vector whose elements are the struct's fields.

The `i`th field of `dA` will be given by $\mathrm{d}A_{i} = {\partial f / \partial A_{i}}$

The problem arises though, that often what people want is not this component-wise Cartesian derivative, but a [**covariant** derivative](https://en.wikipedia.org/wiki/Covariant_derivative), i.e. the mathematically meaningful derivative on the manifold. (Note: I consider the type `MyStruct` to define a manifold whose metric could be constructed via the constructor for `MyStruct`, and I consider `f` to be a function on that manifold).

Here's a concrete example posted by @simsurace in Slack:

```julia
julia> using Enzyme, LinearAlgebra

julia> A = Symmetric(rand(3, 3)); dA = make_zero(A);

julia> Enzyme.autodiff(Reverse, sum, Duplicated(A, dA))
((nothing,),)

julia> dA
3×3 Symmetric{Float64, Matrix{Float64}}:
1.0 2.0 2.0
2.0 1.0 2.0
2.0 2.0 1.0
```
compared to the answer given by Zygote:
```julia
julia> using Zygote

julia> Zygote.gradient(sum, A) |> only
3×3 Symmetric{Float64, FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}:
1.0 1.0 1.0
1.0 1.0 1.0
1.0 1.0 1.0
```

What's going on here? Enzyme is treating `A::Symmetric` as a vector with only one differentiable field (`A.data`), and so it computes

$$
dA_{\mathrm{data}}[a, b] = {\partial \over \partial A_{\mathrm{data}}[a, b]}\sum_{i,j} A[i,j] \\
= {\partial \over \partial A_{\mathrm{data}}[a, b]} \sum_{i,j}
\begin{cases}
A_{\mathrm{data}}[i, j] & j \geq i \\
A_{\mathrm{data}}[j, i], & \text{otherwise}
\end{cases}
$$

Now, since `A.data[i,j]` is never accessed for `j < i`, this evaluates to

$$
dA_{\mathrm{data}}[a, b] = \sum_{i,j} \begin{cases}
\delta_{a, i} \delta_{b,j} & j \geq i \\
\delta_{a, j} \delta_{b,i}, & \text{otherwise}
\end{cases}
= \sum_{i} \delta_{a,i} \delta_{b,i} + \sum_{i, j>i} 2 \delta_{a,i}\delta_{b,j}
= \begin{cases}
1, & a = b \\
2, & a < b \\
0, & \text{otherwise}
\end{cases}
$$

meaning

```julia
dA.data = [1.0 2.0 2.0
0.0 1.0 2.0
0.0 0.0 1.0]
```
The problem is that `dA` itself is a `Symmetric`, so what we end up getting is
```julia
julia> dA
3×3 Symmetric{Float64, Matrix{Float64}}:
1.0 2.0 2.0
2.0 1.0 2.0
2.0 2.0 1.0
```
`dA` has the incorrect mathematical properties if we want think of it as giving an actual derivative. If we interpret this thing as a gradient, it would suggest something incorrect about the slopes of the function `sum` over the space of `Symmetric` 3x3 matrices. If you check, you'll find that the slope of `sum` on this space in every `[i,j]` direction is uniformly `1` (the answer given by `Zygote`). Basically, Zygote / ChainRules try to transform the derivative to then live in the [(co)tangent space](https://en.wikipedia.org/wiki/Tangent_space), whereas Enzyme just calculates the component wise derivatives an stops.

Now, the user *can* extract out the covariant derivative given this Cartesian derivative `dA`, and the primal `A`, but doing so can be difficult, and many users are going to be quite surprised to find that Enzyme is giving them a seemingly "wrong" derivative, and are unlikely to even know how to transform their component wise derivatives into covariant derivatives.

## What are Zygote / ChainRules doing here?

Basically, ChainRules based AD systems are going to default to these Cartesian 'structural' derivatives, unless they have rules put in place to `ProjectTo` the shadow values back to the correct manifold, and they're doing this at the level of rules, so the projections are happening throughout the whole AD process. Here's some resources:

* https://github.com/mcabbott/OddArrays.jl This does a great job of showing various examples that are even more tricky than `Symmetric` and shows how one can try to systematically deal with them
* https://github.com/JuliaDiff/ChainRulesCore.jl/blob/wct/writing-generic-rrules/notes.md
* https://juliadiff.org/ChainRulesCore.jl/stable/design/many_tangents.html
* https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/projectto.html

The problem is that this approach
1) It's pretty hard. It can be computationally expensive, and it requires rules authors never knowing what type of shadow values they're dealing with, and they can end up returning all sorts of crazy stuff to downstream rules. I must say that despite the many difficulties with this approach, I kinda like it, but I'll let people who have actually had to suffer through writing rules and such for it talk about their experiences.
2) You never know if you'll get a Cartesian or a Covariant derivative. If someone implements a type, and you differentiate it without any `ProjectTo` rules being written, you'll end up getting a Cartesian derivative by default. Now, I should mention that this thing is surprisingly robust. i.e. it's pretty smart for general abstract arrays, i.e. writing `Symmetric` myself works correctly without any custom rules or `ProjectTo` methods:
```julia
julia> struct MySymmetric{T} <: AbstractMatrix{T}
data::Matrix{T}
end;

julia> Base.getindex(m::MySymmetric, i::Int, j::Int) = i <= j ? m.data[i, j] : m.data[j, i];

julia> Base.size(m::MySymmetric) = size(m.data);

julia> Zygote.gradient(sum, MySymmetric(rand(3,3))) |> only
3×3 Fill{Float64}, with entries equal to 1.0

julia> Zygote.gradient(MySymmetric(rand(3,3))) do A
s = zero(eltype(A))
for i ∈ eachindex(A)
s += A[i]
end
s
end |> only
3×3 Matrix{Float64}:
1.0 1.0 1.0
1.0 1.0 1.0
1.0 1.0 1.0
```
so this concern is more of a potential concern for non-AbstractArray types if I understand correctly (please chine in if you disagree with this assessment, I'm not saying this is 100% reliable, but I think it is already pretty reliable, and can be made even more reliable).

## What should Enzyme do?

Well, that's certainly up for discussion, but I think maybe the right thing to do would be to develop a set of functions that given a primal and a Cartesian derivative, can calculate a covariant derivative as a post processing pass.

E.g. we could have `autodiff_cartesian` which does something similar to what `autodiff` does currently, and then make it so that autodiff does something like
```julia
autodiff(Reverse, f, Active(x), Duplicated(y, dy_cartesian))
```
turning into something like
```julia
tape, df_cartesian, dx_cartesian = autodiff_cartesian(ReverseWithTangentTape, f, Active(x), Duplicated(y, dy_cartesian))
df = to_tangent_space(tape, f, df_cartesian)
dx = to_tangent_space(tape, x, dx_cartesian)
dy = to_tangent_space(tape, y, dy_cartesian)
return (df, dx,dy)
```
I *think* that in general in order to do this conversion to the tangent space, we'd need a record of which functions were hit (the `tape` and then a series of transformations would be applied to `dy_cartesian` based potentially on that tape and the primal `y` (see https://github.com/mcabbott/OddArrays.jl for examples of cases that need `y` and `dy_tangent`).

I think it'd be good to support a purely `cartesian` mode, and then have a `to_tangent_space` be pretty fussy about rejecting cases it doesn't understand with hard errors, but useful error hints. E.g. something like
![image](https://github.com/EnzymeAD/Enzyme.jl/assets/29157027/5b0e01fc-568b-4843-8908-4b1017c37121)

The design of the `tape` would be pretty hard to do right though.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.