Help with MAP optimization

[sdwfrost]

I'm trying to get MAP optimization working on my toy example. Normally, my observations are Poisson distributed, but to get AD working, I'm using a continuous version of the Poisson that interpolates the CDF. Here's what I have so far, but the parameter values don't seem to change - can someone point me in the direction of my error?

# Libraries
using OrdinaryDiffEq
using SciMLSensitivity
using Distributions
using Random
using Gen
using GenDistributions

# Define continuous Poisson distribution

using StatsFuns, Distributions, ADTypes, FiniteDiff
import DifferentiationInterface as DI

struct ContinuousPoisson{T<:Real,B<:ADTypes.AbstractADType} <: ContinuousUnivariateDistribution
    λ::T
    ad_type::B
end

ContinuousPoisson(λ) = ContinuousPoisson(λ, ADTypes.AutoFiniteDiff())

discretize(d::ContinuousPoisson) = Poisson(d.λ)

Base.minimum(::ContinuousPoisson) = 0
Base.maximum(::ContinuousPoisson) = Inf
Distributions.mode(d::ContinuousPoisson) = d.λ-1
Distributions.cdf(d::ContinuousPoisson, x::Real) = StatsFuns.gammaccdf(x + 1, 1, d.λ)

function Distributions.logcdf(d::ContinuousPoisson, x::Real)
    return StatsFuns.gammalogccdf(x + 1, 1, d.λ)
end

function Distributions.quantile(d::ContinuousPoisson, p::Real)
    upper = Distributions.quantile(discretize(d), p)
    lower = upper - 1
    # NOTE: quantile_bisect is an internal function
    return Distributions.quantile_bisect(d, p, lower, upper)
end

function Distributions.pdf(d::ContinuousPoisson, x::Real)
    return DI.derivative(Base.Fix1(Distributions.cdf, d), d.ad_type, x)
end
        
function Distributions.logpdf(d::ContinuousPoisson, x::Real)
    return log(DI.derivative(Base.Fix1(Distributions.logcdf, d), d.ad_type, x)) + logcdf(d, x)
end

const cpoisson = DistributionsBacked((λ) -> ContinuousPoisson(λ),
                                             (true,),
                                             true,
                                             Float64)

function sir_ode!(du,u,p,t)
        (S,I,R,C) = u
        (β,c,γ) = p
        N = S+I+R
        infection = β*c*I/N*S
        recovery = γ*I
        @inbounds begin
            du[1] = -infection
            du[2] = infection - recovery
            du[3] = recovery
            du[4] = infection
        end
        nothing
end

function cases_from_solution(sol)
    sol_C = Array(sol)[4, :] # Cumulative cases
    sol_X = abs.(sol_C[2:end] - sol_C[1:(end-1)]) # New cases
    return sol_X
end

@gen function sir_ode_model_continuous(l::Int=40,
    N::Float64=1000.0,
    c::Float64=10.0,
    γ::Float64=0.25,
    δt::Float64=1.0)
    i₀ ~ uniform(0.001, 0.1)
    β ~ uniform(0.01, 0.1)
    I = i₀ * N
    u0 = [N - I, I, 0.0, 0.0]
    p = [β, c, γ]
    tspan = (0.0, float(l))
    prob = ODEProblem(sir_ode!, u0, tspan, p)
    sol = solve(prob, Tsit5(), saveat = δt)
    sol_X = cases_from_solution(sol) .- 1.0 # Continuous Poisson domain is from -1 to Inf
    for i in 1:l
        {(:y, i)} ~ cpoisson(sol_X[i])
    end
    return sol
end

p = Gen.choicemap()
p[:β] = 0.05
p[:i₀] = 0.01
fixed_args = (1000.0, 10.0, 0.25, 1.0)
l = 40
(sol, _) = Gen.generate(sir_ode_model_continuous, (l,fixed_args...), p)
Y = [sol[(:y, i)] for i=1:l]
observations = Gen.choicemap()
for (i, y) in enumerate(Y)
    observations[(:y, i)] = y
