# Non-convex loss functions with Flux.jl

**URL:** <https://discourse.julialang.org/t/non-convex-loss-functions-with-flux-jl/88273>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [October 5, 2022, 6:38am UTC](https://discourse.julialang.org/t/non-convex-loss-functions-with-flux-jl/88273 "2022-10-05T06:38:55Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![MattSainsbury-Dale](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mattsainsbury-dale/32/221714_2.png) [@MattSainsbury-Dale](https://discourse.julialang.org/u/MattSainsbury-Dale)\
**Post date:** [October 5, 2022, 6:38am UTC](https://discourse.julialang.org/t/non-convex-loss-functions-with-flux-jl/88273/1 "2022-10-05T06:38:55Z")

</div>

I am interested in training a `Flux` neural network under a non-convex loss function. Specifically, the loss function L\_P(x, y) = (|x - y|)^P for some 0 \< P \< 1, which is implemented in `LossFunctions.jl` as [`LPDistLoss`](https://juliaml.github.io/LossFunctions.jl/stable/losses/distance/#LPDistLoss-1).

Is it possible to use this non-convex loss function in `Flux`? I have attempted to implement L\_P(\cdot, \cdot) by adapting the source code for `Flux.mae` as follows:

```julia
# ofeltype() and _check_sizes() are internal functions from Flux
function LP(ŷ, y; agg = mean, P = ofeltype(ŷ, 0.9))
  _check_sizes(ŷ, y)
  agg(abs.(ŷ .- y).^P)
end

```

However, even with P = 0.9, which is only slightly different to the well-behaved absolute-error loss, some parameters of the neural network become `NaN` during training.

Is this a problem inherent to this loss function, or is there something wrong with my implementation? Thanks in advance for any comments or suggestions.

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [October 6, 2022, 12:15am UTC](https://discourse.julialang.org/t/non-convex-loss-functions-with-flux-jl/88273/2 "2022-10-06T00:15:35Z")

</div>

That definition looks fine to me. Is there an existing working example of training a model with this loss that you can compare against?

---

<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:** [October 6, 2022, 12:59am UTC](https://discourse.julialang.org/t/non-convex-loss-functions-with-flux-jl/88273/3 "2022-10-06T00:59:52Z")

</div>

> [@MattSainsbury-Dale](#):
>
> However, even with P = 0.9 P=0.9P = 0.9 , which is only slightly different to the well-behaved absolute-error loss, some parameters of the neural network become `NaN`

Note that the gradient is quite different near to zero, `x^-0.1` diverges. That could be what goes wrong?

 ![two_powers](https://global.discourse-cdn.com/julialang/original/3X/1/9/195ffac1a639ee8f75d5d513bc91dc49dd83adc3.png)

You could try something like `(abs.(ŷ .- y) .+ 1f-6).^P` to move slightly away from the point. You could also try `P = 1.1` to check

---

<div class="post-metadata">

**Author:** ![MattSainsbury-Dale](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mattsainsbury-dale/32/221714_2.png) [@MattSainsbury-Dale](https://discourse.julialang.org/u/MattSainsbury-Dale)\
**Post date:** [October 6, 2022, 1:12pm UTC](https://discourse.julialang.org/t/non-convex-loss-functions-with-flux-jl/88273/4 "2022-10-06T13:12:21Z")

</div>

Many thanks for your replies.

> @ToucheSir Is there an existing working example of training a model with this loss that you can compare against?

I am not aware of any such examples. I am interested in this loss function because it tends towards the 0-1 loss as P \to 0. There are many well-behaved surrogates for the 0-1 loss that work in a classification setting (i.e., with discrete outputs), but I haven’t found any examples of people doing this for a continuous output, like in my application. (Do you know of any?)

> @mcabbott Note that the gradient is quite different near to zero, `x^-0.1` diverges. That could be what goes wrong?

Thanks for your insights, I think you’ve hit the nail on the head. I tried your suggestion of adding a small positive quantity to |x - y|, and this fixed the `NA` problem. Thanks!

You may be interested that adding a small positive quantity to |x - y| actually has a large effect on the loss function as P \to 0. In particular, it increases its minimum value (it should be zero):

 ![loss](https://global.discourse-cdn.com/julialang/original/3X/0/7/07e4d2db2e09c8b314222d9cbdfd22b5fe251524.png)

---

<div class="post-metadata">

**Author:** ![DanielV](https://avatars.discourse-cdn.com/v4/letter/d/a9a28c/32.png) [@DanielV](https://discourse.julialang.org/u/DanielV)\
**Post date:** [October 18, 2022, 3:19am UTC](https://discourse.julialang.org/t/non-convex-loss-functions-with-flux-jl/88273/5 "2022-10-18T03:19:49Z")

</div>

Note that algorithmically speaking, the minimum value of the loss has minor effect, on both the solution found and the behavior of algorithms.

Might want to focus attention on the locally uniform upper bound on the gradient (which affects both the useful step sizes and the effect of a single loss on a local minimum, under assumption of local strong convexity). So consider using something like SCAD [1] (presented as penalizer, but something like it can be used as a loss), with its L1-like behavior near 0. You can always adjust the constants to iget closer to 0-1, while retaining uniform bounded gradient norm.

[1] [https://fan.princeton.edu/papers/01/penlike.pdf](https://fan.princeton.edu/papers/01/penlike.pdf)