Mode Estimation

After defining a statistical model, in addition to sampling from its distributions, one may be interested in finding the parameter values that maximise for instance the posterior distribution density function or the likelihood. This is called mode estimation. Turing provides support for two mode estimation techniques, maximum likelihood estimation (MLE) and maximum a posterior (MAP) estimation.

To demonstrate mode estimation, let us load Turing and declare a model:

using Turing

@model function gdemo(x)
    s² ~ InverseGamma(2, 3)
    m ~ Normal(0, sqrt(s²))

    for i in eachindex(x)
        x[i] ~ Normal(m, sqrt(s²))
    end
end
Precompiling Turing...
    803.0 ms  ? OptimizationBase
   1356.3 ms  ? Optimization
   2102.0 ms  ? OptimizationOptimJL
Info Given Turing was explicitly requested, output will be shown live 
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
   5434.3 ms  ? Turing
   5719.0 ms  ? Turing → TuringOptimExt
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
Precompiling Optimization...
    791.5 ms  ? OptimizationBase
Info Given Optimization was explicitly requested, output will be shown live 
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
   1395.4 ms  ? Optimization
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
Precompiling OptimizationBase...
Info Given OptimizationBase was explicitly requested, output will be shown live 
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
    766.5 ms  ? OptimizationBase
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: Method definition (::Type{OptimizationBase.OptimizerMissingError})(Any) in module OptimizationBase at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:23 overwritten at /home/runner/.julia/packages/OptimizationBase/sfIfa/src/solve.jl:177.
ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
┌ Warning: Replacing docs for `CommonSolve.solve :: Tuple{SciMLBase.OptimizationProblem, Any, Vararg{Any}}` in module `OptimizationBase`
└ @ Base.Docs docs/Docs.jl:243
WARNING: redefinition of constant OptimizationBase.OPTIMIZER_MISSING_ERROR_MESSAGE. This may fail, cause incorrect answers, or produce other errors.
┌ Warning: Replacing docs for `CommonSolve.init :: Tuple{SciMLBase.OptimizationProblem, Any, Vararg{Any}}` in module `OptimizationBase`
└ @ Base.Docs docs/Docs.jl:243
┌ Warning: Replacing docs for `CommonSolve.solve! :: Tuple{SciMLBase.AbstractOptimizationCache}` in module `OptimizationBase`
└ @ Base.Docs docs/Docs.jl:243
Precompiling OptimizationOptimJL...
    763.6 ms  ? OptimizationBase
    978.4 ms  ? Optimization
Info Given OptimizationOptimJL was explicitly requested, output will be shown live 
┌ Warning: Module Optimization with build ID ffffffff-ffff-ffff-ec3b-fe843fd89b77 is missing from the cache.
│ This may mean Optimization [7f7a1694-90dd-40f0-9382-eb1efda571ba] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:2541
   1118.3 ms  ? OptimizationOptimJL
┌ Warning: Module Optimization with build ID ffffffff-ffff-ffff-ec3b-fe843fd89b77 is missing from the cache.
│ This may mean Optimization [7f7a1694-90dd-40f0-9382-eb1efda571ba] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:2541
Precompiling TuringOptimExt...
    789.0 ms  ? OptimizationBase
    962.7 ms  ? Optimization
   1099.5 ms  ? OptimizationOptimJL
   3714.9 ms  ? Turing
Info Given TuringOptimExt was explicitly requested, output will be shown live 
┌ Warning: Module Turing with build ID ffffffff-ffff-ffff-7940-bef3fac244f3 is missing from the cache.
│ This may mean Turing [fce5fe82-541a-59a6-adf8-730c64b5f9a0] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:2541
    627.3 ms  ? Turing → TuringOptimExt
┌ Warning: Module Turing with build ID ffffffff-ffff-ffff-7940-bef3fac244f3 is missing from the cache.
│ This may mean Turing [fce5fe82-541a-59a6-adf8-730c64b5f9a0] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:2541

gdemo (generic function with 2 methods)

Once the model is defined, we can construct a model instance as we normally would:

