From da81f913af7c8212275432407d0ca774ab3c04ab Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Fri, 5 Jun 2026 16:57:59 +0100 Subject: [PATCH 1/7] =?UTF-8?q?Add=20SubsampledLogDensity;=20rename=20subs?= =?UTF-8?q?ample=20=E2=86=92=20with=5Fbatch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the bespoke `DynamicPPLModelLogDensityFunction` ext wrapper with a generic `SubsampledLogDensity` in main, plus a small DPPL-glue method in the extension. Closes the documentation gap from #259. - `SubsampledLogDensity(prob, make_prob, dataset_size)` wraps any `LogDensityProblem` and rebuilds the inner problem per batch via a user-supplied `(batch, scale) -> prob` callback; type-stable, validates `dataset_size > 0` and `length(batch) <= dataset_size`. - `WeightedLogJoint(scale)` carries the SG correction; the DPPL extension defines its call method on `DynamicPPL.AbstractVarInfo` using the `LogPrior` / `LogLikelihood` / `LogJacobian` accumulators. - Renames the subsampling protocol verb `subsample` → `with_batch` across algorithms, the test model, and the LogReg tutorial. `with_batch` is pure (returns a fresh wrapper) and reads correctly: it applies a batch rather than drawing one. - DPPL users now write a model factory parametric in `N` and supply the data per batch via `|` (conditioning); the old `datapoints=` keyword convention is no longer required. Trade-offs (intentional): - The inner `DynamicPPL.LogDensityFunction` is rebuilt — and its AD prep re-run — on every `with_batch` call, because upstream LDF bakes the model into its prep context. Negligible for ForwardDiff; per-step cost for compiled backends (Mooncake, Enzyme). - DPPL targets now expose `LogDensityOrder{1}` only; the previous wrapper also offered `LogDensityOrder{2}` via a separate Hessian prep. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/src/tutorials/subsampling.md | 57 ++++- ext/AdvancedVIDynamicPPLExt.jl | 202 +----------------- src/AdvancedVI.jl | 9 +- src/algorithms/fisherminbatchmatch.jl | 2 +- src/algorithms/klminnaturalgraddescent.jl | 4 +- src/algorithms/klminsqrtnaturalgraddescent.jl | 4 +- src/algorithms/klminwassfwdbwd.jl | 4 +- src/algorithms/subsampledobjective.jl | 12 +- src/subsampled_logdensity.jl | 69 ++++++ test/integration/dynamicppl.jl | 29 +-- test/models/subsamplednormals.jl | 2 +- 11 files changed, 159 insertions(+), 235 deletions(-) create mode 100644 src/subsampled_logdensity.jl diff --git a/docs/src/tutorials/subsampling.md b/docs/src/tutorials/subsampling.md index 294781e63..71290d545 100644 --- a/docs/src/tutorials/subsampling.md +++ b/docs/src/tutorials/subsampling.md @@ -89,14 +89,14 @@ prob = LogReg(X, y, size(X, 1)) nothing ``` -To enable subsampling, `LogReg` has to implement the method `AdvancedVI.subsample`. +To enable subsampling, `LogReg` has to implement the method `AdvancedVI.with_batch`. For our model, this is fairly simple: We only need to select the rows of `X` and the elements of `y` corresponding to the batch of data points. ```@example subsampling using Accessors using AdvancedVI -function AdvancedVI.subsample(prob::LogReg, idx) +function AdvancedVI.with_batch(prob::LogReg, idx) prob′ = @set prob.X = prob.X[idx, :] return @set prob′.y = prob′.y[idx] end @@ -105,9 +105,9 @@ nothing !!! info - The default implementation of `AdvancedVI.subsample` is `AdvancedVI.subsample(model, idx) = model`. - Therefore, if the specialization of `AdvancedVI.subsample` is not set up properly, `AdvancedVI` will silently use full-batch gradients instead of subsampling. - It is thus useful to check whether the right specialization of `AdvancedVI.subsample` is being called. + The default implementation of `AdvancedVI.with_batch` is `AdvancedVI.with_batch(model, idx) = model`. + Therefore, if the specialization of `AdvancedVI.with_batch` is not set up properly, `AdvancedVI` will silently use full-batch gradients instead of subsampling. + It is thus useful to check whether the right specialization of `AdvancedVI.with_batch` is being called. ## Scalable Inference via AdvancedVI @@ -283,3 +283,50 @@ nothing But remember that subsampling will always be *asymptotically* slower than no subsampling. That is, as the number of iterations increase, there will be a point where no subsampling will overtake subsampling even in terms of wallclock time. Therefore, subsampling is most beneficial when a crude solution to the VI problem suffices. + +## Subsampling with `DynamicPPL` Models + +For `DynamicPPL` models, write the model as a **factory parametric in the +minibatch size `N`**, leaving observations free so they can be supplied via +`|` (conditioning) on each batch. The package extension defines a call method +on [`AdvancedVI.WeightedLogJoint`](@ref) that wires `scale * loglikelihood + +logprior - logjacobian` through `DynamicPPL`'s accumulators; wrap the +resulting LDF factory in [`SubsampledLogDensity`](@ref). + +```julia +using AdvancedVI, ADTypes, DynamicPPL, Distributions, LinearAlgebra, LogDensityProblems + +DynamicPPL.@model function bayes_logreg(X_batch, N) + d = size(X_batch, 2) + β ~ MvNormal(zeros(d), I) + y ~ arraydist([BernoulliLogit(dot(X_batch[i, :], β)) for i in 1:N]) +end + +# `X`, `y_obs` are the full dataset; `n_data = size(X, 1)`. +n_data, d = size(X, 1), size(X, 2) + +# Full-data model used only for varinfo / dim discovery. +model = bayes_logreg(X, n_data) | (y = y_obs,) +vi = DynamicPPL.link!!(DynamicPPL.VarInfo(model), model) + +batchsize = 32 +subsampling = ReshufflingBatchSubsampling(1:n_data, batchsize) +minibatch_model = batch -> bayes_logreg(X[batch, :], length(batch)) | (y = y_obs[batch],) + +make_prob = (batch, scale) -> DynamicPPL.LogDensityFunction( + minibatch_model(batch), AdvancedVI.WeightedLogJoint(scale), vi; adtype=AutoForwardDiff() +) +prob = SubsampledLogDensity(make_prob(1:n_data, 1.0), make_prob, n_data) + +alg = KLMinRepGradProxDescent(AutoForwardDiff(); subsampling) +dim = LogDensityProblems.dimension(prob) +q0 = FullRankGaussian(zeros(dim), LowerTriangular(Matrix{Float64}(0.6 * I, dim, dim))) +q, _, _ = AdvancedVI.optimize(alg, 1000, prob, q0; show_progress=false) +``` + +The variational parameter shape must be **invariant across batches**. Above, +`β` is a `d`-vector regardless of `N`; `N` only controls how many likelihood +terms are summed. Models with per-datapoint latent variables (e.g. `users[:, j]` +in a matrix factorisation) require a variational family whose dimension +changes with `batch`, which is not yet supported by the parameter-space SGD +algorithms in `AdvancedVI`. diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl index 687a36bb0..056661746 100644 --- a/ext/AdvancedVIDynamicPPLExt.jl +++ b/ext/AdvancedVIDynamicPPLExt.jl @@ -1,211 +1,13 @@ module AdvancedVIDynamicPPLExt -using ADTypes: ADTypes using AdvancedVI: AdvancedVI -using AbstractPPL: AbstractPPL using DynamicPPL: DynamicPPL -using LogDensityProblems: LogDensityProblems -using Random -adtype_capabilities(::Type{Nothing}) = LogDensityProblems.LogDensityOrder{0}() - -function adtype_capabilities(::Type{<:ADTypes.AbstractADType}) - return LogDensityProblems.LogDensityOrder{1}() -end - -# `getlogdensity` callable for `DynamicPPL.logdensity_internal`: reads the -# current `loglikeadj` through a Ref so the mutation done by `subsample` is -# observed without rebuilding any AD prep. -struct WeightedLogJoint{R<:Base.RefValue{<:Real}} - loglikeadj_ref::R -end -function (g::WeightedLogJoint)(vi) +function (g::AdvancedVI.WeightedLogJoint)(vi::DynamicPPL.AbstractVarInfo) loglike = DynamicPPL.getloglikelihood(vi) logprior = DynamicPPL.getlogprior(vi) logjac = DynamicPPL.getlogjac(vi) - return g.loglikeadj_ref[] * loglike + logprior - logjac -end - -const _DEFAULT_LDF_ACCS = DynamicPPL.AccumulatorTuple(( - DynamicPPL.LogPriorAccumulator(), - DynamicPPL.LogJacobianAccumulator(), - DynamicPPL.LogLikelihoodAccumulator(), -)) - -function subsample_dynamicpplmodel( - model::DynamicPPL.Model{F,A,D,M,Ta,Td,Ctx,Threaded}, batch -) where {F,A,D,M,Ta,Td,Ctx,Threaded} - new_kwargs = merge(model.defaults, (; datapoints=batch)) - return DynamicPPL.Model{Threaded}(model.f, model.args, new_kwargs, model.context) -end - -# `model` is the original (unsubsampled) source of truth; `subsample` must read -# it (not `model_ref[]`) to get the full-dataset length on every call. -# `model_ref`/`loglikeadj_ref` are mutated in place by `subsample` so the -# closure inside `prep_grad`/`prep_hess` stays valid across subsampling steps. -# `model_ref` is `Ref{Any}` because `subsample_dynamicpplmodel`'s output type -# varies with the batch (a typed Ref would throw on reassignment), and because -# compiled-tape backends would otherwise bake the deref into the tape and miss -# the `subsample` update. -struct DynamicPPLModelLogDensityFunction{ - Model<:DynamicPPL.Model, - LogLikeAdj<:Real, - Ranges<:DynamicPPL.VarNamedTuple, - Strategy<:DynamicPPL.AbstractTransformStrategy, - GetLogDensity, - ADType<:Union{Nothing,ADTypes.AbstractADType}, - PrepGrad, - PrepHess, -} - model::Model - model_ref::Ref{Any} - loglikeadj_ref::Ref{LogLikeAdj} - ranges_and_transforms::Ranges - transform_strategy::Strategy - getlogdensity::GetLogDensity - adtype::ADType - prep_grad::PrepGrad - prep_hess::PrepHess - dim::Int -end - -function DynamicPPLModelLogDensityFunction( - model::DynamicPPL.Model, - varinfo::DynamicPPL.AbstractVarInfo; - use_hessian::Bool=true, - adtype::Union{Nothing,ADTypes.AbstractADType}=nothing, - loglikeadj::Real=1.0, - subsampling::Union{Nothing,AdvancedVI.AbstractSubsampling}=nothing, -) - model_sub = if isnothing(subsampling) - model - else - rng = Random.default_rng() - sub_st = AdvancedVI.init(rng, subsampling) - batch, _, _ = AdvancedVI.step(rng, subsampling, sub_st) - subsample_dynamicpplmodel(model, batch) - end - - ranges_and_transforms, params = DynamicPPL.get_rat_and_samplevec(varinfo.values) - transform_strategy = DynamicPPL.infer_transform_strategy_from_values( - ranges_and_transforms - ) - - model_ref = Ref{Any}(model_sub) - loglikeadj_ref = Ref(float(loglikeadj)) - getlogdensity = WeightedLogJoint(loglikeadj_ref) - f = - params -> DynamicPPL.logdensity_internal( - params, - model_ref[], - getlogdensity, - ranges_and_transforms, - transform_strategy, - _DEFAULT_LDF_ACCS, - ) - cap = adtype_capabilities(typeof(adtype)) - - prep_grad = if cap >= LogDensityProblems.LogDensityOrder{1}() - AbstractPPL.prepare(adtype, f, params) - else - nothing - end - prep_hess = if cap >= LogDensityProblems.LogDensityOrder{1}() && use_hessian - try - AbstractPPL.prepare(adtype, f, params; order=2) - catch err - err isa MethodError || rethrow() - @warn "The selected AD backend does not support `AbstractPPL.prepare(...; order=2)`. AdvancedVI will treat the model as first-order only." - nothing - end - else - nothing - end - return DynamicPPLModelLogDensityFunction{ - typeof(model), - eltype(loglikeadj_ref), - typeof(ranges_and_transforms), - typeof(transform_strategy), - typeof(getlogdensity), - typeof(adtype), - typeof(prep_grad), - typeof(prep_hess), - }( - model, - model_ref, - loglikeadj_ref, - ranges_and_transforms, - transform_strategy, - getlogdensity, - adtype, - prep_grad, - prep_hess, - length(params), - ) -end - -function LogDensityProblems.logdensity(prob::DynamicPPLModelLogDensityFunction, params) - return DynamicPPL.logdensity_internal( - params, - prob.model_ref[], - prob.getlogdensity, - prob.ranges_and_transforms, - prob.transform_strategy, - _DEFAULT_LDF_ACCS, - ) -end - -# `!!` may alias internal buffers of `prep_*`; copy so callers can retain the -# arrays past the next AD call. -function LogDensityProblems.logdensity_and_gradient( - prob::DynamicPPLModelLogDensityFunction, params -) - val, grad = AbstractPPL.value_and_gradient!!(prob.prep_grad, params) - return val, copy(grad) -end - -function LogDensityProblems.logdensity_gradient_and_hessian( - prob::DynamicPPLModelLogDensityFunction, params -) - val, grad, H = AbstractPPL.value_gradient_and_hessian!!(prob.prep_hess, params) - return val, copy(grad), copy(H) -end - -function LogDensityProblems.capabilities( - ::Type{<:DynamicPPLModelLogDensityFunction{Model,L,R,S,G,ADType,PG,PH}} -) where {Model,L,R,S,G,ADType<:Union{Nothing,ADTypes.AbstractADType},PG,PH} - return if PH !== Nothing - LogDensityProblems.LogDensityOrder{2}() - elseif PG !== Nothing - LogDensityProblems.LogDensityOrder{1}() - else - LogDensityProblems.LogDensityOrder{0}() - end -end - -LogDensityProblems.dimension(prob::DynamicPPLModelLogDensityFunction) = prob.dim - -function AdvancedVI.subsample(prob::DynamicPPLModelLogDensityFunction, batch) - model = prob.model # full dataset — `model_ref[]` would already be subsampled - - if !haskey(model.defaults, :datapoints) - throw( - ArgumentError( - "Subsampling is turned on, but the model does not have a `datapoints` keyword argument.", - ), - ) - end - - n_datapoints = length(model.defaults.datapoints) - batchsize = length(batch) - model_sub = subsample_dynamicpplmodel(model, batch) - T = eltype(prob.loglikeadj_ref) - loglikeadj = T(n_datapoints) / T(batchsize) - - prob.model_ref[] = model_sub - prob.loglikeadj_ref[] = loglikeadj - - return prob + return g.scale * loglike + logprior - logjac end end diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 727fb2f5d..8bce4d422 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -301,8 +301,8 @@ export estimate_objective # Subsampling """ - subsample(model, batch) - subsample(q, batch) + with_batch(model, batch) + with_batch(q, batch) Inform `model` or `q` to only use the data points designated by the iterable collection `batch`. For `model`, the log-density should also be adjusted to account for the change in number of data points. @@ -310,13 +310,14 @@ For `model`, the log-density should also be adjusted to account for the change i Implementations may mutate `model`/`q` in place and return the same object, so callers must not assume the returned value is distinct from the input. """ -subsample(model_or_q::Any, ::Any) = model_or_q +with_batch(model_or_q::Any, ::Any) = model_or_q abstract type AbstractSubsampling end include("reshuffling.jl") +include("subsampled_logdensity.jl") -export ReshufflingBatchSubsampling +export ReshufflingBatchSubsampling, SubsampledLogDensity # Main optimization routine function optimize end diff --git a/src/algorithms/fisherminbatchmatch.jl b/src/algorithms/fisherminbatchmatch.jl index f3ea6fd8d..f0c1a6193 100644 --- a/src/algorithms/fisherminbatchmatch.jl +++ b/src/algorithms/fisherminbatchmatch.jl @@ -132,7 +132,7 @@ function step( prob, sub_st, NamedTuple() else batch, sub_st′, sub_inf = step(rng, subsampling, sub_st) - prob_sub = subsample(prob, batch) + prob_sub = with_batch(prob, batch) prob_sub, sub_st′, sub_inf end diff --git a/src/algorithms/klminnaturalgraddescent.jl b/src/algorithms/klminnaturalgraddescent.jl index 596cd372c..69ed852cc 100644 --- a/src/algorithms/klminnaturalgraddescent.jl +++ b/src/algorithms/klminnaturalgraddescent.jl @@ -113,7 +113,7 @@ function step( prob, sub_st, NamedTuple() else batch, sub_st′, sub_inf = step(rng, subsampling, sub_st) - prob_sub = subsample(prob, batch) + prob_sub = with_batch(prob, batch) prob_sub, sub_st′, sub_inf end @@ -184,7 +184,7 @@ function estimate_objective( sub_st = init(rng, sub) return mapreduce(+, 1:length(sub)) do _ batch, sub_st, _ = step(rng, sub, sub_st) - prob_sub = subsample(prob, batch) + prob_sub = with_batch(prob, batch) estimate_objective(rng, obj, q, prob_sub) / length(sub) end end diff --git a/src/algorithms/klminsqrtnaturalgraddescent.jl b/src/algorithms/klminsqrtnaturalgraddescent.jl index 3b29c3b2b..91d622f57 100644 --- a/src/algorithms/klminsqrtnaturalgraddescent.jl +++ b/src/algorithms/klminsqrtnaturalgraddescent.jl @@ -97,7 +97,7 @@ function step( prob, sub_st, NamedTuple() else batch, sub_st′, sub_inf = step(rng, subsampling, sub_st) - prob_sub = subsample(prob, batch) + prob_sub = with_batch(prob, batch) prob_sub, sub_st′, sub_inf end @@ -158,7 +158,7 @@ function estimate_objective( sub_st = init(rng, sub) return mapreduce(+, 1:length(sub)) do _ batch, sub_st, _ = step(rng, sub, sub_st) - prob_sub = subsample(prob, batch) + prob_sub = with_batch(prob, batch) estimate_objective(rng, obj, q, prob_sub) / length(sub) end end diff --git a/src/algorithms/klminwassfwdbwd.jl b/src/algorithms/klminwassfwdbwd.jl index 40577bfc4..8b1e1a00b 100644 --- a/src/algorithms/klminwassfwdbwd.jl +++ b/src/algorithms/klminwassfwdbwd.jl @@ -93,7 +93,7 @@ function step( prob, sub_st, NamedTuple() else batch, sub_st′, sub_inf = step(rng, subsampling, sub_st) - prob_sub = subsample(prob, batch) + prob_sub = with_batch(prob, batch) prob_sub, sub_st′, sub_inf end @@ -153,7 +153,7 @@ function estimate_objective( sub_st = init(rng, sub) return mapreduce(+, 1:length(sub)) do _ batch, sub_st, _ = step(rng, sub, sub_st) - prob_sub = subsample(prob, batch) + prob_sub = with_batch(prob, batch) estimate_objective(rng, obj, q, prob_sub) / length(sub) end end diff --git a/src/algorithms/subsampledobjective.jl b/src/algorithms/subsampledobjective.jl index ad1f858ef..1b782acca 100644 --- a/src/algorithms/subsampledobjective.jl +++ b/src/algorithms/subsampledobjective.jl @@ -34,8 +34,8 @@ function init( # This is necessary to ensure that `init` sees the type "conditioned" on a minibatch # so that any prepared AD evaluator inside it sees the correct batch-subsampled type. batch, _, _ = step(rng, subsampling, sub_st, true) - prob_sub = subsample(prob, batch) - q_init_sub = subsample(q_init, batch) + prob_sub = with_batch(prob, batch) + q_init_sub = with_batch(q_init, batch) params_sub, re_sub = Optimisers.destructure(q_init_sub) obj_st = AdvancedVI.init( @@ -51,8 +51,8 @@ function estimate_objective( sub_st = init(rng, subsampling) return mapreduce(+, 1:length(subsampling)) do _ batch, sub_st, _ = step(rng, subsampling, sub_st) - prob_sub = subsample(prob, batch) - q_sub = subsample(q, batch) + prob_sub = with_batch(prob, batch) + q_sub = with_batch(q, batch) estimate_objective(rng, objective, q_sub, prob_sub; kwargs...) / length(subsampling) end end @@ -77,8 +77,8 @@ function estimate_gradient!( q = restructure(params) batch, sub_st′, sub_inf = step(rng, subsampling, sub_st, true) - prob_sub = subsample(prob, batch) - q_sub = subsample(q, batch) + prob_sub = with_batch(prob, batch) + q_sub = with_batch(q, batch) params_sub, re_sub = Optimisers.destructure(q_sub) obj_st′ = set_objective_state_problem(obj_st, prob_sub) diff --git a/src/subsampled_logdensity.jl b/src/subsampled_logdensity.jl new file mode 100644 index 000000000..713355fad --- /dev/null +++ b/src/subsampled_logdensity.jl @@ -0,0 +1,69 @@ +""" + SubsampledLogDensity(prob, make_prob, dataset_size) + +`LogDensityProblems`-compatible wrapper that supports [`with_batch`](@ref): +`with_batch(prob, batch)` returns a fresh wrapper whose inner problem is +`make_prob(batch, dataset_size / length(batch))`. `make_prob` must return +objects of the same concrete type as the initial `prob` so the wrapper stays +type-stable. The inner problem's capabilities and dimension are surfaced. + +`dataset_size` must equal the size of the dataset that `batch` indexes into: +the rescaling `dataset_size / length(batch)` is only an unbiased estimator +when these are consistent. `with_batch` checks `length(batch) <= dataset_size` +to catch the obvious misuse; a `batch` drawn from a different dataset will +silently scale the gradient by the wrong factor. +""" +struct SubsampledLogDensity{P,F} + prob::P + make_prob::F + dataset_size::Int + function SubsampledLogDensity{P,F}(prob::P, make_prob::F, dataset_size::Int) where {P,F} + # Caught here to prevent silent zero-gradient (or sign flip) downstream. + dataset_size > 0 || + throw(ArgumentError("`dataset_size` must be positive, got $dataset_size.")) + return new{P,F}(prob, make_prob, dataset_size) + end +end +function SubsampledLogDensity(prob, make_prob, dataset_size::Integer) + return SubsampledLogDensity{typeof(prob),typeof(make_prob)}( + prob, make_prob, Int(dataset_size) + ) +end + +function LogDensityProblems.logdensity(prob::SubsampledLogDensity, x) + return LogDensityProblems.logdensity(prob.prob, x) +end + +function LogDensityProblems.logdensity_and_gradient(prob::SubsampledLogDensity, x) + return LogDensityProblems.logdensity_and_gradient(prob.prob, x) +end + +function LogDensityProblems.dimension(prob::SubsampledLogDensity) + return LogDensityProblems.dimension(prob.prob) +end + +function LogDensityProblems.capabilities(::Type{<:SubsampledLogDensity{P}}) where {P} + return LogDensityProblems.capabilities(P) +end + +function with_batch(prob::SubsampledLogDensity{P,F}, batch) where {P,F} + length(batch) <= prob.dataset_size || throw( + ArgumentError( + "`length(batch) = $(length(batch))` exceeds `dataset_size = $(prob.dataset_size)`; " * + "the batch must come from the same dataset that `dataset_size` describes.", + ), + ) + new_inner = prob.make_prob(batch, prob.dataset_size / length(batch)) + return SubsampledLogDensity{P,F}(new_inner, prob.make_prob, prob.dataset_size) +end + +""" + WeightedLogJoint(scale) + +Callable returning `scale * loglikelihood + logprior - logjacobian` of a +varinfo. The call method is backend-specific; package extensions register +overloads for the varinfo types they support. +""" +struct WeightedLogJoint{T<:Real} + scale::T +end diff --git a/test/integration/dynamicppl.jl b/test/integration/dynamicppl.jl index a3ed01d20..8b30bb0fb 100644 --- a/test/integration/dynamicppl.jl +++ b/test/integration/dynamicppl.jl @@ -4,9 +4,9 @@ return x ~ MvNormal(μ, I) end - DynamicPPL.@model function normal_subsampled(μs; datapoints=1:size(μs, 2)) - for i in datapoints - x ~ MvNormal(μs[:, i], I) + DynamicPPL.@model function normal_minibatch(μs_batch, N) + for i in 1:N + x ~ MvNormal(μs_batch[:, i], I) end end @@ -17,8 +17,9 @@ vi = DynamicPPL.VarInfo(model) vi = DynamicPPL.link!!(vi, model) - ext = Base.get_extension(AdvancedVI, :AdvancedVIDynamicPPLExt) - prob = ext.DynamicPPLModelLogDensityFunction(model, vi; adtype=AD) + prob = DynamicPPL.LogDensityFunction( + model, DynamicPPL.getlogjoint_internal, vi; adtype=AD + ) alg = KLMinRepGradProxDescent(AD) d = LogDensityProblems.dimension(prob) @@ -35,16 +36,17 @@ μs = 3 * randn(2, n_data) μ_true = mean(μs; dims=2)[:, 1] - model = normal_subsampled(μs) - vi = DynamicPPL.VarInfo(model) - vi = DynamicPPL.link!!(vi, model) + model = normal_minibatch(μs, n_data) + vi = DynamicPPL.link!!(DynamicPPL.VarInfo(model), model) - dataset = 1:n_data batchsize = 2 - subsampling = ReshufflingBatchSubsampling(dataset, batchsize) + subsampling = ReshufflingBatchSubsampling(1:n_data, batchsize) + minibatch_model = batch -> normal_minibatch(μs[:, batch], length(batch)) - ext = Base.get_extension(AdvancedVI, :AdvancedVIDynamicPPLExt) - prob = ext.DynamicPPLModelLogDensityFunction(model, vi; adtype=AD, subsampling) + make_prob = (batch, scale) -> DynamicPPL.LogDensityFunction( + minibatch_model(batch), AdvancedVI.WeightedLogJoint(scale), vi; adtype=AD + ) + prob = SubsampledLogDensity(make_prob(1:n_data, 1.0), make_prob, n_data) alg = KLMinRepGradProxDescent(AD; subsampling) d = LogDensityProblems.dimension(prob) @@ -54,5 +56,8 @@ Δλ0 = sum(abs2, q0.location - μ_true) Δλ = sum(abs2, q.location - μ_true) @test Δλ ≤ Δλ0 / 2 + + @test_throws ArgumentError SubsampledLogDensity(prob.prob, make_prob, 0) + @test_throws ArgumentError AdvancedVI.with_batch(prob, 1:(n_data + 1)) end end diff --git a/test/models/subsamplednormals.jl b/test/models/subsamplednormals.jl index 8b7a8971a..b2c1774b6 100644 --- a/test/models/subsamplednormals.jl +++ b/test/models/subsamplednormals.jl @@ -42,7 +42,7 @@ function LogDensityProblems.capabilities(::Type{SubsampledNormals{D,F,C}}) where return C() end -function AdvancedVI.subsample(m::SubsampledNormals, idx) +function AdvancedVI.with_batch(m::SubsampledNormals, idx) n_data = length(m.dists) return SubsampledNormals(m.dists[idx], n_data / length(idx), m.capability) end From cd0744b483be174eede281ab101d684c0673084f Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Fri, 5 Jun 2026 17:05:44 +0100 Subject: [PATCH 2/7] Format with JuliaFormatter v1 (blue style) Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/src/tutorials/subsampling.md | 29 ++++++++++++++++------------- test/integration/dynamicppl.jl | 10 +++++++--- 2 files changed, 23 insertions(+), 16 deletions(-) diff --git a/docs/src/tutorials/subsampling.md b/docs/src/tutorials/subsampling.md index 71290d545..bbc117d2a 100644 --- a/docs/src/tutorials/subsampling.md +++ b/docs/src/tutorials/subsampling.md @@ -289,8 +289,7 @@ Therefore, subsampling is most beneficial when a crude solution to the VI proble For `DynamicPPL` models, write the model as a **factory parametric in the minibatch size `N`**, leaving observations free so they can be supplied via `|` (conditioning) on each batch. The package extension defines a call method -on [`AdvancedVI.WeightedLogJoint`](@ref) that wires `scale * loglikelihood + -logprior - logjacobian` through `DynamicPPL`'s accumulators; wrap the +on [`AdvancedVI.WeightedLogJoint`](@ref) that wires `scale * loglikelihood + logprior - logjacobian` through `DynamicPPL`'s accumulators; wrap the resulting LDF factory in [`SubsampledLogDensity`](@ref). ```julia @@ -299,28 +298,32 @@ using AdvancedVI, ADTypes, DynamicPPL, Distributions, LinearAlgebra, LogDensityP DynamicPPL.@model function bayes_logreg(X_batch, N) d = size(X_batch, 2) β ~ MvNormal(zeros(d), I) - y ~ arraydist([BernoulliLogit(dot(X_batch[i, :], β)) for i in 1:N]) + return y ~ arraydist([BernoulliLogit(dot(X_batch[i, :], β)) for i in 1:N]) end # `X`, `y_obs` are the full dataset; `n_data = size(X, 1)`. n_data, d = size(X, 1), size(X, 2) # Full-data model used only for varinfo / dim discovery. -model = bayes_logreg(X, n_data) | (y = y_obs,) -vi = DynamicPPL.link!!(DynamicPPL.VarInfo(model), model) +model = bayes_logreg(X, n_data) | (y=y_obs,) +vi = DynamicPPL.link!!(DynamicPPL.VarInfo(model), model) -batchsize = 32 -subsampling = ReshufflingBatchSubsampling(1:n_data, batchsize) -minibatch_model = batch -> bayes_logreg(X[batch, :], length(batch)) | (y = y_obs[batch],) - -make_prob = (batch, scale) -> DynamicPPL.LogDensityFunction( - minibatch_model(batch), AdvancedVI.WeightedLogJoint(scale), vi; adtype=AutoForwardDiff() -) +batchsize = 32 +subsampling = ReshufflingBatchSubsampling(1:n_data, batchsize) +minibatch_model = batch -> bayes_logreg(X[batch, :], length(batch)) | (y=y_obs[batch],) + +make_prob = + (batch, scale) -> DynamicPPL.LogDensityFunction( + minibatch_model(batch), + AdvancedVI.WeightedLogJoint(scale), + vi; + adtype=AutoForwardDiff(), + ) prob = SubsampledLogDensity(make_prob(1:n_data, 1.0), make_prob, n_data) alg = KLMinRepGradProxDescent(AutoForwardDiff(); subsampling) dim = LogDensityProblems.dimension(prob) -q0 = FullRankGaussian(zeros(dim), LowerTriangular(Matrix{Float64}(0.6 * I, dim, dim))) +q0 = FullRankGaussian(zeros(dim), LowerTriangular(Matrix{Float64}(0.6 * I, dim, dim))) q, _, _ = AdvancedVI.optimize(alg, 1000, prob, q0; show_progress=false) ``` diff --git a/test/integration/dynamicppl.jl b/test/integration/dynamicppl.jl index 8b30bb0fb..aa5de37c3 100644 --- a/test/integration/dynamicppl.jl +++ b/test/integration/dynamicppl.jl @@ -43,9 +43,13 @@ subsampling = ReshufflingBatchSubsampling(1:n_data, batchsize) minibatch_model = batch -> normal_minibatch(μs[:, batch], length(batch)) - make_prob = (batch, scale) -> DynamicPPL.LogDensityFunction( - minibatch_model(batch), AdvancedVI.WeightedLogJoint(scale), vi; adtype=AD - ) + make_prob = + (batch, scale) -> DynamicPPL.LogDensityFunction( + minibatch_model(batch), + AdvancedVI.WeightedLogJoint(scale), + vi; + adtype=AD, + ) prob = SubsampledLogDensity(make_prob(1:n_data, 1.0), make_prob, n_data) alg = KLMinRepGradProxDescent(AD; subsampling) From 38535415a323eee69bebde2486728f9496b8dfe9 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Fri, 5 Jun 2026 17:11:20 +0100 Subject: [PATCH 3/7] Fix DynamicPPL subsampling test so SG correction is exercised MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous model wrote `x ~ MvNormal(μs_batch[:, i], I)` N times for the same `x`. DPPL routes ALL contributions of a sampled variable into `LogPriorAccumulator`, so `LogLikelihood` stayed at 0.0 and the `scale * loglike + logprior - logjac` correction multiplied zero. The test passed only because the per-batch gradient direction happened to point at the right answer. Inverted the model to a proper hierarchical setup: a latent `μ` with a weak prior, observations conditioned on `μ`. Now the per-batch contributions land in `LogLikelihoodAccumulator`, where the SG-correction scale is actually applied. Co-Authored-By: Claude Opus 4.7 (1M context) --- test/integration/dynamicppl.jl | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/test/integration/dynamicppl.jl b/test/integration/dynamicppl.jl index aa5de37c3..2ae35e633 100644 --- a/test/integration/dynamicppl.jl +++ b/test/integration/dynamicppl.jl @@ -4,9 +4,14 @@ return x ~ MvNormal(μ, I) end - DynamicPPL.@model function normal_minibatch(μs_batch, N) + # `μ` is the latent parameter being inferred from observations stored in + # `obs_batch`. The data observations land in `LogLikelihoodAccumulator`, + # which is what the SG-correction scale multiplies — verifying the + # minibatch correction actually exercises the likelihood path. + DynamicPPL.@model function normal_minibatch(obs_batch, N) + μ ~ MvNormal(zeros(size(obs_batch, 1)), 100.0 * I) for i in 1:N - x ~ MvNormal(μs_batch[:, i], I) + obs_batch[:, i] ~ MvNormal(μ, I) end end @@ -33,15 +38,15 @@ @testset "subsampling" begin n_data = 32 - μs = 3 * randn(2, n_data) - μ_true = mean(μs; dims=2)[:, 1] + μ_true = [-2.0, 2.0] + observations = μ_true .+ randn(2, n_data) - model = normal_minibatch(μs, n_data) + model = normal_minibatch(observations, n_data) vi = DynamicPPL.link!!(DynamicPPL.VarInfo(model), model) batchsize = 2 subsampling = ReshufflingBatchSubsampling(1:n_data, batchsize) - minibatch_model = batch -> normal_minibatch(μs[:, batch], length(batch)) + minibatch_model = batch -> normal_minibatch(observations[:, batch], length(batch)) make_prob = (batch, scale) -> DynamicPPL.LogDensityFunction( From 2f6a1f525ce6eacfcb1295e8f8f8c0f42a689c96 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Fri, 5 Jun 2026 17:12:46 +0100 Subject: [PATCH 4/7] Tighten subsampling test target to actual MAP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `μ_true = [-2.0, 2.0]` was a fixed point that q never converges to exactly — q converges to the sample mean of `observations`, which differs by O(1/√n_data). Compute the target from the data so the assertion checks proximity to the actual fixed point. Co-Authored-By: Claude Opus 4.7 (1M context) --- test/integration/dynamicppl.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/integration/dynamicppl.jl b/test/integration/dynamicppl.jl index 2ae35e633..15c09497f 100644 --- a/test/integration/dynamicppl.jl +++ b/test/integration/dynamicppl.jl @@ -38,8 +38,10 @@ @testset "subsampling" begin n_data = 32 - μ_true = [-2.0, 2.0] - observations = μ_true .+ randn(2, n_data) + observations = [-2.0, 2.0] .+ randn(2, n_data) + # MAP target — q converges to the sample mean (weak prior is negligible + # against `n_data` likelihood contributions). + μ_true = mean(observations; dims=2)[:, 1] model = normal_minibatch(observations, n_data) vi = DynamicPPL.link!!(DynamicPPL.VarInfo(model), model) From 7581c0311ae016de623078008ede98e0cbfc6728 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Fri, 5 Jun 2026 17:19:46 +0100 Subject: [PATCH 5/7] Rename with_batch back to subsample Reverts the protocol-verb rename to minimise the diff against main: the algorithm files and `test/models/subsamplednormals.jl` are now byte-identical to main. Only the new abstraction (`SubsampledLogDensity`, `WeightedLogJoint`) and the rewired DPPL ext/test/doc remain changed. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/src/tutorials/subsampling.md | 10 +++++----- src/AdvancedVI.jl | 6 +++--- src/algorithms/fisherminbatchmatch.jl | 2 +- src/algorithms/klminnaturalgraddescent.jl | 4 ++-- src/algorithms/klminsqrtnaturalgraddescent.jl | 4 ++-- src/algorithms/klminwassfwdbwd.jl | 4 ++-- src/algorithms/subsampledobjective.jl | 12 ++++++------ src/subsampled_logdensity.jl | 8 ++++---- test/integration/dynamicppl.jl | 2 +- test/models/subsamplednormals.jl | 2 +- 10 files changed, 27 insertions(+), 27 deletions(-) diff --git a/docs/src/tutorials/subsampling.md b/docs/src/tutorials/subsampling.md index bbc117d2a..03ce43cb2 100644 --- a/docs/src/tutorials/subsampling.md +++ b/docs/src/tutorials/subsampling.md @@ -89,14 +89,14 @@ prob = LogReg(X, y, size(X, 1)) nothing ``` -To enable subsampling, `LogReg` has to implement the method `AdvancedVI.with_batch`. +To enable subsampling, `LogReg` has to implement the method `AdvancedVI.subsample`. For our model, this is fairly simple: We only need to select the rows of `X` and the elements of `y` corresponding to the batch of data points. ```@example subsampling using Accessors using AdvancedVI -function AdvancedVI.with_batch(prob::LogReg, idx) +function AdvancedVI.subsample(prob::LogReg, idx) prob′ = @set prob.X = prob.X[idx, :] return @set prob′.y = prob′.y[idx] end @@ -105,9 +105,9 @@ nothing !!! info - The default implementation of `AdvancedVI.with_batch` is `AdvancedVI.with_batch(model, idx) = model`. - Therefore, if the specialization of `AdvancedVI.with_batch` is not set up properly, `AdvancedVI` will silently use full-batch gradients instead of subsampling. - It is thus useful to check whether the right specialization of `AdvancedVI.with_batch` is being called. + The default implementation of `AdvancedVI.subsample` is `AdvancedVI.subsample(model, idx) = model`. + Therefore, if the specialization of `AdvancedVI.subsample` is not set up properly, `AdvancedVI` will silently use full-batch gradients instead of subsampling. + It is thus useful to check whether the right specialization of `AdvancedVI.subsample` is being called. ## Scalable Inference via AdvancedVI diff --git a/src/AdvancedVI.jl b/src/AdvancedVI.jl index 8bce4d422..5166f6758 100644 --- a/src/AdvancedVI.jl +++ b/src/AdvancedVI.jl @@ -301,8 +301,8 @@ export estimate_objective # Subsampling """ - with_batch(model, batch) - with_batch(q, batch) + subsample(model, batch) + subsample(q, batch) Inform `model` or `q` to only use the data points designated by the iterable collection `batch`. For `model`, the log-density should also be adjusted to account for the change in number of data points. @@ -310,7 +310,7 @@ For `model`, the log-density should also be adjusted to account for the change i Implementations may mutate `model`/`q` in place and return the same object, so callers must not assume the returned value is distinct from the input. """ -with_batch(model_or_q::Any, ::Any) = model_or_q +subsample(model_or_q::Any, ::Any) = model_or_q abstract type AbstractSubsampling end diff --git a/src/algorithms/fisherminbatchmatch.jl b/src/algorithms/fisherminbatchmatch.jl index f0c1a6193..f3ea6fd8d 100644 --- a/src/algorithms/fisherminbatchmatch.jl +++ b/src/algorithms/fisherminbatchmatch.jl @@ -132,7 +132,7 @@ function step( prob, sub_st, NamedTuple() else batch, sub_st′, sub_inf = step(rng, subsampling, sub_st) - prob_sub = with_batch(prob, batch) + prob_sub = subsample(prob, batch) prob_sub, sub_st′, sub_inf end diff --git a/src/algorithms/klminnaturalgraddescent.jl b/src/algorithms/klminnaturalgraddescent.jl index 69ed852cc..596cd372c 100644 --- a/src/algorithms/klminnaturalgraddescent.jl +++ b/src/algorithms/klminnaturalgraddescent.jl @@ -113,7 +113,7 @@ function step( prob, sub_st, NamedTuple() else batch, sub_st′, sub_inf = step(rng, subsampling, sub_st) - prob_sub = with_batch(prob, batch) + prob_sub = subsample(prob, batch) prob_sub, sub_st′, sub_inf end @@ -184,7 +184,7 @@ function estimate_objective( sub_st = init(rng, sub) return mapreduce(+, 1:length(sub)) do _ batch, sub_st, _ = step(rng, sub, sub_st) - prob_sub = with_batch(prob, batch) + prob_sub = subsample(prob, batch) estimate_objective(rng, obj, q, prob_sub) / length(sub) end end diff --git a/src/algorithms/klminsqrtnaturalgraddescent.jl b/src/algorithms/klminsqrtnaturalgraddescent.jl index 91d622f57..3b29c3b2b 100644 --- a/src/algorithms/klminsqrtnaturalgraddescent.jl +++ b/src/algorithms/klminsqrtnaturalgraddescent.jl @@ -97,7 +97,7 @@ function step( prob, sub_st, NamedTuple() else batch, sub_st′, sub_inf = step(rng, subsampling, sub_st) - prob_sub = with_batch(prob, batch) + prob_sub = subsample(prob, batch) prob_sub, sub_st′, sub_inf end @@ -158,7 +158,7 @@ function estimate_objective( sub_st = init(rng, sub) return mapreduce(+, 1:length(sub)) do _ batch, sub_st, _ = step(rng, sub, sub_st) - prob_sub = with_batch(prob, batch) + prob_sub = subsample(prob, batch) estimate_objective(rng, obj, q, prob_sub) / length(sub) end end diff --git a/src/algorithms/klminwassfwdbwd.jl b/src/algorithms/klminwassfwdbwd.jl index 8b1e1a00b..40577bfc4 100644 --- a/src/algorithms/klminwassfwdbwd.jl +++ b/src/algorithms/klminwassfwdbwd.jl @@ -93,7 +93,7 @@ function step( prob, sub_st, NamedTuple() else batch, sub_st′, sub_inf = step(rng, subsampling, sub_st) - prob_sub = with_batch(prob, batch) + prob_sub = subsample(prob, batch) prob_sub, sub_st′, sub_inf end @@ -153,7 +153,7 @@ function estimate_objective( sub_st = init(rng, sub) return mapreduce(+, 1:length(sub)) do _ batch, sub_st, _ = step(rng, sub, sub_st) - prob_sub = with_batch(prob, batch) + prob_sub = subsample(prob, batch) estimate_objective(rng, obj, q, prob_sub) / length(sub) end end diff --git a/src/algorithms/subsampledobjective.jl b/src/algorithms/subsampledobjective.jl index 1b782acca..ad1f858ef 100644 --- a/src/algorithms/subsampledobjective.jl +++ b/src/algorithms/subsampledobjective.jl @@ -34,8 +34,8 @@ function init( # This is necessary to ensure that `init` sees the type "conditioned" on a minibatch # so that any prepared AD evaluator inside it sees the correct batch-subsampled type. batch, _, _ = step(rng, subsampling, sub_st, true) - prob_sub = with_batch(prob, batch) - q_init_sub = with_batch(q_init, batch) + prob_sub = subsample(prob, batch) + q_init_sub = subsample(q_init, batch) params_sub, re_sub = Optimisers.destructure(q_init_sub) obj_st = AdvancedVI.init( @@ -51,8 +51,8 @@ function estimate_objective( sub_st = init(rng, subsampling) return mapreduce(+, 1:length(subsampling)) do _ batch, sub_st, _ = step(rng, subsampling, sub_st) - prob_sub = with_batch(prob, batch) - q_sub = with_batch(q, batch) + prob_sub = subsample(prob, batch) + q_sub = subsample(q, batch) estimate_objective(rng, objective, q_sub, prob_sub; kwargs...) / length(subsampling) end end @@ -77,8 +77,8 @@ function estimate_gradient!( q = restructure(params) batch, sub_st′, sub_inf = step(rng, subsampling, sub_st, true) - prob_sub = with_batch(prob, batch) - q_sub = with_batch(q, batch) + prob_sub = subsample(prob, batch) + q_sub = subsample(q, batch) params_sub, re_sub = Optimisers.destructure(q_sub) obj_st′ = set_objective_state_problem(obj_st, prob_sub) diff --git a/src/subsampled_logdensity.jl b/src/subsampled_logdensity.jl index 713355fad..5726cfc41 100644 --- a/src/subsampled_logdensity.jl +++ b/src/subsampled_logdensity.jl @@ -1,15 +1,15 @@ """ SubsampledLogDensity(prob, make_prob, dataset_size) -`LogDensityProblems`-compatible wrapper that supports [`with_batch`](@ref): -`with_batch(prob, batch)` returns a fresh wrapper whose inner problem is +`LogDensityProblems`-compatible wrapper that supports [`subsample`](@ref): +`subsample(prob, batch)` returns a fresh wrapper whose inner problem is `make_prob(batch, dataset_size / length(batch))`. `make_prob` must return objects of the same concrete type as the initial `prob` so the wrapper stays type-stable. The inner problem's capabilities and dimension are surfaced. `dataset_size` must equal the size of the dataset that `batch` indexes into: the rescaling `dataset_size / length(batch)` is only an unbiased estimator -when these are consistent. `with_batch` checks `length(batch) <= dataset_size` +when these are consistent. `subsample` checks `length(batch) <= dataset_size` to catch the obvious misuse; a `batch` drawn from a different dataset will silently scale the gradient by the wrong factor. """ @@ -46,7 +46,7 @@ function LogDensityProblems.capabilities(::Type{<:SubsampledLogDensity{P}}) wher return LogDensityProblems.capabilities(P) end -function with_batch(prob::SubsampledLogDensity{P,F}, batch) where {P,F} +function subsample(prob::SubsampledLogDensity{P,F}, batch) where {P,F} length(batch) <= prob.dataset_size || throw( ArgumentError( "`length(batch) = $(length(batch))` exceeds `dataset_size = $(prob.dataset_size)`; " * diff --git a/test/integration/dynamicppl.jl b/test/integration/dynamicppl.jl index 15c09497f..9eb7c6f50 100644 --- a/test/integration/dynamicppl.jl +++ b/test/integration/dynamicppl.jl @@ -69,6 +69,6 @@ @test Δλ ≤ Δλ0 / 2 @test_throws ArgumentError SubsampledLogDensity(prob.prob, make_prob, 0) - @test_throws ArgumentError AdvancedVI.with_batch(prob, 1:(n_data + 1)) + @test_throws ArgumentError AdvancedVI.subsample(prob, 1:(n_data + 1)) end end diff --git a/test/models/subsamplednormals.jl b/test/models/subsamplednormals.jl index b2c1774b6..8b7a8971a 100644 --- a/test/models/subsamplednormals.jl +++ b/test/models/subsamplednormals.jl @@ -42,7 +42,7 @@ function LogDensityProblems.capabilities(::Type{SubsampledNormals{D,F,C}}) where return C() end -function AdvancedVI.with_batch(m::SubsampledNormals, idx) +function AdvancedVI.subsample(m::SubsampledNormals, idx) n_data = length(m.dists) return SubsampledNormals(m.dists[idx], n_data / length(idx), m.capability) end From c4aaa5f5a37a97a64339e52813bbc9b5cce295c2 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Fri, 5 Jun 2026 17:27:36 +0100 Subject: [PATCH 6/7] Note that argument-passed observations also work for subsampling Both `bayes_logreg(X_batch, N) | (y = y_batch,)` (conditioning) and `bayes_logreg(X_batch, y_batch, N)` (argument) route per-batch logpdf contributions into `LogLikelihoodAccumulator`, so the SG correction applies identically. Adds a side note to the subsampling tutorial. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/src/tutorials/subsampling.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/src/tutorials/subsampling.md b/docs/src/tutorials/subsampling.md index 03ce43cb2..bcf9ca4dc 100644 --- a/docs/src/tutorials/subsampling.md +++ b/docs/src/tutorials/subsampling.md @@ -327,6 +327,25 @@ q0 = FullRankGaussian(zeros(dim), LowerTriangular(Matrix{Float64}(0.6 * I, dim, q, _, _ = AdvancedVI.optimize(alg, 1000, prob, q0; show_progress=false) ``` +!!! note "Conditioning vs. model arguments" + + Observations may be supplied either by conditioning a free random variable + (as above) or by passing them as a model argument: + + ```julia + DynamicPPL.@model function bayes_logreg(X_batch, y_batch, N) + β ~ MvNormal(zeros(size(X_batch, 2)), I) + return y_batch ~ arraydist([BernoulliLogit(dot(X_batch[i, :], β)) for i in 1:N]) + end + + minibatch_model = batch -> bayes_logreg(X[batch, :], y_obs[batch], length(batch)) + ``` + + Both forms route the per-batch contributions into `LogLikelihoodAccumulator`, + so the SG correction applies identically. Use whichever reads more + naturally — typically arguments for densely-observed data and conditioning + when the same model shape is also used outside the SG-VI loop. + The variational parameter shape must be **invariant across batches**. Above, `β` is a `d`-vector regardless of `N`; `N` only controls how many likelihood terms are summed. Models with per-datapoint latent variables (e.g. `users[:, j]` From 3c12ec49650409714e9e32bf36776b923fff7114 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Fri, 5 Jun 2026 17:49:16 +0100 Subject: [PATCH 7/7] Drop unresolvable @ref cross-references `SubsampledLogDensity` and `WeightedLogJoint` are not listed under any `@docs`/`@autodocs` block in the manual (35 other docstrings have the same status), so Documenter can't resolve `@ref`. Plain code-style backticks render the same monospace name without breaking the build. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/src/tutorials/subsampling.md | 4 ++-- src/subsampled_logdensity.jl | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/src/tutorials/subsampling.md b/docs/src/tutorials/subsampling.md index bcf9ca4dc..1bc2492b1 100644 --- a/docs/src/tutorials/subsampling.md +++ b/docs/src/tutorials/subsampling.md @@ -289,8 +289,8 @@ Therefore, subsampling is most beneficial when a crude solution to the VI proble For `DynamicPPL` models, write the model as a **factory parametric in the minibatch size `N`**, leaving observations free so they can be supplied via `|` (conditioning) on each batch. The package extension defines a call method -on [`AdvancedVI.WeightedLogJoint`](@ref) that wires `scale * loglikelihood + logprior - logjacobian` through `DynamicPPL`'s accumulators; wrap the -resulting LDF factory in [`SubsampledLogDensity`](@ref). +on `AdvancedVI.WeightedLogJoint` that wires `scale * loglikelihood + logprior - logjacobian` through `DynamicPPL`'s accumulators; wrap the +resulting LDF factory in `SubsampledLogDensity`. ```julia using AdvancedVI, ADTypes, DynamicPPL, Distributions, LinearAlgebra, LogDensityProblems diff --git a/src/subsampled_logdensity.jl b/src/subsampled_logdensity.jl index 5726cfc41..b442dcc78 100644 --- a/src/subsampled_logdensity.jl +++ b/src/subsampled_logdensity.jl @@ -1,7 +1,7 @@ """ SubsampledLogDensity(prob, make_prob, dataset_size) -`LogDensityProblems`-compatible wrapper that supports [`subsample`](@ref): +`LogDensityProblems`-compatible wrapper that supports `subsample`: `subsample(prob, batch)` returns a fresh wrapper whose inner problem is `make_prob(batch, dataset_size / length(batch))`. `make_prob` must return objects of the same concrete type as the initial `prob` so the wrapper stays