Merge pull request #126 from ornithos/logsoftmaxgrad

Improve numerical stability of logsoftmax gradient
FluxML · Jun 12, 2019 · a80bdff · a80bdff
2 parents f5fce7a + d219cdb
commit a80bdff
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/softmax.jl b/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)
diff --git a/test/activation.jl b/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