Update variational inference interface to match AdvancedVI@0.5

HISTORY.md

@@ -1,5 +1,18 @@
 # 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 have been introduced in `AdvancedVI` v0.5. As a result, the method `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.
+  - The default hyperparameters of `DoG`and `DoWG` have been altered.
+  - The depricated `AdvancedVI@0.2`-era interface is now removed.
+
 # 0.41.0
 
 ## DynamicPPL 0.38
@@ -62,7 +75,7 @@ Note that if the initial sample is included, the corresponding sampler statistic
 Due to a technical limitation of MCMCChains.jl, this causes all indexing into MCMCChains to return `Union{Float64, Missing}` or similar.
 If you want the old behaviour, you can discard the first sample (e.g. using `discard_initial=1`).
 
-# 0.40.5
+# 0.4# 0.40.5
 
 Bump Optimization.jl compatibility to include v5.
 

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"

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,

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

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"

test/variational/advi.jl → 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