# Gradient of a gradient of a FastChain

**URL:** <https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876>\
**Category:** General Usage\
**Tags:** question, zygote, ad, neural-network\
**Created:** [December 31, 2021, 7:56pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876 "2021-12-31T19:56:50Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![mahdiar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mahdiar/32/32352_2.png) [@mahdiar](https://discourse.julialang.org/u/mahdiar)\
**Post date:** [December 31, 2021, 7:56pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/1 "2021-12-31T19:56:50Z")

</div>

What is the best way to differentiate an artificial neural network built by FastChain?

The objective is to make a gradient of a neural network output with respect to its output. The result should be differentiable.

Here is a dummy example.

```julia
using DiffEqFlux, Zygote
nn = FastChain(FastDense(1,32,tanh), FastDense(32,32,tanh), FastDense(32,1))
θ = initial_params(nn)
Zygote.gradient(nn[1],0.1,θ)

```

It results in the following error

```julia
MethodError: no method matching getindex(::FastChain{Tuple{FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#92"{Vector{Float32}}}, FastDense{typeof(tanh), DiffEqFlux.var"#initial_params#92"{Vector{Float32}}}, FastDense{typeof(identity), DiffEqFlux.var"#initial_params#92"{Vector{Float32}}}}}, ::Int64)

Stacktrace:
 [1] top-level scope
   @ In[3]:5
 [2] eval
   @ ./boot.jl:373 [inlined]
 [3] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)
   @ Base ./loading.jl:1196

```

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [January 1, 2022, 12:34am UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/2 "2022-01-01T00:34:28Z")

</div>

Gradient of what? You’re not taking the gradient of a function.

```julia
Zygote.gradient(x -> nn(x,θ)[1],[0.1])

```

---

<div class="post-metadata">

**Author:** ![mahdiar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mahdiar/32/32352_2.png) [@mahdiar](https://discourse.julialang.org/u/mahdiar)\
**Post date:** [January 1, 2022, 4:29pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/3 "2022-01-01T16:29:18Z")

</div>

Thanks for providing the correct syntax to use. There still is a problem with `DiffEqFlux.(sciml_train)`, it might be related to differentiability or again a mistake in syntax.

```julia
using DiffEqFlux, Zygote, LinearAlgebra
nn = FastChain(FastDense(1,32,tanh), FastDense(32,32,tanh), FastDense(32,1))
θ = initial_params(nn)
cost(nn, θ) = reduce(vcat,first.(Zygote.gradient(x -> nn(x,θ)[1],[0.1])))
loss(θ) = cost(nn,θ)
DiffEqFlux.sciml_train(loss, ADAM, θ, maxiters=2)

```

Any idea how to fix this?

Btw, the input of the `loss()` function is the optimization parameter vector \theta. Is there a way to use `cost()` function directly in `DiffEqFlux.(sciml_train)`?

---

<div class="post-metadata">

**Author:** ![Elrod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/elrod/32/22461_2.png) [@Elrod](https://discourse.julialang.org/u/Elrod)\
**Post date:** [January 1, 2022, 4:33pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/4 "2022-01-01T16:33:27Z")

</div>

The loss should be a scalar.

---

<div class="post-metadata">

**Author:** ![mahdiar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mahdiar/32/32352_2.png) [@mahdiar](https://discourse.julialang.org/u/mahdiar)\
**Post date:** [January 1, 2022, 8:08pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/5 "2022-01-01T20:08:20Z")

</div>

Isn’t it a scalar in the previous example?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [January 1, 2022, 9:00pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/6 "2022-01-01T21:00:14Z")

</div>

Are you trying to take a gradient of a gradient of a FastChain?

```julia
using DiffEqFlux, Zygote, LinearAlgebra
nn = FastChain(FastDense(1,32,tanh), FastDense(32,32,tanh), FastDense(32,1))
θ = initial_params(nn)
cost(nn, θ) = Zygote.gradient(x -> nn(x,θ)[1],[0.1])[1]
loss(θ) = cost(nn,θ)
DiffEqFlux.sciml_train(loss, θ, ADAM(0.1), maxiters=2)

```

That’s very different from a gradient of a FastChain.

---

<div class="post-metadata">

**Author:** ![mahdiar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mahdiar/32/32352_2.png) [@mahdiar](https://discourse.julialang.org/u/mahdiar)\
**Post date:** [January 1, 2022, 9:10pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/7 "2022-01-01T21:10:23Z")

</div>

To be precise, I want to use a gradient of a FastChain in a loss function. Many of the times AD was failing at the `DiffEqFlux.sciml_train()`, which basically means that the gradient of a gradient is failing somewhere. Any idea how to fix the simple example?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [January 1, 2022, 9:23pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/8 "2022-01-01T21:23:47Z")

</div>

Okay yes, this is a gradient of a gradient example, not a gradient example. The gradient of a gradient of a FastChain won’t work because the FastChain adjoint uses mutation. It could be specialized to handle this case, but because this is almost never an efficient way to calculate the second derivative (forward-over-adjoint is just better in almost all respects) I’m not sure it’s a high priority to specialize this.

(And BTW, taking a gradient of gradient is a TensorFlow misnomer. It’s actually Jacobian of a gradient unless it’s a scalar function and thus the second derivative. Otherwise the sizes don’t align. TensorFlow silently makes gradient = sum of Jacobian as I describe here [Gradient of Gradient in Zygote - #3 by ChrisRackauckas](https://discourse.julialang.org/t/gradient-of-gradient-in-zygote/52685/3) You should really double check whether that summation is the interpretation you wanted)

---

<div class="post-metadata">

**Author:** ![Elrod](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/elrod/32/22461_2.png) [@Elrod](https://discourse.julialang.org/u/Elrod)\
**Post date:** [January 2, 2022, 4:32am UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/9 "2022-01-02T04:32:27Z")

</div>

`reduce(vcat, ...)`

Won’t return a scalar in general.  
But requiring a gradient of a gradient with the fast chains is a secondary issue, as Chris R pointed out.

---

<div class="post-metadata">

**Author:** ![cortner](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cortner/32/204_2.png) [@cortner](https://discourse.julialang.org/u/cortner)\
**Post date:** [January 2, 2022, 5:53am UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/10 "2022-01-02T05:53:31Z")

</div>

> [@ChrisRackauckas](#):
>
> It could be specialized to handle this case, but because this is almost never an efficient way to calculate the second derivative (forward-over-adjoint is just better in almost all respects) I’m not sure it’s a high priority to specialize this

Interesting. Can you point to a source where I can understand this? I’ve written rrules for rrules with great success (I thought …) with performance reasonably close to the first `rrule` which indicated to me that any other approach couldn’t really be more performant… But maybe I have a special case, or maybe I misunderstood something?

EDIT: never mind, I see your link above has some discussion … I will peruse it.

---

<div class="post-metadata">

**Author:** ![mahdiar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mahdiar/32/32352_2.png) [@mahdiar](https://discourse.julialang.org/u/mahdiar)\
**Post date:** [January 2, 2022, 4:13pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/11 "2022-01-02T16:13:29Z")

</div>

> [@ChrisRackauckas](#):
>
> Okay yes, this is a gradient of a gradient example, not a gradient example. The gradient of a gradient of a FastChain won’t work because the FastChain adjoint uses mutation. It could be specialized to handle this case, but because this is almost never an efficient way to calculate the second derivative (forward-over-adjoint is just better in almost all respects) I’m not sure it’s a high priority to specialize this.
> 
> (And BTW, taking a gradient of gradient is a TensorFlow misnomer. It’s actually Jacobian of a gradient unless it’s a scalar function and thus the second derivative. Otherwise the sizes don’t align. TensorFlow silently makes gradient = sum of Jacobian as I describe here [Gradient of Gradient in Zygote - #3 by ChrisRackauckas](https://discourse.julialang.org/t/gradient-of-gradient-in-zygote/52685/3) You should really double check whether that summation is the interpretation you wanted)

Is [this topic](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569) relevant? The model in their example is based on `Chain` instead of `FastChaint`.

Thanks for the additional information on TensorFlow. Maybe this is one of the reasons why people in PINNs use multiple neural networks for estimating each variable? Sum of Jacobian is clearly not the best approach.

---

<div class="post-metadata">

**Author:** ![mahdiar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mahdiar/32/32352_2.png) [@mahdiar](https://discourse.julialang.org/u/mahdiar)\
**Post date:** [January 2, 2022, 4:21pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/12 "2022-01-02T16:21:24Z")

</div>

> [@Elrod](#):
>
> `reduce(vcat, ...)`
> 
> Won’t return a scalar in general.  
> But requiring a gradient of a gradient with the fast chains is a secondary issue, as Chris R pointed out.

Understood. But, in this example, there is only one element in `first.(Zygote.gradient(x -> nn(x,θ)[1],[0.1]))`, so the output of `reduce(vcat, ...)` is scalar. Is this still a problem here? Removing `reduce(vcat, ...)` changes the error.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [January 2, 2022, 4:51pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/13 "2022-01-02T16:51:37Z")

</div>

> [@mahdiar](#):
>
> Is [this topic](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569) relevant? The model in their example is based on `Chain` instead of `FastChaint` .

It’s the same issue as in both cases what’s happening is that the adjoint definition itself is not Zygote differentiable, which is where the double differentiation issue comes into play. But then the solution is the same in both cases that you probably want to use Forward-over-Adjoint anyways (replying to @cortner, IIRC there’s a lengthy discussion in Griewank’s AD book about how double reverse mode is almost never optimal), and that thread shows how to do FoA second derivatives via mixing ForwardDiff.jl into Zygote.jl. So my suggestion is the same. In this case, it’s actually not too hard to fix the double adjoint, but you still wouldn’t want to use it if I fixed that so… ehh… maybe later.

> [@mahdiar](#):
>
> Thanks for the additional information on TensorFlow. Maybe this is one of the reasons why people in PINNs use multiple neural networks for estimating each variable? Sum of Jacobian is clearly not the best approach.

Yes, we discuss this a bit in the [weird paper-y write-up thingy we wrote on NeuralPDE.jl](https://arxiv.org/abs/2107.09443) (not quite a paper, not quite a review, but full of relevant information). Splitting into separate neural networks decreases the asymptotic cost of differentiation for systems of PDEs. You could still use a few tricks to make the neural networks share weights/layers in this form BTW (I should write an example for how to do that), but putting the networks together both has a higher cost and a tendency to compute things that aren’t needed (for example, if you want AD to compute the 2nd derivative it would compute the second derivative w.r.t. every output, even if only one dependent variable was undergoing diffusion).

So yes, there’s a lot going on in AD space to improve PINNs. Mixing the networks it not asymptotically good. Reverse over reverse is not good, so you want to mix forward with reverse. Standard AD mixing without extra tricks is not good to higher order (see the stuff on [Taylor-mode AD](https://openreview.net/forum?id=SkxEF3FNPH): you need to use things like that in order to make 3rd derivatives scale much better, and even the heat equation has 3rd derivatives when you consider the one derivative for the loss function, this is why [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl) exists). The Julia Lab will be spending a good part of next year to demonstrate how handling all of these more effectively together improves PINN training performance, but it’s not all ready right now.

---

<div class="post-metadata">

**Author:** ![mahdiar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mahdiar/32/32352_2.png) [@mahdiar](https://discourse.julialang.org/u/mahdiar)\
**Post date:** [January 2, 2022, 6:07pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/14 "2022-01-02T18:07:54Z")

</div>

> [@ChrisRackauckas](#):
>
> It’s the same issue as in both cases what’s happening is that the adjoint definition itself is not Zygote differentiable, which is where the double differentiation issue comes into play. But then the solution is the same in both cases that you probably want to use Forward-over-Adjoint anyways (replying to @cortner, IIRC there’s a lengthy discussion in Griewank’s AD book about how double reverse mode is almost never optimal), and that thread shows how to do FoA second derivatives via mixing ForwardDiff.jl into Zygote.jl. So my suggestion is the same. In this case, it’s actually not too hard to fix the double adjoint, but you still wouldn’t want to use it if I fixed that so… ehh… maybe later.

How is it possible to use FoA with `sciml_train()`? Based on the simple example provided earlier, the cost function is using `Zygote.gradient()`. Is there an option in `sciml_train()` to use ForwardDiff in the optimization process?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [January 3, 2022, 2:58am UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/15 "2022-01-03T02:58:02Z")

</div>

It’s the same thing. Just use ForwardDiff.jl in the cost function.

---

<div class="post-metadata">

**Author:** ![mahdiar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mahdiar/32/32352_2.png) [@mahdiar](https://discourse.julialang.org/u/mahdiar)\
**Post date:** [January 3, 2022, 4:54pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/16 "2022-01-03T16:54:41Z")

</div>

Do you mean something like this:

```julia
using DiffEqFlux, ForwardDiff, LinearAlgebra
nn = FastChain(FastDense(1,32,tanh), FastDense(32,32,tanh), FastDense(32,1))
θ = initial_params(nn)
cost(θ_) = ForwardDiff.gradient(x -> nn(x,θ_)[1],[0.1])[1]
DiffEqFlux.sciml_train(cost, θ, ADAM(0.1), maxiters=2)

```

It results in the following error.

```julia
MethodError: Cannot `convert` an object of type Nothing to an object of type Float32

```

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [January 5, 2022, 2:52am UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/17 "2022-01-05T02:52:04Z")

</div>

```julia
using DiffEqFlux, ForwardDiff, LinearAlgebra
nn = FastChain(FastDense(1,32,tanh), FastDense(32,32,tanh), FastDense(32,1))
θ = initial_params(nn)
function cost(θ_)
  f = x -> nn(x,θ_)[1]
  x = [0f-1]
  sum(ForwardDiff.gradient(f, x))
end
Zygote.gradient(cost,θ) # nothing

```

That shouldn’t be a `nothing`. It looks like it’s an issue with Zygote’s overload of ForwardDiff.gradient. MWE:

```julia
using Zygote
θ = rand(2,2)
function cost(θ_)
  f = x -> sum(θ_*x)
  x = [1.0,2.0]
  sum(ForwardDiff.gradient(f, x))
end
Zygote.gradient(cost,θ) # nothing

```

@mcabbott it looks like this would’ve been covered by the same thing as [https://github.com/FluxML/Zygote.jl/issues/953#issuecomment-841882071](https://github.com/FluxML/Zygote.jl/issues/953#issuecomment-841882071) . Maybe you have an idea what’s going on here?

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [January 5, 2022, 3:16am UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/18 "2022-01-05T03:16:34Z")

</div>

Yes, `ForwardDiff.gradient(f, x)` does not keep a gradient with respect to `f`.

Maybe that could be done better, there was a thread of attempts somewhere.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [January 5, 2022, 10:57am UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/19 "2022-01-05T10:57:38Z")

</div>

Well interesting. @simeonschaub is this a good Diffractor case or still too early?

---

<div class="post-metadata">

**Author:** ![mahdiar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mahdiar/32/32352_2.png) [@mahdiar](https://discourse.julialang.org/u/mahdiar)\
**Post date:** [February 8, 2022, 3:58pm UTC](https://discourse.julialang.org/t/gradient-of-a-gradient-of-a-fastchain/73876/20 "2022-02-08T15:58:28Z")

</div>

Is there any update for this thread? Do we have a friendly way to use the derivative of a neural network in the cost function?