Minor Touches for ScoreGradELBO

.github/workflows/CI.yml

@@ -19,8 +19,8 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.7'
           - '1.10'
+          - '1'
         os:
           - ubuntu-latest
           - macOS-latest

Project.toml

@@ -14,6 +14,7 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
+OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -49,6 +50,7 @@ Functors = "0.4"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
 Mooncake = "0.4"
+OnlineStats = "1"
 Optimisers = "0.2.16, 0.3"
 ProgressMeter = "1.6"
 Random = "1"

src/AdvancedVI.jl

@@ -6,6 +6,7 @@ using Accessors
 
 using Random
 using Distributions
+using OnlineStats
 
 using Functors
 using Optimisers

src/objectives/elbo/entropy.jl

@@ -37,10 +37,3 @@ function estimate_entropy(
         -logpdf(q, mc_sample)
     end
 end
-
-function estimate_entropy_maybe_stl(
-    entropy_estimator::AbstractEntropyEstimator, samples, q, q_stop
-)
-    q_maybe_stop = maybe_stop_entropy_score(entropy_estimator, q, q_stop)
-    return estimate_entropy(entropy_estimator, samples, q_maybe_stop)
-end

src/objectives/elbo/repgradelbo.jl

@@ -45,6 +45,13 @@ function Base.show(io::IO, obj::RepGradELBO)
     return print(io, ")")
 end
 
+function estimate_entropy_maybe_stl(
+    entropy_estimator::AbstractEntropyEstimator, samples, q, q_stop
+)
+    q_maybe_stop = maybe_stop_entropy_score(entropy_estimator, q, q_stop)
+    return estimate_entropy(entropy_estimator, samples, q_maybe_stop)
+end
+
 function estimate_energy_with_samples(prob, samples)
     return mean(Base.Fix1(LogDensityProblems.logdensity, prob), eachsample(samples))
 end

src/objectives/elbo/scoregradelbo.jl

@@ -1,3 +1,4 @@
+
 """
 	ScoreGradELBO(n_samples; kwargs...)
 
@@ -25,9 +26,8 @@ To reduce the variance of the gradient estimator, we use a baseline computed fro
 - `n_samples::Int`: Number of Monte Carlo samples used to estimate the ELBO.
 
 # Keyword Arguments
-- `entropy`: The estimator for the entropy term. (Type `<: AbstractEntropyEstimator`; Default: `ClosedFormEntropy()`)
+- `entropy`: The estimator for the entropy term. (Type `<: AbstractEntropyEstimator`; Default: `FullMonteCarloEntropy()`)
 - `baseline_window_size::Int`: The window size to use to compute the baseline. (Default: `10`)
-- `baseline_history::Vector{Float64}`: The history of the baseline. (Default: `Float64[]`)
 
 # Requirements
 - The variational approximation ``q_{\\lambda}`` implements `rand` and `logpdf`.
@@ -36,21 +36,24 @@ To reduce the variance of the gradient estimator, we use a baseline computed fro
 
 Depending on the options, additional requirements on ``q_{\\lambda}`` may apply.
 """
-struct ScoreGradELBO{EntropyEst<:AbstractEntropyEstimator} <:
-       AdvancedVI.AbstractVariationalObjective
+struct ScoreGradELBO{EntropyEst<:AbstractEntropyEstimator} <: AbstractVariationalObjective
     entropy::EntropyEst
     n_samples::Int
     baseline_window_size::Int
-    baseline_history::Vector{Float64}
 end
 
 function ScoreGradELBO(
     n_samples::Int;
-    entropy::AbstractEntropyEstimator=ClosedFormEntropy(),
+    entropy::AbstractEntropyEstimator=MonteCarloEntropy(),
     baseline_window_size::Int=10,
-    baseline_history::Vector{Float64}=Float64[],
 )
-    return ScoreGradELBO(entropy, n_samples, baseline_window_size, baseline_history)
+    return ScoreGradELBO(entropy, n_samples, baseline_window_size)
+end
+
+function init(
+    ::Random.AbstractRNG, obj::ScoreGradELBO, prob, params::AbstractVector{T}, restructure
+) where {T<:Real}
+    return MovingWindow(T, obj.baseline_window_size)
 end
 
 function Base.show(io::IO, obj::ScoreGradELBO)
@@ -63,28 +66,11 @@ function Base.show(io::IO, obj::ScoreGradELBO)
     return print(io, ")")
 end
 
-function compute_control_variate_baseline(history, window_size)
-    if length(history) == 0
-        return 1.0
-    end
-    min_index = max(1, length(history) - window_size)
-    return mean(history[min_index:end])
-end
-
-function estimate_energy_with_samples(
-    prob, samples_stop, samples_logprob, samples_logprob_stop, baseline
-)
-    fv = Base.Fix1(LogDensityProblems.logdensity, prob).(eachsample(samples_stop))
-    fv_mean = mean(fv)
-    score_grad = mean(@. samples_logprob * (fv - baseline))
-    score_grad_stop = mean(@. samples_logprob_stop * (fv - baseline))
-    return fv_mean + (score_grad - score_grad_stop)
-end
-
 function estimate_objective(
     rng::Random.AbstractRNG, obj::ScoreGradELBO, q, prob; n_samples::Int=obj.n_samples
 )
-    samples, entropy = reparam_with_entropy(rng, q, q, obj.n_samples, obj.entropy)
+    samples = rand(rng, q, n_samples)
+    entropy = estimate_entropy(obj.entropy, samples, q)
     energy = map(Base.Fix1(LogDensityProblems.logdensity, prob), eachsample(samples))
     return mean(energy) + entropy
 end
@@ -94,18 +80,19 @@ function estimate_objective(obj::ScoreGradELBO, q, prob; n_samples::Int=obj.n_sa
 end
 
 function estimate_scoregradelbo_ad_forward(params′, aux)
-    @unpack rng, obj, problem, adtype, restructure, q_stop = aux
-    baseline = compute_control_variate_baseline(
-        obj.baseline_history, obj.baseline_window_size
-    )
+    @unpack rng, obj, problem, adtype, restructure, samples, q_stop, baseline = aux
     q = restructure_ad_forward(adtype, restructure, params′)
-    samples_stop = rand(rng, q_stop, obj.n_samples)
-    entropy = estimate_entropy_maybe_stl(obj.entropy, samples_stop, q, q_stop)
-    samples_logprob = logpdf.(Ref(q), AdvancedVI.eachsample(samples_stop))
-    samples_logprob_stop = logpdf.(Ref(q_stop), AdvancedVI.eachsample(samples_stop))
-    energy = estimate_energy_with_samples(
-        problem, samples_stop, samples_logprob, samples_logprob_stop, baseline
-    )
+
+    ℓq = logpdf.(Ref(q), AdvancedVI.eachsample(samples))
+    ℓq_stop = logpdf.(Ref(q_stop), AdvancedVI.eachsample(samples))
+    ℓπ = map(Base.Fix1(LogDensityProblems.logdensity, problem), eachsample(samples))
+    ℓπ_mean = mean(ℓπ)
+    score_grad = mean(@. ℓq * (ℓπ - baseline))
+    score_grad_stop = mean(@. ℓq_stop * (ℓπ - baseline))
+
+    energy = ℓπ_mean + (score_grad - score_grad_stop)
+    entropy = estimate_entropy(obj.entropy, samples, q)
+
     elbo = energy + entropy
     return -elbo
 end
@@ -120,20 +107,32 @@ function AdvancedVI.estimate_gradient!(
     restructure,
     state,
 )
+    baseline_buf = state
+    baseline_history = OnlineStats.value(baseline_buf)
+    baseline = if isempty(baseline_history)
+        zero(eltype(params))
+    else
+        mean(baseline_history)
+    end
     q_stop = restructure(params)
+    samples = rand(rng, q_stop, obj.n_samples)
     aux = (
         rng=rng,
         adtype=adtype,
         obj=obj,
         problem=prob,
         restructure=restructure,
+        baseline=baseline,
+        samples=samples,
         q_stop=q_stop,
     )
     AdvancedVI.value_and_gradient!(
         adtype, estimate_scoregradelbo_ad_forward, params, aux, out
     )
     nelbo = DiffResults.value(out)
     stat = (elbo=-nelbo,)
-    push!(obj.baseline_history, -nelbo)
-    return out, nothing, stat
+    if obj.baseline_window_size > 0
+        fit!(baseline_buf, -nelbo)
+    end
+    return out, baseline_buf, stat
 end

test/inference/repgradelbo_locationscale.jl

@@ -15,7 +15,7 @@ if @isdefined(Enzyme)
     )
 end
 
-@testset "inference ScoreGradELBO VILocationScale" begin
+@testset "inference RepGradELBO VILocationScale" begin
     @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
                                                                      [Float64, Float32],
         (modelname, modelconstr) in

test/inference/scoregradelbo_distributionsad.jl

@@ -9,20 +9,16 @@ if @isdefined(Mooncake)
     AD_scoregradelbo_distributionsad[:Mooncake] = AutoMooncake(; config=Mooncake.Config())
 end
 
-#if @isdefined(Enzyme)
-#    AD_scoregradelbo_distributionsad[:Enzyme] = AutoEnzyme()
-#end
+if @isdefined(Enzyme)
+    AD_scoregradelbo_distributionsad[:Enzyme] = AutoEnzyme()
+end
 
 @testset "inference ScoreGradELBO DistributionsAD" begin
     @testset "$(modelname) $(objname) $(realtype) $(adbackname)" for realtype in
                                                                      [Float64, Float32],
         (modelname, modelconstr) in Dict(:Normal => normal_meanfield),
         n_montecarlo in [1, 10],
-        (objname, objective) in Dict(
-            :ScoreGradELBOClosedFormEntropy => ScoreGradELBO(n_montecarlo),
-            :ScoreGradELBOStickingTheLanding =>
-                ScoreGradELBO(n_montecarlo; entropy=StickingTheLandingEntropy()),
-        ),
+        (objname, objective) in Dict(:ScoreGradELBO => ScoreGradELBO(n_montecarlo)),
         (adbackname, adtype) in AD_scoregradelbo_distributionsad
 
         seed = (0x38bef07cf9cc549d)

test/inference/scoregradelbo_locationscale.jl

@@ -21,11 +21,7 @@ end
         (modelname, modelconstr) in
         Dict(:Normal => normal_meanfield, :Normal => normal_fullrank),
         n_montecarlo in [1, 10],
-        (objname, objective) in Dict(
-            :ScoreGradELBOClosedFormEntropy => ScoreGradELBO(n_montecarlo),
-            :ScoreGradELBOStickingTheLanding =>
-                ScoreGradELBO(n_montecarlo; entropy=StickingTheLandingEntropy()),
-        ),
+        (objname, objective) in Dict(:ScoreGradELBO => ScoreGradELBO(n_montecarlo)),
         (adbackname, adtype) in AD_locationscale
 
         seed = (0x38bef07cf9cc549d)

test/inference/scoregradelbo_locationscale_bijectors.jl

@@ -5,9 +5,9 @@ AD_scoregradelbo_locationscale_bijectors = Dict(
     #:Zygote => AutoZygote(),
 )
 
-#if @isdefined(Tapir)
-#    AD_scoregradelbo_locationscale_bijectors[:Tapir] = AutoTapir(; safe_mode=false)
-#end
+if @isdefined(Mooncake)
+    AD_scoregradelbo_locationscale_bijectors[:Mooncake] = AutoMooncake(; config=nothing)
+end
 
 if @isdefined(Enzyme)
     AD_scoregradelbo_locationscale_bijectors[:Enzyme] = AutoEnzyme()
@@ -19,11 +19,7 @@ end
         (modelname, modelconstr) in
         Dict(:NormalLogNormalMeanField => normallognormal_meanfield),
         n_montecarlo in [1, 10],
-        (objname, objective) in Dict(
-            #:ScoreGradELBOClosedFormEntropy => ScoreGradELBO(n_montecarlo), # not supported yet.
-            :ScoreGradELBOStickingTheLanding =>
-                ScoreGradELBO(n_montecarlo; entropy=StickingTheLandingEntropy()),
-        ),
+        (objname, objective) in Dict(:ScoreGradELBO => ScoreGradELBO(n_montecarlo)),
         (adbackname, adtype) in AD_scoregradelbo_locationscale_bijectors
 
         seed = (0x38bef07cf9cc549d)

test/interface/scoregradelbo.jl

@@ -25,4 +25,21 @@ using Test
         elbo = estimate_objective(obj, q0, model; n_samples=10^4)
         @test elbo ≈ elbo_ref rtol = 0.2
     end
+
+    @testset "baseline_window" begin
+        T = 100
+        adtype = AutoForwardDiff()
+
+        obj = ScoreGradELBO(10)
+        _, _, stats, _ = optimize(rng, model, obj, q0, T; show_progress=false, adtype)
+        @test isfinite(last(stats).elbo)
+
+        obj = ScoreGradELBO(10; baseline_window_size=0)
+        _, _, stats, _ = optimize(rng, model, obj, q0, T; show_progress=false, adtype)
+        @test isfinite(last(stats).elbo)
+
+        obj = ScoreGradELBO(10; baseline_window_size=1)
+        _, _, stats, _ = optimize(rng, model, obj, q0, T; show_progress=false, adtype)
+        @test isfinite(last(stats).elbo)
+    end
 end