From 1375f70945331b334a46bdd771ce307b85e16452 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sun, 13 Sep 2026 18:41:05 -0400 Subject: [PATCH] AD changes --- docs/src/performance_tips.md | 18 ++++- src/multi_choice_models/DDM.jl | 87 +++++++++++++++------- test/Project.toml | 1 + test/multi_choice_models/ddm_tests.jl | 102 ++++++++++++++++++++++++++ 4 files changed, 180 insertions(+), 28 deletions(-) diff --git a/docs/src/performance_tips.md b/docs/src/performance_tips.md index 30e23ff0..4e8c0f7d 100644 --- a/docs/src/performance_tips.md +++ b/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. \ No newline at end of file +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. \ No newline at end of file diff --git a/src/multi_choice_models/DDM.jl b/src/multi_choice_models/DDM.jl index b84e0ba8..44790b58 100644 --- a/src/multi_choice_models/DDM.jl +++ b/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) # diff --git a/test/Project.toml b/test/Project.toml index efaa9327..388d9b56 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/multi_choice_models/ddm_tests.jl b/test/multi_choice_models/ddm_tests.jl index 2e02e9d4..b0e59df7 100644 --- a/test/multi_choice_models/ddm_tests.jl +++ b/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