Compute DDM log density in log space for stable AD and Turing sampling

docs/src/performance_tips.md

@@ -12,4 +12,20 @@ Turing provides three general recommendations for developing performant code:
 2. Use Multivariate distributions in place of Univariate distributions when applicable. 
 3. Use forward mode automatic differentiation when your model has a small number of parameters (i.e., 5-10), and use reverse mode automatic differentiation for larger models. 
 
-See the [Turing documentation](https://turinglang.org/docs/tutorials/docs-13-using-turing-performance-tips/) for more details. Note that the Turing ecosystem provides a benchmarking package, [TuringBenchmarking.jl](https://turinglang.org/TuringBenchmarking.jl/dev/) to aid in the selection of an automatic differentiation backend. 
+See the [Turing documentation](https://turinglang.org/docs/tutorials/docs-13-using-turing-performance-tips/) for more details. Note that the Turing ecosystem provides a benchmarking package, [TuringBenchmarking.jl](https://turinglang.org/TuringBenchmarking.jl/dev/) to aid in the selection of an automatic differentiation backend.
+
+## Automatic Differentiation Backends
+
+Do not use a compiled ReverseDiff tape (`AutoReverseDiff(; compile = true)`) with the `DDM`.
+
+A compiled ReverseDiff tape reuses the calculation steps recorded on its first run. The DDM changes its calculation and number of series terms as parameters change, so reusing those steps can give incorrect gradients without an error. This also affects other models whose series length depends on their parameters.
+
+The following timings measure one gradient evaluation for a hierarchical DDM log likelihood:
+
+| parameters | observations | ForwardDiff | ReverseDiff (compiled) | Mooncake |
+|---|---|---|---|---|
+| 53 | 1,000 | 3.8 ms | 9.8 ms | **1.2 ms** |
+| 153 | 3,000 | 38.2 ms | 34.7 ms | **3.8 ms** |
+| 403 | 8,000 | 290.1 ms | 98.4 ms | **11.3 ms** |
+
+Mooncake was fastest in these benchmarks and supports calculations that change with the parameters. Use `AutoMooncake` for larger models and `AutoForwardDiff` for small ones. Uncompiled ReverseDiff (`AutoReverseDiff(; compile = false)`) also works, but is slower.

src/multi_choice_models/DDM.jl

@@ -68,20 +68,52 @@ end
 # Wabersich & Vandekerckhove (2014) #
 #####################################
 
-function pdf(d::DDM{T}, choice, rt; ϵ::Real = 1.0e-12) where {T <: Real}
-    @argcheck d.τ < rt
+function pdf(d::DDM, choice, rt; ϵ::Real = 1.0e-12)
+    return exp(logpdf(d, choice, rt; ϵ))
+end
+
+# At α = 0 or z = 0 or 1, the process stops at τ, so the density is zero for rt > τ.
+# Round before comparing so ForwardDiff checks only the value, ignoring derivatives.
+# Keep the rounded values as floats: converting a large α to Int can overflow.
+_is_degenerate(d::DDM) = ceil(d.α) ≤ 0 || ceil(d.z) ≤ 0 || floor(d.z) ≥ 1
+
+function logpdf(d::DDM, choice, rt; ϵ::Real = 1.0e-12)
+    (ν, α, z, τ) = params(d)
+    # Return -Inf with zero gradient so the sampler can reject this proposal.
+    if rt ≤ τ || _is_degenerate(d)
+        return promote_type(typeof(ν), typeof(rt))(-Inf)
+    end
     if choice == 1
-        (ν, α, z, τ) = params(d)
-        return _pdf(DDM(-ν, α, 1 - z, τ), rt; ϵ)
+        return _logpdf(DDM(-ν, α, 1 - z, τ), rt; ϵ)
     end
-    return _pdf(d, rt; ϵ)
+    return _logpdf(d, rt; ϵ)
 end
 
-# probability density function over the lower boundary
-function _pdf(d::DDM, t::Real; ϵ::Real = 1.0e-12)
+logpdf(d::DDM, data::Tuple) = logpdf(d, data...)
+
+# Log density at the lower boundary.
+# Compute in log space to keep tail values and gradients finite when pdf is too small.
+function _logpdf(d::DDM, t::Real; ϵ::Real = 1.0e-12)
     (ν, α, z, τ) = params(d)
     u = (t - τ) / α^2 #use normalized time
 
+    # Very large α can make u round to zero. Avoid dividing by zero below.
+    # As in _is_degenerate, compare the value without its derivatives.
+    ceil(u) ≤ 0 && return promote_type(typeof(u), typeof(z))(-Inf)
+
+    K_s, K_l = _n_terms_pdf(u, ϵ)
+
+    log_p = -α * z * ν - 0.5 * ν^2 * (t - τ) - 2 * log(α)
+
+    # Use the series that needs fewer terms.
+    if K_s < K_l
+        return log_p + _log_small_time_pdf(u, z, ceil(Int, K_s))
+    end
+    return log_p + _log_large_time_pdf(u, z, ceil(Int, K_l))
+end
+
+# Terms needed for each series at accuracy ϵ (Navarro & Fuss, 2009).
+function _n_terms_pdf(u::Real, ϵ::Real)
     K_s = 2.0
     K_l = 1 / (π * sqrt(u))
     # number of terms needed for large-time expansion
@@ -92,43 +124,44 @@ function _pdf(d::DDM, t::Real; ϵ::Real = 1.0e-12)
     if (2 * sqrt(2 * π * u) * ϵ) < 1
         K_s = max(2 + sqrt(-2u * log(2ϵ * sqrt(2 * π * u))), sqrt(u) + 1)
     end
-
-    p = exp((-α * z * ν) - (0.5 * (ν^2) * (t - τ))) / (α^2)
-
-    # decision rule for infinite sum algorithm
-    if K_s < K_l
-        return p * _small_time_pdf(u, z, ceil(Int, K_s))
-    end
-    return p * _large_time_pdf(u, z, ceil(Int, K_l))
+    return K_s, K_l
 end
 
-# small-time expansion
-function _small_time_pdf(u::T, z::T, K::Int) where {T <: Real}
+# log of the small-time expansion
+function _log_small_time_pdf(u::Real, z::Real, K::Int)
+    T = promote_type(typeof(u), typeof(z))
     inf_sum = zero(T)
 
+    # Factor out the largest exponential (k = 0) to avoid underflow.
+    shift = -z^2 / (2u)
+
     k_series = (-floor(Int, 0.5 * (K - 1))):ceil(Int, 0.5 * (K - 1))
     for k ∈ k_series
-        inf_sum += ((2k + z) * exp(-((2k + z)^2 / (2u))))
+        # Subtract before dividing so k = 0 gives zero even when u is tiny.
+        a = 2k + z
+        inf_sum += (a * exp((z^2 - a^2) / (2u)))
     end
 
-    return inf_sum / sqrt(2π * u^3)
+    inf_sum ≤ 0 && return T(-Inf)
+    return shift - 0.5 * log(2π) - 1.5 * log(u) + log(inf_sum)
 end
 
-# large-time expansion
-function _large_time_pdf(u::T, z::T, K::Int) where {T <: Real}
+# log of the large-time expansion
+function _log_large_time_pdf(u::Real, z::Real, K::Int)
+    T = promote_type(typeof(u), typeof(z))
     inf_sum = zero(T)
 
+    # Factor out the largest exponential (k = 1) to avoid underflow.
+    shift = -0.5 * π^2 * u
+
     for k ∈ 1:K
-        inf_sum += (k * exp(-0.5 * (k^2 * π^2 * u)) * sin(k * π * z))
+        inf_sum += (k * exp(-0.5 * (k^2 - 1) * π^2 * u) * sin(k * π * z))
     end
 
-    return π * inf_sum
+    inf_sum ≤ 0 && return T(-Inf)
+    return shift + log(π) + log(inf_sum)
 end
 
-logpdf(d::DDM, choice, rt; ϵ::Real = 1.0e-12) = log(pdf(d, choice, rt; ϵ))
-
-logpdf(d::DDM, data::Tuple) = logpdf(d, data...)
-
 #########################################
 # Cumulative density function           #
 # Blurton, Kesselmeier, & Gondan (2012) #

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 KernelDensity = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b"

test/multi_choice_models/ddm_tests.jl

@@ -292,6 +292,108 @@
         @test test_val4 ≈ 0.4450956 atol = 1e-5
     end
 
+    @safetestset "logpdf" begin
+        using SequentialSamplingModels
+        using Test
+
+        # logpdf agrees with log(pdf) for ordinary density values.
+        for ν ∈ -3:0.5:3, α ∈ [0.4, 0.8, 1.5, 2.5], z ∈ [0.2, 0.5, 0.8],
+            rt ∈ [0.35, 0.5, 1.0, 2.0], choice ∈ (1, 2)
+
+            dist = DDM(; ν, α, z, τ = 0.3)
+            @test logpdf(dist, choice, rt) ≈ log(pdf(dist, choice, rt)) atol = 1e-10
+        end
+
+        # logpdf stays finite when pdf is too small to represent accurately.
+        # Reference values use the original small-time series at 400-bit precision.
+        @test pdf(DDM(; ν = 6.0, α = 6.0, z = 0.2, τ = 0.2), 2, 0.201) < 1e-300
+        for (ν, α, z, τ, rt, true_value) ∈ [
+            (6.0, 6.0, 0.2, 0.2, 0.201, -717.592984057937),
+            (5.0, 4.0, 0.95, 0.05, 0.06, -733.8011821874904),
+            (4.0, 1.5, 0.8, 0.2, 0.201, -715.182984057937)
+        ]
+            @test logpdf(DDM(; ν, α, z, τ), 2, rt) ≈ true_value atol = 1e-9
+        end
+
+        # outside the support
+        dist = DDM(; ν = 1.0, α = 0.8, z = 0.5, τ = 0.3)
+        @test logpdf(dist, 1, 0.3) == -Inf
+        @test logpdf(dist, 1, 0.2) == -Inf
+        @test pdf(dist, 1, 0.2) == 0
+    end
+
+    @safetestset "automatic differentiation" begin
+        using ForwardDiff
+        using SequentialSamplingModels
+        using Test
+
+        f(p, choice, rt) = logpdf(DDM(p[1], p[2], p[3], p[4]), choice, rt)
+
+        # Gradients stay finite even when pdf is too small to differentiate reliably.
+        for (ν, α, z, τ, rt) ∈ [
+            (6.0, 6.0, 0.2, 0.2, 0.201),
+            (5.0, 4.0, 0.95, 0.05, 0.06),
+            (4.0, 1.5, 0.8, 0.2, 0.201),
+            (6.0, 0.3, 0.05, 0.05, 10.05)
+        ],
+            choice ∈ (1, 2)
+
+            grad = ForwardDiff.gradient(p -> f(p, choice, rt), [ν, α, z, τ])
+            @test all(isfinite, grad)
+        end
+
+        # Compare gradients with finite differences across the switch between series.
+        for α ∈ range(0.5, 1.5, length = 21), choice ∈ (1, 2)
+            p = [1.0, α, 0.5, 0.3]
+            ad = ForwardDiff.gradient(x -> f(x, choice, 0.6), p)
+            h = 1e-6
+            fd = map(eachindex(p)) do i
+                lo = copy(p)
+                hi = copy(p)
+                lo[i] -= h
+                hi[i] += h
+                (f(hi, choice, 0.6) - f(lo, choice, 0.6)) / (2h)
+            end
+            @test ad ≈ fd atol = 1e-4
+        end
+
+        # out of support the gradient is zero rather than NaN
+        @test all(iszero, ForwardDiff.gradient(p -> f(p, 1, 0.2), [1.0, 0.8, 0.5, 0.3]))
+    end
+
+    @safetestset "degenerate parameters" begin
+        using ForwardDiff
+        using SequentialSamplingModels
+        using Test
+
+        f(p, choice, rt) = logpdf(DDM(p[1], p[2], p[3], p[4]), choice, rt)
+
+        # Boundary parameters give logpdf = -Inf and zero gradients.
+        for p ∈ [[1.0, 0.0, 0.5, 0.3], [1.0, 0.8, 0.0, 0.3]], choice ∈ (1, 2)
+            @test logpdf(DDM(p[1], p[2], p[3], p[4]), choice, 0.6) == -Inf
+            @test pdf(DDM(p[1], p[2], p[3], p[4]), choice, 0.6) == 0
+            @test all(iszero, ForwardDiff.gradient(x -> f(x, choice, 0.6), p))
+        end
+        # The constructor rejects ForwardDiff inputs at z = 1; check only logpdf here.
+        for choice ∈ (1, 2)
+            @test logpdf(DDM(; ν = 1.0, α = 0.8, z = 1.0, τ = 0.3), choice, 0.6) == -Inf
+        end
+
+        # Parameters just inside the boundaries still give finite log densities.
+        for z ∈ [1e-16, 1 - 1e-16], choice ∈ (1, 2)
+            @test isfinite(logpdf(DDM(; ν = 1.0, α = 0.8, z, τ = 0.3), choice, 0.6))
+        end
+        @test isfinite(logpdf(DDM(; ν = 1.0, α = 1e-16, z = 0.5, τ = 0.3), 2, 0.6))
+
+        # Huge α makes normalized time round to zero; this used to throw InexactError.
+        @test logpdf(DDM(; ν = 1.0, α = 1e300, z = 0.5, τ = 0.3), 2, 0.6) == -Inf
+        # A large α still gives a finite log density when normalized time is nonzero.
+        @test isfinite(logpdf(DDM(; ν = 1.0, α = 1e19, z = 0.5, τ = 0.3), 2, 0.6))
+        # Tiny decision times used to give NaN from Inf - Inf in the k = 0 term.
+        @test logpdf(DDM(; ν = 1.0, α = 1.0, z = 0.5, τ = 0.0), 2, 5e-324) == -Inf
+        @test logpdf(DDM(; ν = 1.0, α = 1.0, z = 0.5, τ = 0.0), 2, 1e-200) ≈ -1.25e199
+    end
+
     @safetestset "simulate" begin
         using SequentialSamplingModels
         using Test