Improve accuracy and performance of ChainRules derivatives

Project.toml

@@ -1,7 +1,7 @@
 name = "LogExpFunctions"
 uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 authors = ["StatsFun.jl contributors, Tamas K. Papp <tkpapp@gmail.com>"]
-version = "1.0.4"
+version = "1.0.5"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

ext/LogExpFunctionsChainRulesCoreExt.jl

@@ -114,23 +114,55 @@ function ChainRulesCore.rrule(::typeof(xexpy), x::Real, y::Real)
     return Ω, xexpy_pullback
 end
 
-ChainRulesCore.@scalar_rule(logistic(x::Real), (Ω * (1 - Ω),))
+# logistic(x) * logistic(-x), avoiding cancellation in `1 - logistic(x)`
+function _∂_logistic(x::Real)
+    e = exp(-abs(x))
+    return e / (1 + e)^2
+end
+# 2x / (1 + x^2), avoiding overflow of `x^2` for large `|x|`
+_∂_log1psq(x::Real) = abs(x) < 1 ? 2 * x / (1 + x^2) : 2 / (x + inv(x))
+
+ChainRulesCore.@scalar_rule(logistic(x::Real), (_∂_logistic(x),))
 ChainRulesCore.@scalar_rule(logit(x::Real), (inv(x * (1 - x)),))
 ChainRulesCore.@scalar_rule(logcosh(x::Real), tanh(x))
 ChainRulesCore.@scalar_rule(logabssinh(x::Real), coth(x))
-ChainRulesCore.@scalar_rule(logabstanh(x::Real), inv(cosh(x) * sinh(x)))
-ChainRulesCore.@scalar_rule(log1psq(x::Real), (2 * x / (1 + x^2),))
+ChainRulesCore.@scalar_rule(logabstanh(x::Real), 2 / sinh(2 * x))
+ChainRulesCore.@scalar_rule(log1psq(x::Real), (_∂_log1psq(x),))
 ChainRulesCore.@scalar_rule(log1pexp(x::Real), (logistic(x),))
-ChainRulesCore.@scalar_rule(log1mexp(x::Real), (-exp(x - Ω),))
+ChainRulesCore.@scalar_rule(log1mexp(x::Real), (-inv(expm1(-x)),))
 ChainRulesCore.@scalar_rule(log2mexp(x::Real), (-exp(x - Ω),))
-ChainRulesCore.@scalar_rule(logexpm1(x::Real), (exp(x - Ω),))
+ChainRulesCore.@scalar_rule(logexpm1(x::Real), (-inv(expm1(-x)),))
 ChainRulesCore.@scalar_rule(log1pmx(x::Real), (-x / (x + 1),))
 ChainRulesCore.@scalar_rule(logmxp1(x::Real), ((1 - x) / x,))
+ChainRulesCore.@scalar_rule(loglogistic(x::Real), (logistic(-x),))
+ChainRulesCore.@scalar_rule(log1mlogistic(x::Real), (-logistic(x),))
+ChainRulesCore.@scalar_rule(logitexp(x::Real), (-inv(expm1(x)),))
+ChainRulesCore.@scalar_rule(logit1mexp(x::Real), (inv(expm1(x)),))
 
-ChainRulesCore.@scalar_rule(logaddexp(x::Real, y::Real), (exp(x - Ω), exp(y - Ω)))
+# the one with respect to the larger argument as `1 - t` (no cancellation since
+# 0 ≤ t ≤ 1/2 for `logaddexp` and t ≤ 0 for `logsubexp`).
+# The handling of `x == y` mirrors the primal implementations.
+# (0 ≤ t ≤ 1/2 for `logaddexp` and t ≤ 0 for `logsubexp`, so `1 - t` is accurate).
+# `x - y` is `NaN` if `x == y == ±Inf`.
+function _∂_logaddexp(x::Real, y::Real)
+    Δ = x == y && isinf(x) ? zero(x - y) : abs(x - y)
+    t = logistic(-Δ)
+    return x >= y ? (1 - t, t) : (t, 1 - t)
+end
+function _∂_logsubexp(x::Real, y::Real)
+    Δ = x == y == -Inf ? zero(x - y) : abs(x - y)
+    t = -inv(expm1(Δ))
+    return x >= y ? (1 - t, t) : (t, 1 - t)
+end
+ChainRulesCore.@scalar_rule(
+    logaddexp(x::Real, y::Real),
+    @setup((∂x, ∂y) = _∂_logaddexp(x, y)),
+    (∂x, ∂y),
+)
 ChainRulesCore.@scalar_rule(
     logsubexp(x::Real, y::Real),
-    (x > y ? exp(x - Ω) : -exp(x - Ω), x > y ? -exp(y - Ω) : exp(y - Ω)),
+    @setup((∂x, ∂y) = _∂_logsubexp(x, y)),
+    (∂x, ∂y),
 )
 
 function ChainRulesCore.frule((_, Δx), ::typeof(logsumexp), x::AbstractArray{<:Real}; dims=:)
