Flux.jl: back! calculates gradient incorrectly for two layer perceptron with relu

[philip-bl]

Hello. I've only started trying to use Julia and FluxML for deep learning purposes recently, and most of the basic stuff doesn't work :( Here I demonstrate two identical programs, one is using Pytorch, the other one is using FluxML. They produce identical results until back propagation is performed. The gradients calculated by back propagation are not identical, and I believe julia's gradient are incorrect.

https://gist.github.com/philip-bl/588ffe4c6b1e740a087b0af4ad03e561

In case it matters, here's my versioninfo() output:

Julia Version 1.1.0
Commit 80516ca202 (2019-01-21 21:24 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Core(TM) i7-4800MQ CPU @ 2.70GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.1 (ORCJIT, haswell)

and here is output of ] status

    Status `~/.julia/environments/v1.1/Project.toml`
  [fbb218c0] BSON v0.2.3
  [5ae59095] Colors v0.9.5
  [587475ba] Flux v0.8.2
  [7073ff75] IJulia v1.18.1
  [6218d12a] ImageMagick v0.7.1
  [916415d5] Images v0.17.3
  [9920b226] MLDataPattern v0.5.0
  [eb30cadb] MLDatasets v0.3.0
  [91a5bcdd] Plots v0.24.0
  [5e47fb64] TestImages v0.4.1
  [c2297ded] ZMQ v1.0.0

[philip-bl]

Hmm, if I run the same Julia code on http://juliabox.com/ with versioninfo()

Julia Version 1.0.3
Commit 099e826241 (2018-12-18 01:34 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Xeon(R) CPU E5-2673 v4 @ 2.30GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.0 (ORCJIT, broadwell)
Environment:
  JULIABOX = true
  JULIA_PKG_SERVER = https://pkg.juliacomputing.com
  JULIA = /opt/julia-0.6/bin/julia
  JULIA_KERNELS = ['julia-0.6', 'julia-1.0', 'julia-1.0k']
  JULIA_PKG_TOKEN_PATH = /mnt/juliabox/.julia/token.toml
  JULIABOX_ROLE = 
  JULIA_NUM_THREADS = 4

then I get NaNs in gradients. I don't know why it is so.

[ornithos]

Nice find. Looks like you're running into numerical instability issues from a suboptimal implementation of the logsoftmax derivative.

(FWIW this problem is unlikely to be encountered too regularly with usual initialisations / step-size / regularisation of NNs: the final layer inputs to the softmax in the variable buzz are large:

Tracked 3×5 Array{Float64,2}:
 259.0  136.0   357.0  178.0   455.0
 610.0  307.0   844.0  405.0  1078.0
 961.0  478.0  1331.0  632.0  1701.0

and causing substantial saturation of the softmax function. If your comment about "most of the basic stuff" not working refers to other issues too, please give more info.)

Nevertheless, it indicates that the tests of logsoftmax are too weak. For example:

using Flux
using Flux.Tracker: ngradient
f(u) = sum(x->x^2, logsoftmax(u))
testinput = param(randn(4,2))

Flux∇ = Flux.gradient(f, testinput)
Num∇ = ngradient(f, testinput.data)
all(isapprox.(Tracker.data.(Flux∇), Num∇, rtol = 1e-5, atol = 1e-5))
julia> true

testinput *= 100
Flux∇ = Flux.gradient(f, testinput)
Num∇ = ngradient(f, testinput.data)
all(isapprox.(Tracker.data.(Flux∇), Num∇, rtol = 1e-5, atol = 1e-5))
julia> false

and hence with different tests, this would have been picked up. The current derivative definition is in NNlib here. I'll try to get a PR in this evening if I have time.

[MikeInnes]

Should be closed by https://github.com/FluxML/NNlib.jl/pull/126