end

p0 = Gen.choicemap()
p0[:β] = 0.1
p0[:i₀] = 0.025
(tr,) = Gen.generate(sir_ode_model_continuous, (l,fixed_args...), merge(observations, p0))

niter = 10
for i in 1:niter
    global tr
    tr = map_optimize(tr, select(:β,:i₀), max_step_size=0.1, min_step_size=0.0001)
    tr, = update(tr, (l,fixed_args...), (NoChange(),), choicemap(:β => tr[:β], :i₀ => tr[:i₀]))
end

tr[:β],tr[:i₀]

[ztangent]

Hello! Looking at this code, it's not obvious to me if anything is wrong. Have you tried using map_optimize with a simpler model? Also, I think it might be the case that it's enough to call map_optimize only once, and that any subsequent calls will not change the parameters (unless you e.g. added more data).

[ztangent]

Hmm, okay, my best guess is that the gradient code for cpoisson isn't working correctly. And so it may be worth manually implementing logpdf_grad for cpoisson instead of relying on DistributionsBacked. Another thing you could try is to use a regular discrete poisson (and use discrete observations), for which we know logpdf_grad is implemented correctly.

[sdwfrost]

I originally tried a regular Poisson, but I get an error that it can't work out the gradient:

# Libraries
using OrdinaryDiffEq
using SciMLSensitivity
using Distributions
using Random
using Gen

function sir_ode!(du,u,p,t)
        (S,I,R,C) = u
        (β,c,γ) = p
        N = S+I+R
        infection = β*c*I/N*S
        recovery = γ*I
        @inbounds begin
            du[1] = -infection
            du[2] = infection - recovery
            du[3] = recovery
            du[4] = infection
        end
        nothing
end

function cases_from_solution(sol)
    sol_C = Array(sol)[4, :] # Cumulative cases
    sol_X = abs.(sol_C[2:end] - sol_C[1:(end-1)]) # New cases
    return sol_X
end

@gen function sir_ode_model_discrete(l::Int=40,
    N::Float64=1000.0,
    c::Float64=10.0,
    γ::Float64=0.25,
    δt::Float64=1.0)
    i₀ ~ uniform(0.001, 0.1)
    β ~ uniform(0.01, 0.1)
    I = i₀ * N
    u0 = [N - I, I, 0.0, 0.0]
    p = [β, c, γ]
    tspan = (0.0, float(l))
    prob = ODEProblem(sir_ode!, u0, tspan, p)
    sol = solve(prob, Tsit5(), saveat = δt)
    sol_X = cases_from_solution(sol)
    for i in 1:l
        {(:y, i)} ~ poisson(sol_X[i])
    end
    return sol
end

p = Gen.choicemap()
p[:β] = 0.05
p[:i₀] = 0.01
fixed_args = (1000.0, 10.0, 0.25, 1.0)
l = 40
(sol, _) = Gen.generate(sir_ode_model_discrete, (l,fixed_args...), p)
Y = [sol[(:y, i)] for i=1:l]
observations = Gen.choicemap()
for (i, y) in enumerate(Y)
    observations[(:y, i)] = y
end

p0 = Gen.choicemap()
p0[:β] = 0.051
p0[:i₀] = 0.009
(tr,) = Gen.generate(sir_ode_model_discrete, (l,fixed_args...), merge(observations, p0))
tr_opt = map_optimize(tr, select(:β,:i₀), max_step_size=0.1, min_step_size=0.0001)

This generates:

