Numerical errors on Turing v0.7.4

[itsdfish]

Hi all-

After updating from Turing .7.3. to .7.4, I encountered tons of numerical errors (> 3000) for several of my models, causing a convergence failure. These models worked for previous versions of the Turing. I compared .7.3 and .7.4 using Julia 1.3.

Here is my setup for .7.3:

(v1.3) pkg> st Turing
    Status `~/.julia/environments/v1.3/Project.toml`
  [0bf59076] AdvancedHMC v0.2.14
  [31c24e10] Distributions v0.21.11
  [189a3867] Reexport v0.2.0
  [4c63d2b9] StatsFuns v0.8.0
  [fce5fe82] Turing v0.7.3 ⚲

and for .7.4 :

(v1.3) pkg> st Turing
    Status `~/.julia/environments/v1.3/Project.toml`
  [0bf59076] AdvancedHMC v0.2.14
  [31c24e10] Distributions v0.21.11
  [189a3867] Reexport v0.2.0
  [4c63d2b9] StatsFuns v0.9.2
  [fce5fe82] Turing v0.7.4

Both versions use AdvancedHMC v0.2.14. So that does not appear to be the problem or at least the exclusive problem.

An example model can be found here and it can be ran with the file called Run baseline Fan model.jl

[itsdfish]

By the way, I discovered that the problem resolved when I downgrade to StatsFuns v0.8.0 and the errors start with v0.9.0. Its not clear to me what in that release is causing trouble with Turing.

[mohamed82008]

If you have a minimal working example, we can look further into it.

[itsdfish]

Sorry about that. Thanks for looking into this more. Here is a much smaller example, indicating the problem is probably with logsumexp:

using Turing, Random, Parameters, StatsFuns
import Distributions: logpdf, rand

struct MyDist{T1} <: ContinuousUnivariateDistribution
  μ::T1
end

function logpdf(d::MyDist{<:T}, data::Array{<:T1,1}) where {T,T1}
    μ = d.μ
    N = length(data)
    LLs = Array{typeof(μ),1}(undef,N)
    for i in 1:N 
        LLs[i] = logpdf(Normal(μ, 1.0), data[i])
    end 
    #return sum(LLs)
    return logsumexp(LLs)
end

Random.seed!(1101)

@model model(y) = begin
    N = length(y)
    μ ~ Normal(0,1)
    y ~ MyDist(μ)
end
y = rand(Normal(0,1),10)

chain = sample(model(y) ,NUTS(1000, .8),  2000)
println(chain)

As before, the model works with StatsFuns v0.8.0. In addition, the model works when logsumexp is replaced with sum regardless of version. I suspect this is the culprit.

[mohamed82008]

@xukai92 any idea what could be happening here?

[mohamed82008]

The problem is that the new logsumexp is not auto-differentiable:

using ForwardDiff, StatsFuns

# New StatsFuns
julia> ForwardDiff.gradient(x -> StatsFuns.logsumexp(x), rand(4))
4-element Array{Float64,1}:
 NaN
 NaN
 NaN
 NaN

# Old StatsFuns
julia> ForwardDiff.gradient(x -> StatsFuns.logsumexp(x), rand(4))
4-element Array{Float64,1}:
 0.1280594678275393
 0.3043198459775647
 0.30254704646302555
 0.2650736397318705

[itsdfish]

The problem begins at the call to reduce in logsumexp:

using ForwardDiff

function max_reduce(X::AbstractArray{T}; dims=:) where {T<:Real}
    u = reduce(max, X, dims=dims, init=log(zero(T)))
end
ForwardDiff.gradient(x ->max_reduce(x),  rand(4))

init=log(zero(T)) seems like it might be problematic.

[itsdfish]

Indeed, removing log from the initial value solves the problem (using v0.8.0):

using ForwardDiff, StatsFuns, Random

function logsumexp1(X::AbstractArray{T}; dims=:) where {T<:Real}
    u = reduce(max, X, dims=dims, init=zero(T))
    u isa AbstractArray || isfinite(u) || return float(u)
    let u=u # avoid https://github.com/JuliaLang/julia/issues/15276
        # TODO: remove the branch when JuliaLang/julia#31020 is merged.
        if u isa AbstractArray
            u .+ log.(sum(exp.(X .- u); dims=dims))
        else
            u + log(sum(x -> exp(x-u), X))
        end
    end
end

Random.seed!(100)

y1 = ForwardDiff.gradient(x -> StatsFuns.logsumexp(x), rand(4))

Random.seed!(100)

y2 = ForwardDiff.gradient(x -> logsumexp1(x), rand(4))

y1 == y2

[itsdfish]

I submitted a pull request with StatsFuns that fixes the error. Thanks for your help!