# Instantiate the gdemo model with our data.
data = [1.5, 2.0]
model = gdemo(data)
DynamicPPL.Model{typeof(gdemo), (:x,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}(gdemo, (x = [1.5, 2.0],), NamedTuple(), DynamicPPL.DefaultContext())

Finding the maximum aposteriori or maximum likelihood parameters is as simple as

# Generate a MLE estimate.
mle_estimate = maximum_likelihood(model)

# Generate a MAP estimate.
map_estimate = maximum_a_posteriori(model)
ModeResult with maximized lp of -4.62
[0.9074074072016184, 1.1666666667008674]

The estimates are returned as instances of the ModeResult type. It has the fields values for the parameter values found and lp for the log probability at the optimum, as well as f for the objective function and optim_result for more detailed results of the optimisation procedure.

@show mle_estimate.values
@show mle_estimate.lp;
mle_estimate.values = [0.06250000009880316, 1.7500000002538696]
mle_estimate.lp = -0.06528834416956464

Controlling the optimisation process

Under the hood maximum_likelihood and maximum_a_posteriori use the Optimization.jl package, which provides a unified interface to many other optimisation packages. By default Turing typically uses the LBFGS method from Optim.jl to find the mode estimate, but we can easily change that:

using OptimizationOptimJL: NelderMead
@show maximum_likelihood(model, NelderMead())

using OptimizationNLopt: NLopt.LD_TNEWTON_PRECOND_RESTART
@show maximum_likelihood(model, LD_TNEWTON_PRECOND_RESTART());
maximum_likelihood(model, NelderMead()) = [0.062492750810579485, 1.750000846443205]
Precompiling OptimizationNLopt...
    824.0 ms  ? OptimizationBase
   1058.6 ms  ? Optimization
Info Given OptimizationNLopt was explicitly requested, output will be shown live 
┌ Warning: Module Optimization with build ID ffffffff-ffff-ffff-ec3b-fe843fd89b77 is missing from the cache.
│ This may mean Optimization [7f7a1694-90dd-40f0-9382-eb1efda571ba] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:2541
    805.9 ms  ? OptimizationNLopt
┌ Warning: Module Optimization with build ID ffffffff-ffff-ffff-ec3b-fe843fd89b77 is missing from the cache.
│ This may mean Optimization [7f7a1694-90dd-40f0-9382-eb1efda571ba] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:2541
┌ Warning: The selected optimization algorithm requires second order derivatives, but `SecondOrder` ADtype was not provided. 
│         So a `SecondOrder` with ADTypes.AutoForwardDiff() for both inner and outer will be created, this can be suboptimal and not work in some cases so 
│         an explicit `SecondOrder` ADtype is recommended.
└ @ OptimizationBase ~/.julia/packages/OptimizationBase/sfIfa/src/cache.jl:49
maximum_likelihood(model, LD_TNEWTON_PRECOND_RESTART()) = [0.062499999912016886, 1.7499999998263396]

The above are just two examples, Optimization.jl supports many more.

We can also help the optimisation by giving it a starting point we know is close to the final solution, or by specifying an automatic differentiation method

import Mooncake

maximum_likelihood(
    model, NelderMead(); initial_params=[0.1, 2], adtype=AutoMooncake()
)
ModeResult with maximized lp of -0.07
[0.062494553692639856, 1.7500042095865365]

When providing values to arguments like initial_params the parameters are typically specified in the order in which they appear in the code of the model, so in this case first s² then m. More precisely it’s the order returned by Turing.Inference.getparams(model, DynamicPPL.VarInfo(model)).

We can also do constrained optimisation, by providing either intervals within which the parameters must stay, or costraint functions that they need to respect. For instance, here’s how one can find the MLE with the constraint that the variance must be less than 0.01 and the mean must be between -1 and 1.:

maximum_likelihood(model; lb=[0.0, -1.0], ub=[0.01, 1.0])
ModeResult with maximized lp of -59.73
[0.009999999897804623, 0.9999999958099894]

The arguments for lower (lb) and upper (ub) bounds follow the arguments of Optimization.OptimizationProblem, as do other parameters for providing constraints, such as cons. Any extraneous keyword arguments given to maximum_likelihood or maximum_a_posteriori are passed to Optimization.solve. Some often useful ones are maxiters for controlling the maximum number of iterations and abstol and reltol for the absolute and relative convergence tolerances:

badly_converged_mle = maximum_likelihood(
    model, NelderMead(); maxiters=10, reltol=1e-9
)
ModeResult with maximized lp of -0.08
[0.06109376297463792, 1.7833750157096984]

We can check whether the optimisation converged using the optim_result field of the result:

@show badly_converged_mle.optim_result;
badly_converged_mle.optim_result = retcode: Failure
u: [-2.79534549698983, 1.7833750157096984]
Final objective value:     0.08378174908238689

For more details, such as a full list of possible arguments, we encourage the reader to read the docstring of the function Turing.Optimisation.estimate_mode, which is what maximum_likelihood and maximum_a_posteriori call, and the documentation of Optimization.jl.

Analyzing your mode estimate

Turing extends several methods from StatsBase that can be used to analyze your mode estimation results. Methods implemented include vcov, informationmatrix, coeftable, params, and coef, among others.

For example, let’s examine our ML estimate from above using coeftable:

using StatsBase: coeftable
coeftable(mle_estimate)
Coef. Std. Error z Pr(> z )
s² 0.0625 0.0625 1.0 0.317311 -0.0599977 0.184998
m 1.75 0.176777 9.89949 4.18383e-23 1.40352 2.09648

Standard errors are calculated from the Fisher information matrix (inverse Hessian of the log likelihood or log joint). Note that standard errors calculated in this way may not always be appropriate for MAP estimates, so please be cautious in interpreting them.

Sampling with the MAP/MLE as initial states

You can begin sampling your chain from an MLE/MAP estimate by extracting the vector of parameter values and providing it to the sample function with the keyword initial_params. For example, here is how to sample from the full posterior using the MAP estimate as the starting point:

map_estimate = maximum_a_posteriori(model)
chain = sample(model, NUTS(), 1_000; initial_params=map_estimate.values.array)
Back to top