Merge pull request #126 from ornithos/logsoftmaxgrad

Improve numerical stability of logsoftmax gradient

Author: MikeInnes

src/softmax.jl

@@ -81,5 +81,5 @@ function logsoftmax!(out::AbstractVecOrMat, xs::AbstractVecOrMat)
     end
     return out
 end
-∇logsoftmax(Δ, xs) = ∇softmax(Δ ./ max.(eps(eltype(xs)),softmax(xs)), xs)
+∇logsoftmax(Δ, xs) = Δ - sum(Δ, dims=1) .* softmax(xs)
 ∇logsoftmax!(Δ, xs) = ∇softmax!(Δ, Δ, xs)

test/activation.jl

@@ -100,7 +100,7 @@ end
         xs = Float32[1 2 3; 1000 2000 3000]
         @test logsoftmax(xs) ≈ [-999 -1998 -2997; 0 0 0.]
 
-        @test NNlib.∇logsoftmax(ones(size(xs)), xs) ≈ zeros(Float32, size(xs))
+        @test NNlib.∇logsoftmax(ones(size(xs)), xs) ≈ Float32[1 1 1; -1 -1 -1] 
         @test NNlib.∇softmax(ones(size(xs)), xs) ≈ zeros(Float32, size(xs))
 
         # These values precalculated using PyTorch's nn.LogSoftmax