update vi interface to match [email protected]

Author: Red-Portal

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

@@ -117,6 +117,9 @@ export
     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,61 @@ 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;
+        algorithm::AdvancedVI.AbstractVariationalAlgorithm = KLMinRepGradProxDescent(DEFAULT_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` 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.
+- `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.
+- `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.
 
 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()
-    ),
-    show_progress::Bool=PROGRESS[],
-    optimizer=AdvancedVI.DoWG(),
-    averager=AdvancedVI.PolynomialAveraging(),
-    operator=AdvancedVI.ProximalLocationScaleEntropy(),
+    max_iter::Int,
+    args...;
     adtype::ADTypes.AbstractADType=DEFAULT_ADTYPE,
+    algorithm=KLMinRepGradProxDescent(adtype; n_samples=10),
+    show_progress::Bool=PROGRESS[],
     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

src/variational/deprecated.jl

@@ -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/Project.toml

@@ -13,7 +13,7 @@ include("test_utils/models.jl")
 include("test_utils/numerical_tests.jl")
 include("test_utils/sampler.jl")
 
-Turing.setprogress!(false)
+#Turing.setprogress!(false)
 included_paths, excluded_paths = parse_args(ARGS)
 
 # Filter which tests to run and collect timing and allocations information to show in a
@@ -30,55 +30,59 @@ macro timeit_include(path::AbstractString)
 end
 
 @testset "Turing" verbose = true begin
-    @testset "Aqua" begin
-        @timeit_include("Aqua.jl")
-    end
+    # @testset "Aqua" begin
+    #     @timeit_include("Aqua.jl")
+    # end
 
-    @testset "AD" verbose = true begin
-        @timeit_include("ad.jl")
-    end
+    # @testset "AD" verbose = true begin
+    #     @timeit_include("ad.jl")
+    # end
 
-    @testset "essential" verbose = true begin
-        @timeit_include("essential/container.jl")
-    end
+    # @testset "essential" verbose = true begin
+    #     @timeit_include("essential/container.jl")
+    # end
 
-    @testset "samplers (without AD)" verbose = true begin
-        @timeit_include("mcmc/particle_mcmc.jl")
-        @timeit_include("mcmc/emcee.jl")
-        @timeit_include("mcmc/ess.jl")
-        @timeit_include("mcmc/is.jl")
-    end
+    # @testset "samplers (without AD)" verbose = true begin
+    #     @timeit_include("mcmc/particle_mcmc.jl")
+    #     @timeit_include("mcmc/emcee.jl")
+    #     @timeit_include("mcmc/ess.jl")
+    #     @timeit_include("mcmc/is.jl")
+    # end
 
     @timeit TIMEROUTPUT "inference" begin
-        @testset "inference with samplers" verbose = true begin
-            @timeit_include("mcmc/gibbs.jl")
-            @timeit_include("mcmc/hmc.jl")
-            @timeit_include("mcmc/Inference.jl")
-            @timeit_include("mcmc/sghmc.jl")
-            @timeit_include("mcmc/external_sampler.jl")
-            @timeit_include("mcmc/mh.jl")
-            @timeit_include("ext/dynamichmc.jl")
-            @timeit_include("mcmc/repeat_sampler.jl")
-        end
+        # @testset "inference with samplers" verbose = true begin
+        #     @timeit_include("mcmc/gibbs.jl")
+        #     @timeit_include("mcmc/hmc.jl")
+        #     @timeit_include("mcmc/Inference.jl")
+        #     @timeit_include("mcmc/sghmc.jl")
+        #     @timeit_include("mcmc/external_sampler.jl")
+        #     @timeit_include("mcmc/mh.jl")
+        #     @timeit_include("ext/dynamichmc.jl")
+        #     @timeit_include("mcmc/repeat_sampler.jl")
+        # end
 
         @testset "variational algorithms" begin
             @timeit_include("variational/advi.jl")
         end
 
-        @testset "mode estimation" verbose = true begin
-            @timeit_include("optimisation/Optimisation.jl")
-            @timeit_include("ext/OptimInterface.jl")
-        end
+        # @testset "mode estimation" verbose = true begin
+        #     @timeit_include("optimisation/Optimisation.jl")
+        #     @timeit_include("ext/OptimInterface.jl")
+        # end
     end
 
-    @testset "stdlib" verbose = true begin
-        @timeit_include("stdlib/distributions.jl")
-        @timeit_include("stdlib/RandomMeasures.jl")
-    end
+    # @testset "variational optimisers" begin
+    #     @timeit_include("variational/optimisers.jl")
+    # end
 
-    @testset "utilities" begin
-        @timeit_include("mcmc/utilities.jl")
-    end
+    # @testset "stdlib" verbose = true begin
+    #     @timeit_include("stdlib/distributions.jl")
+    #     @timeit_include("stdlib/RandomMeasures.jl")
+    # end
+
+    # @testset "utilities" begin
+    #     @timeit_include("mcmc/utilities.jl")
+    # end
 end
 
 show(TIMEROUTPUT; compact=true, sortby=:firstexec)