diff --git a/HISTORY.md b/HISTORY.md index dc66f1f496..6e23fd6243 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,5 +1,55 @@ # 0.42.0 +## Breaking Changes + +**AdvancedVI 0.5** + +Turing.jl v0.42 updates `AdvancedVI.jl` compatibility to 0.5. +Most of the changes introduced in `AdvancedVI.jl@0.5` are structural, with some changes spilling out into the interface. +The summary of the changes below are the things that affect the end-users of Turing. +For a more comprehensive list of changes, please refer to the [changelogs](https://github.com/TuringLang/AdvancedVI.jl/blob/main/HISTORY.md) in `AdvancedVI`. + +A new level of interface for defining different variational algorithms has been introduced in `AdvancedVI` v0.5. As a result, the function `Turing.vi` now receives a keyword argument `algorithm`. The object `algorithm <: AdvancedVI.AbstractVariationalAlgorithm` should now contain all the algorithm-specific configurations. Therefore, keyword arguments of `vi` that were algorithm-specific such as `objective`, `operator`, `averager` and so on, have been moved as fields of the relevant `<: AdvancedVI.AbstractVariationalAlgorithm` structs. +For example, + +```julia +vi(model, q, n_iters; objective=RepGradELBO(10), operator=AdvancedVI.ClipScale()) +``` + +is now + +```julia +vi( + model, + q, + n_iters; + algorithm=KLMinRepGradDescent(adtype; n_samples=10, operator=AdvancedVI.ClipScale()), +) +``` + +Similarly, + +```julia +vi( + model, + q, + n_iters; + objective=RepGradELBO(10; entropy=AdvancedVI.ClosedFormEntropyZeroGradient()), + operator=AdvancedVI.ProximalLocationScaleEntropy(), +) +``` + +is now + +```julia +vi(model, q, n_iters; algorithm=KLMinRepGradProxDescent(adtype; n_samples=10)) +``` + +Additionally, + + - The default hyperparameters of `DoG`and `DoWG` have been altered. + - The deprecated `AdvancedVI@0.2`-era interface is now removed. + # 0.41.0 ## DynamicPPL 0.38 diff --git a/Project.toml b/Project.toml index 6db0028523..bbddd9508f 100644 --- a/Project.toml +++ b/Project.toml @@ -55,7 +55,7 @@ Accessors = "0.1" AdvancedHMC = "0.3.0, 0.4.0, 0.5.2, 0.6, 0.7, 0.8" AdvancedMH = "0.8" AdvancedPS = "0.7" -AdvancedVI = "0.4" +AdvancedVI = "0.5" BangBang = "0.4.2" Bijectors = "0.14, 0.15" Compat = "4.15.0" diff --git a/src/Turing.jl b/src/Turing.jl index 58a58eb2af..a4f40df259 100644 --- a/src/Turing.jl +++ b/src/Turing.jl @@ -116,10 +116,12 @@ export externalsampler, # Variational inference - AdvancedVI vi, - ADVI, q_locationscale, q_meanfield_gaussian, q_fullrank_gaussian, + KLMinRepGradProxDescent, + KLMinRepGradDescent, + KLMinScoreGradDescent, # ADTypes AutoForwardDiff, AutoReverseDiff, diff --git a/src/variational/VariationalInference.jl b/src/variational/VariationalInference.jl index d516319684..630d3b62f2 100644 --- a/src/variational/VariationalInference.jl +++ b/src/variational/VariationalInference.jl @@ -1,21 +1,24 @@ module Variational -using DynamicPPL +using AdvancedVI: + AdvancedVI, KLMinRepGradDescent, KLMinRepGradProxDescent, KLMinScoreGradDescent using ADTypes +using Bijectors: Bijectors using Distributions +using DynamicPPL using LinearAlgebra using LogDensityProblems using Random +using ..Turing: DEFAULT_ADTYPE, PROGRESS -import ..Turing: DEFAULT_ADTYPE, PROGRESS - -import AdvancedVI -import Bijectors - -export vi, q_locationscale, q_meanfield_gaussian, q_fullrank_gaussian - -include("deprecated.jl") +export vi, + q_locationscale, + q_meanfield_gaussian, + q_fullrank_gaussian, + KLMinRepGradProxDescent, + KLMinRepGradDescent, + KLMinScoreGradDescent """ q_initialize_scale( @@ -248,76 +251,66 @@ end """ vi( [rng::Random.AbstractRNG,] - model::DynamicPPL.Model; + model::DynamicPPL.Model, q, - n_iterations::Int; - objective::AdvancedVI.AbstractVariationalObjective = AdvancedVI.RepGradELBO( - 10; entropy = AdvancedVI.ClosedFormEntropyZeroGradient() + max_iter::Int; + adtype::ADTypes.AbstractADType=DEFAULT_ADTYPE, + algorithm::AdvancedVI.AbstractVariationalAlgorithm = KLMinRepGradProxDescent( + adtype; n_samples=10 ), show_progress::Bool = Turing.PROGRESS[], - optimizer::Optimisers.AbstractRule = AdvancedVI.DoWG(), - averager::AdvancedVI.AbstractAverager = AdvancedVI.PolynomialAveraging(), - operator::AdvancedVI.AbstractOperator = AdvancedVI.ProximalLocationScaleEntropy(), - adtype::ADTypes.AbstractADType = Turing.DEFAULT_ADTYPE, kwargs... ) -Approximating the target `model` via variational inference by optimizing `objective` with the initialization `q`. +Approximate the target `model` via the variational inference algorithm `algorithm` by starting from the initial variational approximation `q`. This is a thin wrapper around `AdvancedVI.optimize`. +The default `algorithm`, `KLMinRepGradProxDescent` ([relevant docs](https://turinglang.org/AdvancedVI.jl/dev/klminrepgradproxdescent/)), assumes `q` uses `AdvancedVI.MvLocationScale`, which can be constructed by invoking `q_fullrank_gaussian` or `q_meanfield_gaussian`. +For other variational families, refer to `AdvancedVI` to determine the best algorithm and options. # Arguments - `model`: The target `DynamicPPL.Model`. - `q`: The initial variational approximation. -- `n_iterations`: Number of optimization steps. +- `max_iter`: Maximum number of steps. # Keyword Arguments -- `objective`: Variational objective to be optimized. +- `adtype`: Automatic differentiation backend to be applied to the log-density. The default value for `algorithm` also uses this backend for differentiation the variational objective. +- `algorithm`: Variational inference algorithm. - `show_progress`: Whether to show the progress bar. -- `optimizer`: Optimization algorithm. -- `averager`: Parameter averaging strategy. -- `operator`: Operator applied after each optimization step. -- `adtype`: Automatic differentiation backend. See the docs of `AdvancedVI.optimize` for additional keyword arguments. # Returns -- `q`: Variational distribution formed by the last iterate of the optimization run. -- `q_avg`: Variational distribution formed by the averaged iterates according to `averager`. -- `state`: Collection of states used for optimization. This can be used to resume from a past call to `vi`. -- `info`: Information generated during the optimization run. +- `q`: Output variational distribution of `algorithm`. +- `state`: Collection of states used by `algorithm`. This can be used to resume from a past call to `vi`. +- `info`: Information generated while executing `algorithm`. """ function vi( rng::Random.AbstractRNG, model::DynamicPPL.Model, q, - n_iterations::Int; - objective=AdvancedVI.RepGradELBO( - 10; entropy=AdvancedVI.ClosedFormEntropyZeroGradient() + max_iter::Int, + args...; + adtype::ADTypes.AbstractADType=DEFAULT_ADTYPE, + algorithm::AdvancedVI.AbstractVariationalAlgorithm=KLMinRepGradProxDescent( + adtype; n_samples=10 ), show_progress::Bool=PROGRESS[], - optimizer=AdvancedVI.DoWG(), - averager=AdvancedVI.PolynomialAveraging(), - operator=AdvancedVI.ProximalLocationScaleEntropy(), - adtype::ADTypes.AbstractADType=DEFAULT_ADTYPE, kwargs..., ) return AdvancedVI.optimize( rng, - LogDensityFunction(model), - objective, + algorithm, + max_iter, + LogDensityFunction(model; adtype), q, - n_iterations; + args...; show_progress=show_progress, - adtype, - optimizer, - averager, - operator, kwargs..., ) end -function vi(model::DynamicPPL.Model, q, n_iterations::Int; kwargs...) - return vi(Random.default_rng(), model, q, n_iterations; kwargs...) +function vi(model::DynamicPPL.Model, q, max_iter::Int; kwargs...) + return vi(Random.default_rng(), model, q, max_iter; kwargs...) end end diff --git a/src/variational/deprecated.jl b/src/variational/deprecated.jl deleted file mode 100644 index 9a9f4777b5..0000000000 --- a/src/variational/deprecated.jl +++ /dev/null @@ -1,61 +0,0 @@ - -import DistributionsAD -export ADVI - -Base.@deprecate meanfield(model) q_meanfield_gaussian(model) - -struct ADVI{AD} - "Number of samples used to estimate the ELBO in each optimization step." - samples_per_step::Int - "Maximum number of gradient steps." - max_iters::Int - "AD backend used for automatic differentiation." - adtype::AD -end - -function ADVI( - samples_per_step::Int=1, - max_iters::Int=1000; - adtype::ADTypes.AbstractADType=ADTypes.AutoForwardDiff(), -) - Base.depwarn( - "The type ADVI will be removed in future releases. Please refer to the new interface for `vi`", - :ADVI; - force=true, - ) - return ADVI{typeof(adtype)}(samples_per_step, max_iters, adtype) -end - -function vi(model::DynamicPPL.Model, alg::ADVI; kwargs...) - Base.depwarn( - "This specialization along with the type `ADVI` will be deprecated in future releases. Please refer to the new interface for `vi`.", - :vi; - force=true, - ) - q = q_meanfield_gaussian(Random.default_rng(), model) - objective = AdvancedVI.RepGradELBO( - alg.samples_per_step; entropy=AdvancedVI.ClosedFormEntropy() - ) - operator = AdvancedVI.IdentityOperator() - _, q_avg, _, _ = vi(model, q, alg.max_iters; objective, operator, kwargs...) - return q_avg -end - -function vi( - model::DynamicPPL.Model, - alg::ADVI, - q::Bijectors.TransformedDistribution{<:DistributionsAD.TuringDiagMvNormal}; - kwargs..., -) - Base.depwarn( - "This specialization along with the type `ADVI` will be deprecated in future releases. Please refer to the new interface for `vi`.", - :vi; - force=true, - ) - objective = AdvancedVI.RepGradELBO( - alg.samples_per_step; entropy=AdvancedVI.ClosedFormEntropy() - ) - operator = AdvancedVI.IdentityOperator() - _, q_avg, _, _ = vi(model, q, alg.max_iters; objective, operator, kwargs...) - return q_avg -end diff --git a/test/Project.toml b/test/Project.toml index 435f8cc5f2..8d819a6749 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -44,7 +44,7 @@ AbstractMCMC = "5" AbstractPPL = "0.11, 0.12, 0.13" AdvancedMH = "0.6, 0.7, 0.8" AdvancedPS = "0.7" -AdvancedVI = "0.4" +AdvancedVI = "0.5" Aqua = "0.8" BangBang = "0.4" Bijectors = "0.14, 0.15" @@ -53,6 +53,7 @@ Combinatorics = "1" Distributions = "0.25" DistributionsAD = "0.6.3" DynamicHMC = "2.1.6, 3.0" +DynamicPPL = "0.38" FiniteDifferences = "0.10.8, 0.11, 0.12" ForwardDiff = "0.10.12 - 0.10.32, 0.10, 1" HypothesisTests = "0.11" diff --git a/test/variational/advi.jl b/test/variational/vi.jl similarity index 60% rename from test/variational/advi.jl rename to test/variational/vi.jl index ed8f745df2..b426f0e6a3 100644 --- a/test/variational/advi.jl +++ b/test/variational/vi.jl @@ -10,12 +10,16 @@ using Distributions: Dirichlet, Normal using LinearAlgebra using MCMCChains: Chains using Random +using ReverseDiff using StableRNGs: StableRNG using Test: @test, @testset using Turing using Turing.Variational @testset "ADVI" begin + adtype = AutoReverseDiff() + operator = AdvancedVI.ClipScale() + @testset "q initialization" begin m = gdemo_default d = length(Turing.DynamicPPL.VarInfo(m)[:]) @@ -41,86 +45,50 @@ using Turing.Variational @testset "default interface" begin for q0 in [q_meanfield_gaussian(gdemo_default), q_fullrank_gaussian(gdemo_default)] - _, q, _, _ = vi(gdemo_default, q0, 100; show_progress=Turing.PROGRESS[]) + q, _, _ = vi(gdemo_default, q0, 100; show_progress=Turing.PROGRESS[], adtype) c1 = rand(q, 10) end end - @testset "custom interface $name" for (name, objective, operator, optimizer) in [ - ( - "ADVI with closed-form entropy", - AdvancedVI.RepGradELBO(10), - AdvancedVI.ProximalLocationScaleEntropy(), - AdvancedVI.DoG(), - ), - ( - "ADVI with proximal entropy", - AdvancedVI.RepGradELBO(10; entropy=AdvancedVI.ClosedFormEntropyZeroGradient()), - AdvancedVI.ClipScale(), - AdvancedVI.DoG(), - ), - ( - "ADVI with STL entropy", - AdvancedVI.RepGradELBO(10; entropy=AdvancedVI.StickingTheLandingEntropy()), - AdvancedVI.ClipScale(), - AdvancedVI.DoG(), - ), + @testset "custom algorithm $name" for (name, algorithm) in [ + ("KLMinRepGradProxDescent", KLMinRepGradProxDescent(adtype; n_samples=10)), + ("KLMinRepGradDescent", KLMinRepGradDescent(adtype; operator, n_samples=10)), ] T = 1000 - q, q_avg, _, _ = vi( + q, _, _ = vi( gdemo_default, q_meanfield_gaussian(gdemo_default), T; - objective, - optimizer, - operator, + algorithm, + adtype, show_progress=Turing.PROGRESS[], ) - N = 1000 - c1 = rand(q_avg, N) c2 = rand(q, N) end - @testset "inference $name" for (name, objective, operator, optimizer) in [ - ( - "ADVI with closed-form entropy", - AdvancedVI.RepGradELBO(10), - AdvancedVI.ProximalLocationScaleEntropy(), - AdvancedVI.DoG(), - ), - ( - "ADVI with proximal entropy", - RepGradELBO(10; entropy=AdvancedVI.ClosedFormEntropyZeroGradient()), - AdvancedVI.ClipScale(), - AdvancedVI.DoG(), - ), - ( - "ADVI with STL entropy", - AdvancedVI.RepGradELBO(10; entropy=AdvancedVI.StickingTheLandingEntropy()), - AdvancedVI.ClipScale(), - AdvancedVI.DoG(), - ), + @testset "inference $name" for (name, algorithm) in [ + ("KLMinRepGradProxDescent", KLMinRepGradProxDescent(adtype; n_samples=10)), + ("KLMinRepGradDescent", KLMinRepGradDescent(adtype; operator, n_samples=10)), ] rng = StableRNG(0x517e1d9bf89bf94f) T = 1000 - q, q_avg, _, _ = vi( + q, _, _ = vi( rng, gdemo_default, q_meanfield_gaussian(gdemo_default), T; - optimizer, + algorithm, + adtype, show_progress=Turing.PROGRESS[], ) N = 1000 - for q_out in [q_avg, q] - samples = transpose(rand(rng, q_out, N)) - chn = Chains(reshape(samples, size(samples)..., 1), ["s", "m"]) + samples = transpose(rand(rng, q, N)) + chn = Chains(reshape(samples, size(samples)..., 1), ["s", "m"]) - check_gdemo(chn; atol=0.5) - end + check_gdemo(chn; atol=0.5) end # regression test for: @@ -143,7 +111,7 @@ using Turing.Variational @test all(x0 .≈ x0_inv) # And regression for https://github.com/TuringLang/Turing.jl/issues/2160. - _, q, _, _ = vi(rng, m, q_meanfield_gaussian(m), 1000) + q, _, _ = vi(rng, m, q_meanfield_gaussian(m), 1000; adtype) x = rand(rng, q, 1000) @test mean(eachcol(x)) ≈ [0.5, 0.5] atol = 0.1 end @@ -158,7 +126,7 @@ using Turing.Variational end model = demo_issue2205() | (y=1.0,) - _, q, _, _ = vi(rng, model, q_meanfield_gaussian(model), 1000) + q, _, _ = vi(rng, model, q_meanfield_gaussian(model), 1000; adtype) # True mean. mean_true = 1 / 2 var_true = 1 / 2