@@ -185,7 +217,7 @@ function ChainRulesCore.rrule(::typeof(softmax), x::AbstractArray{<:Real}; dims=
     return Ω, softmax_pullback
 end
 
-ChainRulesCore.@scalar_rule(cloglog(x), (-inv((1 - x) * log1p(-x)),))
-ChainRulesCore.@scalar_rule(cexpexp(x), (-xexpx(-exp(x)),))
+ChainRulesCore.@scalar_rule(cloglog(x::Real), (-inv((1 - x) * log1p(-x)),))
+ChainRulesCore.@scalar_rule(cexpexp(x::Real), (-xexpx(-exp(x)),))
 
 end # module

test/chainrules.jl

@@ -48,6 +48,11 @@
         test_frule(log1psq, x)
         test_rrule(log1psq, x)
     end
+    # no overflow of `x^2`
+    for x in (-1e200, 1e200)
+        @test last(frule((NoTangent(), 1.0), log1psq, x)) ≈ 2 / x rtol=4*eps()
+        @test last(last(rrule(log1psq, x))(1.0)) ≈ 2 / x rtol=4*eps()
+    end
 
     # test all `Float64` and `Float32` branches of `logistic`
     for x in (-821.4, -23.5, 12.3, 41.2)
@@ -58,6 +63,13 @@
         test_frule(logistic, x; rtol=1f-3, atol=1f-3)
         test_rrule(logistic, x; rtol=1f-3, atol=1f-3)
     end
+    # derivative is accurate for large `|x|` in a relative sense (no cancellation)
+    for x in (40.0, -40.0, 30f0, -30f0)
+        e = exp(-abs(big(x)))
+        ∂x = e / (1 + e)^2
+        @test last(frule((NoTangent(), one(x)), logistic, x)) ≈ ∂x rtol=4*eps(typeof(x))
+        @test last(last(rrule(logistic, x))(one(x))) ≈ ∂x rtol=4*eps(typeof(x))
+    end
 
     for x in (-randexp(), randexp())
         test_frule(logcosh, x)
@@ -94,6 +106,25 @@
         test_frule(logexpm1, x; rtol=1f-3, atol=1f-3)
         test_rrule(logexpm1, x; rtol=1f-3, atol=1f-3)
     end
+    # special case (manually since FiniteDifferences/ChainRulesTestUtils fails at Inf)
+    @test frule((NoTangent(), 1.0), logexpm1, Inf) === (Inf, 1.0)
+    @test last(last(rrule(logexpm1, Inf))(1.0)) === 1.0
+
+    # accuracy for small `|x|` (where `Ω` is large in magnitude)
+    for (f, x) in ((log1mexp, -1e-12), (logexpm1, 6.039463507942533e-8))
+        ∂x = -inv(expm1(-big(x)))
+        @test last(frule((NoTangent(), 1.0), f, x)) ≈ ∂x rtol=4*eps()
+        @test last(last(rrule(f, x))(1.0)) ≈ ∂x rtol=4*eps()
+    end
+
+    for x in (-randexp(), randexp())
+        test_scalar(loglogistic, x)
+        test_scalar(log1mlogistic, x)
+    end
+    for x in (-10.2, -3.3, -0.3)
+        test_scalar(logitexp, x)
+        test_scalar(logit1mexp, x)
+    end
 
     test_scalar(log1pmx, rand())
 
@@ -106,6 +137,55 @@
         test_frule(logsubexp, x, y)
         test_rrule(logsubexp, x, y)
     end
+    # nearly equal arguments (manually since finite differences are inaccurate close to
+    # the singularity at `x == y`)
+    for (x, y) in ((1.0, 1.0 + 1e-6), (1.0 + 1e-6, 1.0))
+        d = big(x) - big(y)
+        ∂x, ∂y = -inv(expm1(-d)), -inv(expm1(d))
+        @test last(frule((NoTangent(), 1.0, 0.0), logsubexp, x, y)) ≈ ∂x rtol=4*eps()
+        @test last(frule((NoTangent(), 0.0, 1.0), logsubexp, x, y)) ≈ ∂y rtol=4*eps()
+        _, back = rrule(logsubexp, x, y)
+        _, x̄, ȳ = back(1.0)
+        @test x̄ ≈ ∂x rtol=4*eps()
+        @test ȳ ≈ ∂y rtol=4*eps()
+    end
+    # accuracy for arguments of large magnitude
+    let (x, y) = (300.0, 301.0)
+        d = big(x) - big(y)
+        ∂x, ∂y = logistic(d), logistic(-d)
+        _, back = rrule(logaddexp, x, y)
+        _, x̄, ȳ = back(1.0)
+        @test x̄ ≈ ∂x rtol=4*eps()
+        @test ȳ ≈ ∂y rtol=4*eps()
+    end
+    # special cases (manually since FiniteDifferences/ChainRulesTestUtils fails at Inf)
+    for (f, ∂x, ∂y) in ((logaddexp, 1.0, 0.0), (logsubexp, 1.0, -0.0))
+        Ω, back = rrule(f, Inf, 1.0)
+        @test Ω === Inf
+        @test back(1.0) === (NoTangent(), ∂x, ∂y)
+        Ω, back = rrule(f, 1.0, Inf)
+        @test Ω === Inf
+        @test back(1.0) === (NoTangent(), ∂y, ∂x)
+        @test last(frule((NoTangent(), 1.0, 0.0), f, Inf, 1.0)) === ∂x
+    end
+    # equal and `NaN` arguments (`x == y == ±Inf`: #61)
+    for T in (Float32, Float64)
+        for x in (T(-Inf), one(T), T(Inf))
+            _, back = rrule(logaddexp, x, x)
+            @test back(one(T)) === (NoTangent(), T(0.5), T(0.5))
+        end
+        # limits as `x` approaches `y` from above
+        for x in (T(-Inf), one(T))
+            _, back = rrule(logsubexp, x, x)
+            @test back(one(T)) === (NoTangent(), T(Inf), T(-Inf))
+        end
+        _, back = rrule(logsubexp, T(Inf), T(Inf))
+        @test all(isnan, back(one(T))[2:3])
+        for f in (logaddexp, logsubexp)
+            _, back = rrule(f, T(NaN), one(T))
+            @test all(isnan, back(one(T))[2:3])
+        end
+    end
 
     for x in (randn(10), randn(10, 8)), dims in (:, 1, 1:2, 2)
         dims isa Colon || all(d <= ndims(x) for d in dims) || continue