ERROR: Gradient required but not available for argument 1 of Gen.Poisson()
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] special_reverse_exec!(instruction::ReverseDiff.SpecialInstruction{…})
   @ Gen ~/.julia/packages/Gen/ME5el/src/dynamic/backprop.jl:26
 [3] reverse_exec!(instruction::ReverseDiff.SpecialInstruction{Gen.Poisson, Tuple{…}, ReverseDiff.TrackedReal{…}, Nothing})
   @ ReverseDiff ~/.julia/packages/ReverseDiff/p1MzG/src/tape.jl:93
 [4] reverse_pass!(tape::Vector{ReverseDiff.AbstractInstruction})
   @ ReverseDiff ~/.julia/packages/ReverseDiff/p1MzG/src/tape.jl:87
 [5] reverse_pass!
   @ ~/.julia/packages/Gen/ME5el/src/backprop.jl:12 [inlined]
 [6] choice_gradients(trace::Gen.DynamicDSLTrace{DynamicDSLFunction{Any}}, selection::DynamicSelection, retval_grad::Nothing)
   @ Gen ~/.julia/packages/Gen/ME5el/src/dynamic/backprop.jl:382
 [7] map_optimize(trace::Gen.DynamicDSLTrace{…}, selection::DynamicSelection; max_step_size::Float64, tau::Float64, min_step_size::Float64, verbose::Bool)
   @ Gen ~/.julia/packages/Gen/ME5el/src/inference/map_optimize.jl:15

[ztangent]

Okay this looks like it was a bug on our end! It looks like this line was incorrectly implemented:
https://github.com/probcomp/Gen.jl/blob/81bb2cf49c3e05a0a862f7c23fe0d14ccbc295e3/src/modeling_library/distributions/poisson.jl#L27C1-L27C41

I can make a PR to fix it (thanks for finding it), but in the meantime you can manually workaround this issue by redefining Gen.has_argument_grads in your script:

Gen.has_argument_grads(::Gen.Poisson) = (true,)

Hopefully that is the only issue!

[sdwfrost]

Thanks @ztangent! That fixed the compilation error, but now I have the same issue as with my continuous Poisson example, where the parameter values aren't updated.

