Add SubsampledLogDensity

docs/src/tutorials/subsampling.md

@@ -283,3 +283,72 @@ 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` 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
+
+DynamicPPL.@model function bayes_logreg(X_batch, N)
+    d = size(X_batch, 2)
+    β ~ MvNormal(zeros(d), I)
+    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)
+
+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)
+```
+
+!!! 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]`
+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`.

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

src/AdvancedVI.jl

@@ -315,8 +315,9 @@ subsample(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

src/subsampled_logdensity.jl

@@ -0,0 +1,69 @@
+"""
+    SubsampledLogDensity(prob, make_prob, dataset_size)
+
+`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
+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. `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.
+"""
+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 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)`; " *
+            "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