add minibatch subsampling (doubly stochastic) objective

[Red-Portal]

This is a draft for the subsampling variational objective, which addresses #38 . Any perspectives/concerns/comments are welcome! The current plan is to only implement random reshuffling. As I recently showed that there is no point in implementing independent subsampling. Although importance sampling of datapoints could be an interesting addition, it will require custom DynamicPPL.Contexts.

The key design decisions is the following function:

# This function/signature will be moved to src/AdvancedVI.jl
"""
    subsample(model, batch)

# Arguments
- `model`: Model subject to subsampling. Could be the target model or the variational approximation.
- `batch`: Data points or indices corresponding to the subsampled "batch."

# Returns 
- `sub`: Subsampled model.
"""
prob_sub = subsample(prob, batch)
q_sub = subsample(q, batch)

Given a previous slack DM thread with @yebai , this interface could be straightforwardly implemented by Turing models as

@model function pPCA(X::AbstractMatrix{<:Real}, k::Int; data_or_indices = 1:size(X,1))
    N, D = size(X)
    N_sub = length(batch_idx)

    W ~ filldist(Normal(), D, k)
    Z ~ filldist(Normal(), k, N)

    # Subsampling
    # Some AD backends are not happy about `view`.
    # In that case, this step will commit a copy and, therefore, shall not be considered free.
    Z_sub = view(Z, :, idx)
    X_sub = view(X, :, idx)

    genes_mean = W * Z_sub
    return X_sub ~ arraydist([MvNormal(m, Eye(N_sub)) for m in eachcol(genes_mean')])
end;

where data_or_indices could be made a reserved keyword argument for Turing models. Then, I think

using Accessors

function subsample(m::DynamicPPL.Model, batch)
    n, b = length(m.defaults), length(batch)
    m = @set m.defaults = batch
    m = @set m.context = MiniBatchContext(context=m.context; b, n)
    m
end

should generally work?

My current guess would be that subsample(m::DynamicPPL.Model, batch) would have to end up in the main Turing repository.

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 29 lines in your changes missing coverage. Please review.

Project coverage is 82.92%. Comparing base (1b36c6e) to head (644d314).
Report is 12 commits behind head on master.

Files	Patch %	Lines
src/objectives/subsampling.jl	0.00%	29 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master      #84       +/-   ##
===========================================
- Coverage   96.09%   82.92%   -13.18%     
===========================================
  Files          11       12        +1     
  Lines         205      246       +41     
===========================================
+ Hits          197      204        +7     
- Misses          8       42       +34     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[Red-Portal]

related: TuringLang/DynamicPPL.jl#633

[Red-Portal]

Here is a comparison of the convergence speed of full-batch v.s. subsampling with respect to wall-clock time.

[yebai]

I suggest delaying this to AdvancedVI v0.4 so that the syntax in DynamicPPL is implemented.

[Red-Portal]

@yebai Sounds good to me.

[mhauru]

I don't understand the theory nor the wider context of AdvancedVI, so my review is quite shallow, but I don't see any significant problems here. I left a few small local proposals.

[mhauru]

Suggested change

	Notice that, when computing the log-density, we multiple by a constant `likeadj`.
	Notice that, when computing the log-density, we multiply by a constant `likeadj`.

[mhauru]

Are these .svg files missing from the repo? I haven't looked at the built docs, just don't see them in the PR.

[mhauru]

Do I understand correctly that subsampling is a more general operation than just a VI thing? If that's the case, could this be moved to DynamicPPL, or even AbstractPPL?

Also, I wonder if an empty function without methods would make more sense. Is returning the unmodified original model a reasonable fallback? I could imagine it confusing users, who would call it and get a return value, not realising it's actually just the original model.

[Red-Portal]

@yebai Any comments on the current direction on the PPL side?

[mhauru]

Could the varargs be explained in the docstring?

[mhauru]

Could you comment on what happens if batchsize does not divide length(data), or whether that's significant at all?