To follow up with your other suggestion, I then tried a toy example with a Poisson with a conjugate gamma prior; this time, the optimization does do something (the optimized parameters are different than the initial conditions), but it gives the wrong answer (it should be close to 5 (4.992260077399226 in the code below) but comes out as 7.84044921875:

using Distributions
using Random
using Gen
Random.seed!(1234)

Gen.has_argument_grads(::Gen.Poisson) = (true,)

function poisson_map(Y, α, β)
    n = length(Y)
    λ = (α + sum(Y) -1)/(β+n)
    return λ
end

@gen function poisson_model(l::Int=40, α::Float64=0.1, β::Float64=0.1)
    λ ~ gamma(α,β)
    for i in 1:l
        {(:y, i)} ~ poisson(λ)
    end
    nothing
end

p = Gen.choicemap()
p[:λ] = 5.0
l = 10000
α = 0.1
β = 0.1
fixed_args = (α, β)
(sol, _) = Gen.generate(poisson_model, (l,fixed_args...), p)
Y = [sol[(:y, i)] for i=1:l]
print(poisson_map(Y,α,β))

observations = Gen.choicemap()
for (i, y) in enumerate(Y)
    observations[(:y, i)] = y
end
p0 = Gen.choicemap()
p0[:λ] = 2.0
(tr,) = Gen.generate(poisson_model, (l,fixed_args...), merge(observations, p0))
tr_opt = map_optimize(tr, select(:λ), max_step_size=0.1, min_step_size=0.0001)
print(tr_opt[:λ])

[ztangent]

Does calling map_optimize multiple times on the toy example help things at all? I think there's a chance that one call to map_optimize may not be enough (internally, it gives up once no further improvements can be made taking a small step size, but resetting to the larger step size can sometimes still help).

Another option might be to initialize λ to a different value, if the posterior happens to have local optima.

If neither of those address the issue, I'm not sure what it could be --- perhaps the gradient returned by our Poisson distribution is incorrect?

[sdwfrost]

I get a few problems with the toy example. I can set \lambda to 4.9, but the MAP returned by map_optimize is incorrect (>5 but not as bad). If I set to something a little larger (say 5.1), I get DomainErrors; the same is true when I try to iterate multiple times - I get DomainErrors that prevent the optimization from returning.

[ztangent]

Hmm, the DomainErrors are likely because λ is being set to something below zero in the process of optimization, such that the density is now undefined since λ is required to be greater than zero. We haven't figured out a great way to address these issues yet, beyond manually implementing a variant of of Gen.Poisson() which is parameterized by log(lambda) directly instead of being parameterized by lambda (see #531 for a related issue).

Here's something to check. If a trace is initialized with the value of λ that is exactly at the optimum, and you call choice_gradients on that trace to get the gradient of the log density with respect to \lambda, is the gradient returned zero? If so, then the gradient code is operating correctly, and the failure to converge to the optimum is due to some limitation of the optimization procedure implemented by map_optimize (which is maybe fixable by changing the step size, etc.).

(As far as I can tell, our implementation of logpdf_grad for both the gamma and poisson distributions is correct, so I think the gradient code should be behaving correctly...)

[sdwfrost]

logpdf_grad for Poisson and gamma look good to me too. The gradient isn't zero at the value returned by map_optimize; to test the robustness of the solver, I wrapped the trace to use with Optimization.jl:

using Optimization
using OptimizationOptimJL
function poisson_map_opt(u,p)
    λ0 = u[1]
    cm = Gen.choicemap()
    cm[:λ] = λ0
    (tr,) = Gen.generate(poisson_model, (l,fixed_args...), merge(observations, cm))
    -get_score(tr)
end

optf = OptimizationFunction(poisson_map_opt)
prob = OptimizationProblem(optf, [2.0], [])
sol = solve(prob, NelderMead())
sol.u[1]

This does a better job and gets a gradient close to zero when plugged back into the trace (I can probably do better if I also wrap the gradient and use a gradient based solver). I am worried though that the MAP estimates from optimizing the score (4.987) are different from the analytic one (4.992, not a huge difference, but there). With the above code, this isn't due to the optimizer getting stuck. I think my analytical MAP is correct, but then that would mean my understanding of get_score is off...

[ztangent]

Ah I guess I was wondering if choice_gradients returns a zero gradient if you set lambda to the analytic MAP? I would hope that sanity check would pass at least.

My guess is that the particular gradient descent algorithm used in map_optimize isn't super well optimized (especially in cases when gradients change wildly with small differences in \lambda, which I expect they are doing here), and if you manually implemented some other gradient-based optimization method using choice_gradients.

Looking at the implementation more carefully, I believe the algorithm used by map_optimize is a very simplified version of backtracking line search --- in particular, map_optimize will return new trace once a it finds a step size that leads to an improvement in get_score, and not when it reaches the local minimum. A more sophisticated implementation of backtracking line search would use a more careful stopping condition. And in either case, you would need to iteratively call map_optimize to eventually reach the local optimum.

[sdwfrost]

As an addendum, the Turing.jl results are pretty similar to Gen.jl, which makes me think it's mostly a numerical issue with the map_estimate solver.

using Turing
@model function poisson_model_turing(Y,α,β)
    # Calculate number of timepoints
    l = length(Y)
    λ ~ Gamma(α,β)
    for i in eachindex(Y)
      Y[i] ~ Poisson(λ)
    end
end

model = poisson_model_turing(Y, α, β)
map_estimate = maximum_a_posteriori(model)

ModeResult with maximized lp of -22040.44
[4.987322677322676]

[sdwfrost]

I still don't know why the results are off the analytical MAP though...

[sdwfrost]

I'm trying to find an answer here. so we can close this thread; my feeling is, is that map_optimize isn't robust, and the issue of MAP estimation would be better served by providing an OptimizationFunction to Optimization.jl (despite the additional cost in dependencies)

[ztangent]

Yes, I think that's right. I'm not sure whether we'll want to add a new dependency vs. implement a better algorithm in-house vs. have a tutorial that shows how users can write their own custom optimization algorithm using choice_gradients, but you're welcome to file an issue that map_optimize doesn't converge to the local optimum in cases like the one you showed!