From 7ad9f251da4ec678692bfb1bb287402d6b14da21 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sat, 23 Aug 2025 16:12:44 -0400 Subject: [PATCH 1/4] Use concrete types for objective states, reorder arguments --- src/algorithms/paramspacesgd/abstractobjective.jl | 4 ++-- src/algorithms/paramspacesgd/paramspacesgd.jl | 2 +- src/algorithms/paramspacesgd/repgradelbo.jl | 13 +++++++++---- src/algorithms/paramspacesgd/scoregradelbo.jl | 9 +++++++-- 4 files changed, 19 insertions(+), 9 deletions(-) diff --git a/src/algorithms/paramspacesgd/abstractobjective.jl b/src/algorithms/paramspacesgd/abstractobjective.jl index e395319bd..bfa36f9d3 100644 --- a/src/algorithms/paramspacesgd/abstractobjective.jl +++ b/src/algorithms/paramspacesgd/abstractobjective.jl @@ -59,7 +59,7 @@ function estimate_objective end export estimate_objective """ - estimate_gradient!(rng, obj, adtype, out, params, restructure, obj_state) + estimate_gradient!(rng, obj, adtype, out, obj_state, params, restructure) Estimate (possibly stochastic) gradients of the variational objective `obj` targeting `prob` with respect to the variational parameters `λ` @@ -68,9 +68,9 @@ Estimate (possibly stochastic) gradients of the variational objective `obj` targ - `obj::AbstractVariationalObjective`: Variational objective. - `adtype::ADTypes.AbstractADType`: Automatic differentiation backend. - `out::DiffResults.MutableDiffResult`: Buffer containing the objective value and gradient estimates. +- `obj_state`: Previous state of the objective. - `params`: Variational parameters to evaluate the gradient on. - `restructure`: Function that reconstructs the variational approximation from `params`. -- `obj_state`: Previous state of the objective. # Returns - `out::MutableDiffResult`: Buffer containing the objective value and gradient estimates. diff --git a/src/algorithms/paramspacesgd/paramspacesgd.jl b/src/algorithms/paramspacesgd/paramspacesgd.jl index ef4116e78..5d01c10ca 100644 --- a/src/algorithms/paramspacesgd/paramspacesgd.jl +++ b/src/algorithms/paramspacesgd/paramspacesgd.jl @@ -91,7 +91,7 @@ function step( params, re = Optimisers.destructure(q) grad_buf, obj_st, info = estimate_gradient!( - rng, objective, adtype, grad_buf, params, re, obj_st, objargs... + rng, objective, adtype, grad_buf, obj_st, params, re, objargs... ) grad = DiffResults.gradient(grad_buf) diff --git a/src/algorithms/paramspacesgd/repgradelbo.jl b/src/algorithms/paramspacesgd/repgradelbo.jl index cb9dd3982..dafbd183f 100644 --- a/src/algorithms/paramspacesgd/repgradelbo.jl +++ b/src/algorithms/paramspacesgd/repgradelbo.jl @@ -23,6 +23,11 @@ struct RepGradELBO{EntropyEst<:AbstractEntropyEstimator} <: AbstractVariationalO n_samples::Int end +struct RepGradELBOState{Problem, ObjADPrep} + problem::Problem + obj_ad_prep::ObjADPrep +end + function init( rng::Random.AbstractRNG, obj::RepGradELBO, @@ -56,7 +61,7 @@ function init( obj_ad_prep = AdvancedVI._prepare_gradient( estimate_repgradelbo_ad_forward, adtype, params, aux ) - return (obj_ad_prep=obj_ad_prep, problem=ad_prob) + return RepGradELBOState(ad_prob, obj_ad_prep) end function RepGradELBO(n_samples::Int; entropy::AbstractEntropyEstimator=ClosedFormEntropy()) @@ -143,9 +148,9 @@ function estimate_gradient!( obj::RepGradELBO, adtype::ADTypes.AbstractADType, out::DiffResults.MutableDiffResult, + state::RepGradELBOState, params, restructure, - state, args..., ) (; obj_ad_prep, problem) = state @@ -162,6 +167,6 @@ function estimate_gradient!( estimate_repgradelbo_ad_forward, out, obj_ad_prep, adtype, params, aux ) nelbo = DiffResults.value(out) - stat = (elbo=(-nelbo),) - return out, state, stat + info = (elbo=(-nelbo),) + return out, state, info end diff --git a/src/algorithms/paramspacesgd/scoregradelbo.jl b/src/algorithms/paramspacesgd/scoregradelbo.jl index 3ce787ba3..d50ad10fe 100644 --- a/src/algorithms/paramspacesgd/scoregradelbo.jl +++ b/src/algorithms/paramspacesgd/scoregradelbo.jl @@ -16,6 +16,11 @@ struct ScoreGradELBO <: AbstractVariationalObjective n_samples::Int end +struct ScoreGradELBOState{Problem, ObjADPrep} + problem::Problem + obj_ad_prep::ObjADPrep +end + function init( rng::Random.AbstractRNG, obj::ScoreGradELBO, @@ -31,7 +36,7 @@ function init( obj_ad_prep = AdvancedVI._prepare_gradient( estimate_scoregradelbo_ad_forward, adtype, params, aux ) - return (obj_ad_prep=obj_ad_prep, problem=prob) + return ScoreGradELBOState(prob, obj_ad_prep) end function Base.show(io::IO, obj::ScoreGradELBO) @@ -83,9 +88,9 @@ function AdvancedVI.estimate_gradient!( obj::ScoreGradELBO, adtype::ADTypes.AbstractADType, out::DiffResults.MutableDiffResult, + state::ScoreGradELBOState, params, restructure, - state, ) q = restructure(params) (; obj_ad_prep, problem) = state From fc89b707debc2a05f32567b32f5ce35fda18f572 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sat, 23 Aug 2025 16:15:21 -0400 Subject: [PATCH 2/4] add history --- HISTORY.md | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 HISTORY.md diff --git a/HISTORY.md b/HISTORY.md new file mode 100644 index 000000000..7fa0f8b2d --- /dev/null +++ b/HISTORY.md @@ -0,0 +1,14 @@ + +# Release 0.5 + +## Interface Changes +An additional layer of indirection, `AbstractAlgorithms` has been added. +Previously, all variational inference algorithms were assumed to run SGD in parameter space. +This desing however, is proving to be too rigid. +Instead, each algorithm is now assumed to implement three simple interfaces: `init`, `step`, and `output`. +Algorithms that run SGD in parameter space now need to implement the `AbstractVarationalObjective` interface of `ParamSpaceSGD <: AbstractAlgorithms`, which is a general implementation of the new interface. +Therefore, the old behavior of `AdvancedVI` is fully inhereted by `ParamSpaceSGD`. + +## Internal Changes +The state of the objectives now use a concrete type. +Related to this, the objective `state` argument in `estimate_gradient!` has been moved to the front to avoid type ambiguities. From a94a1cc8dadd0580221334660bc2abb95c9aca68 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sat, 23 Aug 2025 16:16:11 -0400 Subject: [PATCH 3/4] run formatter --- src/algorithms/paramspacesgd/repgradelbo.jl | 2 +- src/algorithms/paramspacesgd/scoregradelbo.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/algorithms/paramspacesgd/repgradelbo.jl b/src/algorithms/paramspacesgd/repgradelbo.jl index dafbd183f..a9beeff56 100644 --- a/src/algorithms/paramspacesgd/repgradelbo.jl +++ b/src/algorithms/paramspacesgd/repgradelbo.jl @@ -23,7 +23,7 @@ struct RepGradELBO{EntropyEst<:AbstractEntropyEstimator} <: AbstractVariationalO n_samples::Int end -struct RepGradELBOState{Problem, ObjADPrep} +struct RepGradELBOState{Problem,ObjADPrep} problem::Problem obj_ad_prep::ObjADPrep end diff --git a/src/algorithms/paramspacesgd/scoregradelbo.jl b/src/algorithms/paramspacesgd/scoregradelbo.jl index d50ad10fe..55d81d7b3 100644 --- a/src/algorithms/paramspacesgd/scoregradelbo.jl +++ b/src/algorithms/paramspacesgd/scoregradelbo.jl @@ -16,7 +16,7 @@ struct ScoreGradELBO <: AbstractVariationalObjective n_samples::Int end -struct ScoreGradELBOState{Problem, ObjADPrep} +struct ScoreGradELBOState{Problem,ObjADPrep} problem::Problem obj_ad_prep::ObjADPrep end From dfbd0d5d3dd41e5c312657325b82d4f142ceac80 Mon Sep 17 00:00:00 2001 From: Kyurae Kim Date: Sat, 23 Aug 2025 16:18:54 -0400 Subject: [PATCH 4/4] fix run formatter --- HISTORY.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 7fa0f8b2d..682ea5adc 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,8 +1,8 @@ - # Release 0.5 ## Interface Changes -An additional layer of indirection, `AbstractAlgorithms` has been added. + +An additional layer of indirection, `AbstractAlgorithms` has been added. Previously, all variational inference algorithms were assumed to run SGD in parameter space. This desing however, is proving to be too rigid. Instead, each algorithm is now assumed to implement three simple interfaces: `init`, `step`, and `output`. @@ -10,5 +10,6 @@ Algorithms that run SGD in parameter space now need to implement the `AbstractVa Therefore, the old behavior of `AdvancedVI` is fully inhereted by `ParamSpaceSGD`. ## Internal Changes + The state of the objectives now use a concrete type. Related to this, the objective `state` argument in `estimate_gradient!` has been moved to the front to avoid type ambiguities.