# Nested AD with Lux etc

**URL:** <https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573>\
**Category:** Machine Learning\
**Tags:** ad\
**Created:** [April 25, 2024, 9:39pm UTC](https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573 "2024-04-25T21:39:42Z")\
**Posts on this page:** 7\
**Page:** 2

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [April 28, 2024, 3:53pm UTC](https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573/22 "2024-04-28T15:53:52Z")

</div>

Let me clarify. Say we are defining the following loss function

```julia
function loss_function(m, x, ps, st)
    f = x -> m(x, ps, st)
    return sum(abs2, only(Zygote.gradient(f, x)))
end

```

If we do a `Zygote.gradient(loss_function, ....)` it will be very slow in the best case and error in the vast majority of cases, so we want to use JVP over Zygote.gradient to compute the gradients wrt x and ps.

```julia
function CRC.rrule(cfg::..., ::typeof(Zygote.gradient), f::F, x) where {F}
     # Here `f` is potentially a closure
     # Q1: How do we extract the values like `m`, `ps` and `st` in `f`?
     # Let's proceed assuming we did it
     y = Zygote.gradient(f, x)
     function gradient_pullback(Delta)
           # We use our jvp trick to compute dm, dps, dst, dx
           return NoTangent(), NoTangent(), Tangent for `f`, dx
           # Q2: What is the appropriate tangent for `f`
     end
     return y, gradient_pullback
end

```

Since I did not know how to solve this what Lux does and I am proposing for DI to do is., we have a `ParamsStruct`

```julia
function Zygote.gradient(f::ParamsStruct, x)
     return __my_custom_gradient_func(f.func, x, f.ps) # In the forward pass we just call Zygote.gradient
end

# Defining rrule is now simple as
function CRC.rrule(cfg::..., ::typeof(__my_custom_gradient_func), f::F, x, ps) where {F}
     # Use the JVP trick here
end
```

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [April 28, 2024, 10:09pm UTC](https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573/23 "2024-04-28T22:09:59Z")

</div>

