Merge pull request #75 from pevnak/softmax

Softmax

Author: MikeInnes

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,11 @@ 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))
+
+xs = [-0.238639 0.748142 -0.283194 -0.525461 -1.5348 -0.797842; 0.690384 0.211427 0.254794 -0.213572 -0.314174 -0.372663; -1.14637 -0.577988 0.718952 0.91972 -0.620773 0.929977]
+@test isapprox(NNlib.∇logsoftmax(ones(size(xs)), xs), [0.237703 -0.621474 0.448193 0.546047 0.564185 0.632273; -0.930163 0.0519798 0.0549979 0.3799 -0.477112 0.437428; 0.69246 0.569494 -0.503191 -0.925947 -0.0870738 -1.0697]; rtol = 1e-6)
+@test isapprox(NNlib.∇softmax(ones(size(xs)), xs), zeros(size(xs)); atol = 1e-6)
 end