made softmax and logsoftmax more stable with respect to overflow

Author: pevnak

src/logsoftmax.jl

@@ -23,7 +23,7 @@ end
 logsoftmax!(xs) = logsoftmax!(xs, xs)
 logsoftmax(xs) = logsoftmax!(similar(xs), xs)
 
-∇logsoftmax(Δ, xs) = ∇softmax(Δ ./ softmax(xs), xs)
+∇logsoftmax(Δ, xs) = ∇softmax(Δ ./ max.(eps(eltype(xs)),softmax(xs)), xs)
 
 """
     logsoftmax(xs) = log.(exp.(xs) ./ sum(exp.(xs)))

src/softmax.jl

@@ -29,8 +29,8 @@ softmax!(xs) = softmax!(xs, xs)
 softmax(xs) = softmax!(similar(xs), xs)
 
 function ∇softmax!(out::AbstractVecOrMat, Δ::AbstractVecOrMat, xs::AbstractVecOrMat)
-  s = sum(exp, xs, dims=1)
-  out .= exp.(xs)./s.*(Δ .- sum(Δ .* exp.(xs), dims=1)./s)
+  sf = softmax(xs)
+  out .= sf .* (Δ .- sum(Δ .*sf, dims = 1))
 end
 
 ∇softmax!(Δ, xs) = ∇softmax!(Δ, Δ, xs)

test/runtests.jl

@@ -29,4 +29,6 @@ xs = Float32[1, 2, 3000.]
 
 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.∇softmax(ones(size(xs)), xs) ≈ zeros(Float32, size(xs))
 end