> [@avikpal](#):
>
> Since I did not know how to solve this what Lux does and I am proposing for DI to do is., we have a `ParamsStruct`

It seems like a very specific hack, and I wouldn’t want DI to become a collection of such hacks. Can we maybe take a step back and figure out the general principle behind it? Switching the mode for one of two input arguments?

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [April 29, 2024, 12:21am UTC](https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573/24 "2024-04-29T00:21:49Z")

</div>

It is not really a hack nor is it switching the backend entirely. The claim is doing reverse over reverse comes with limited benefits and none of the current AD tools in julia can really handle this (except maybe enzyme). However, Dual Numbers pretty much works with with any of the AD tools that we have.

Let’s say you have an `outer_func` which takes `N` inputs. Then an inner function `inner_func` closes over `M < N` inputs and computes the gradient for `M - N` of the inputs inside `outer_func` and performs some arbitrary operation returning a scalar.

If I call `Zygote.gradient(outer_func, args...)`, I expect gradients wrt all the `args...` (unlike what Zygote does to the forwarddiff operations which drops gradients (with a warning)). If you try reverse mode over the inner gradient function, you are pretty much screwed.

So what we should be doing instead is unwrap the closure and construct Dual Numbers using the input to the pullback operator and compute the inner gradient for **all** the N inputs (note we only changed the outer gradient call with a JVP, see [Nested AD with Lux etc - #12 by stevengj](https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573/12) for why this is true), extract the partials and return the gradients via an rrule. Now we could have easily done it without unwrapping the closure, but it seems messy to return a tangent type corresponding to the closure.

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [April 29, 2024, 12:36am UTC](https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573/25 "2024-04-29T00:36:05Z")

</div>

> [@oxinabox](#):
>
> Diffractor does it under the hood for one of the tests.

Do you have a link to where it does it?

I see what you mean by the NamedTuple part:

```julia
julia> f = let dd=dd
           x -> dd .+ x
       end
#3 (generic function with 1 method)

julia> fieldnames(typeof(f))
(:dd,)

```

---

<div class="post-metadata">

**Author:** ![oxinabox](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oxinabox/32/206603_2.png) [@oxinabox](https://discourse.julialang.org/u/oxinabox)\
**Post date:** [April 29, 2024, 3:23am UTC](https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573/26 "2024-04-29T03:23:57Z")

</div>

> [@avikpal](#):
>
> Do you have a link to where it does it?

Here is [the test](https://github.com/JuliaDiff/Diffractor.jl/blob/d0b3e3e0ea59805d6aadbcf867d370180753a4c1/test/forward_mutation.jl#L69-L84)

It just uses the [normal path for constructing objects.](https://github.com/JuliaDiff/Diffractor.jl/blob/d0b3e3e0ea59805d6aadbcf867d370180753a4c1/src/stage1/recurse_fwd.jl#L14-L27)  
and then the [normal path for differentiating calls.](https://github.com/JuliaDiff/Diffractor.jl/blob/d0b3e3e0ea59805d6aadbcf867d370180753a4c1/src/stage1/forward.jl#L202-L209)

> [@avikpal](#):
>
> ```julia
> function CRC.rrule(cfg::..., ::typeof(Zygote.gradient), f::F, x) where {F}
> 
> ```

> ```
> # Q1: How do we extract the values like `m`, `ps` and `st` in `f`?
> 
> ```

`f.m`, `f.ps`, `f.st`, etc.  
This is, as I think you have worked out, what i mean by closures being callable NamedTuples.

> `# We use our jvp trick to compute dm, dps, dst, dx` …  
> `# Q2: What is the appropriate tangent for `f`

It is `Tangent{F}(m=dm, ps=dps, st=dst)`

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [April 30, 2024, 10:12pm UTC](https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573/27 "2024-04-30T22:12:34Z")

</div>

> [@avikpal](#):
>
> The main issue however are the parameter gradients which commonly show up in most optimization setups

For future reference, here is a manual implementation of forward-over-reverse calculation with parameter gradients (mixed second derivatives). In particular, if you have a scalar-valued h(x,p) = g(\nabla\_x f) for some scalar-valued f(x,p), then one can similarly derive:

\left. \nabla\_p h \right|\_{x,p} = \left. \frac{\partial}{\partial\alpha} \left. \nabla\_p f \right|\_{x + \alpha \left. \nabla g \right|\_{z},p} \right|\_{\alpha = 0} \, ,

where z = \left. \nabla\_x f \right|\_{x,p}.

Here is an example calculation via ForwardDiff over Zygote, along with a finite-difference check:

```julia
julia> using ForwardDiff, Zygote, LinearAlgebra

julia> f(x,p) = sum(p)^2/norm(x) + p[1]*x[2]; # example function

julia> g(∇ₓf) = sum(∇ₓf)^3; # example ℝⁿ → ℝ function

julia> h(x,p) = g(Zygote.gradient(x -> f(x,p), x)[1]); # evaluate h by reverse mode

julia> function ∇ₚh(x,p)
           ∇ₚf(y,q) = Zygote.gradient(u -> f(y,u), q)[1]
           ∇g = Zygote.gradient(g, Zygote.gradient(x -> f(x,p), x)[1])[1]
           return ForwardDiff.derivative(α -> ∇ₚf(x + α*∇g, p), 0)
       end;

julia> x = randn(5); p = randn(4); δp = randn(4) * 1e-8;

julia> h(x,p)
-6.538498714556666e-5

julia> ∇ₚh(x,p)
4-element Vector{Float64}:
  0.0025659422776640596
 -0.0023030596877005173
 -0.0023030596877005173
 -0.0023030596877005173

julia> h(x,p+δp) - h(x,p) # finite-difference check
-5.696379205464251e-11

julia> ∇ₚh(x,p)'δp # exact directional derivative
-5.6963775418992394e-11

```

PS. This also seems like a good example of the clarity benefits of Unicode variable names for mapping math to code.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [May 1, 2024, 6:05am UTC](https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573/28 "2024-05-01T06:05:44Z")

</div>

I still don’t have a very thorough understanding of what’s needed here, but if @avikpal wants to submit a PR to DifferentiationInterface.jl doing that, I’ll happily review and discuss more

[Previous page](https://discourse.julialang.org/t/nested-ad-with-lux-etc/113573.md?page=1)