# Loss functions that involve gradients

**URL:** <https://discourse.julialang.org/t/loss-functions-that-involve-gradients/55293>\
**Category:** Machine Learning\
**Created:** [February 15, 2021, 2:45am UTC](https://discourse.julialang.org/t/loss-functions-that-involve-gradients/55293 "2021-02-15T02:45:23Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![balaji1975](https://avatars.discourse-cdn.com/v4/letter/b/8edcca/32.png) [@balaji1975](https://discourse.julialang.org/u/balaji1975)\
**Post date:** [February 15, 2021, 2:45am UTC](https://discourse.julialang.org/t/loss-functions-that-involve-gradients/55293/1 "2021-02-15T02:45:23Z")

</div>

This is similar to [this question](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569).

I am training a neural net NN(), where the loss function involves gradient of NN(.) wrt to the input. For example:

```julia
m = Flux.Chain(Dense(5,5,relu), Dense(5,5,relu), Dense(5,1))
g(z) = only(m(z))

function loss(x,y)
    w = Flux.gradient(g, x)[1]
    yhat = dot(w, x) / g(x)
    return (yhat - y)^2
end

```

This does not work within the Flux framework. Currently, I train the network 1) “manually” using ForwardDiff.gradient where required and 2) where the network is small, using Optim.

Is there a way we can handle this within Flux? Many thanks!

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [February 15, 2021, 5:40am UTC](https://discourse.julialang.org/t/loss-functions-that-involve-gradients/55293/2 "2021-02-15T05:40:36Z")

</div>

I use such loss functions, but they are finicky. You need to be certain that gradients are all differentiable, which means that you have to in practice re-write some gradients, as they are not second-order gradients friendly. But in theory, zygote can do 2nd order gradients. Be aware that compilation time might be long.

---

<div class="post-metadata">

**Author:** ![balaji1975](https://avatars.discourse-cdn.com/v4/letter/b/8edcca/32.png) [@balaji1975](https://discourse.julialang.org/u/balaji1975)\
**Post date:** [February 16, 2021, 2:29am UTC](https://discourse.julialang.org/t/loss-functions-that-involve-gradients/55293/3 "2021-02-16T02:29:49Z")

</div>

Thanks Tomas.

---

<div class="post-metadata">

**Author:** ![lazarusA](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lazarusa/32/6571_2.png) [@lazarusA](https://discourse.julialang.org/u/lazarusA)\
**Post date:** [November 17, 2022, 9:02pm UTC](https://discourse.julialang.org/t/loss-functions-that-involve-gradients/55293/4 "2022-11-17T21:02:15Z")

</div>

@balaji1975 how do you do the training with ForwadDiff.gradient?

what code follows after your loss? Specifically how do you do the call for gradients and the update later of parameters?

---

<div class="post-metadata">

**Author:** ![balaji1975](https://avatars.discourse-cdn.com/v4/letter/b/8edcca/32.png) [@balaji1975](https://discourse.julialang.org/u/balaji1975)\
**Post date:** [November 22, 2022, 2:08am UTC](https://discourse.julialang.org/t/loss-functions-that-involve-gradients/55293/5 "2022-11-22T02:08:03Z")

</div>

@lazarusA

It has been a long time, memory is sketchy. I ended up using just Flux with a custom loss function.

Essentially, i was trying to find a neural network N that minimized the error:  
 ![Screenshot 2022-11-22 at 10.05.33 AM](https://global.discourse-cdn.com/julialang/original/3X/7/f/7f086afdda08c8602979c6b01b2f364d0bc55148.png)
