Black-box Pseudodata Variational Inference
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Recent advances in coreset methods have shown that a selection of representativedatapoints can replace massive volumes of data for Bayesian inference, preservingthe relevant statistical information and signiﬁcantly accelerating subsequent downstream tasks.
Existing variational coreset constructions rely on either selectingsubsets of the observed datapoints, or jointly performing approximate inferenceand optimizing pseudodata in the observed space akin to inducing points methodsin Gaussian Processes.
So far, both approaches are limited by complexities inevaluating their objectives for general purpose models, and require generatingsamples from a typically intractable posterior over the coreset throughout inferenceand testing.
In this work, we present a black-box variational inference algorithmfor coresets that overcomes these constraints and enables principled application ofvariational coresets to intractable models, such as Bayesian neural networks.
Weapply our techniques to supervised learning problems, and compare it with existingapproaches in the literature for data summarization and inference.
Introduction
Recent work [25] has demonstrated that Bayesian posteriors on largescale datasets can be sparsely represented via surrogate densities deﬁnedthrough a set of learnable weighted pseudo-observations in data space.
In a similar vein with inducing points in Gaussian Processes [38], thesetarget to represent parsimoniously the sufﬁcient statistics of the observeddataset that are required for inference in the model, and learned via a variational scheme termed Pseudodata Sparse Variational Inference (PSVI).
Replacing the original dataset via the reduced set of pseudo-observationsunlocks the potential of scaling up downstream learning tasks, compressing for efﬁcient storage and visualisation, and performing rapid modelexploration.
See Fig.
1 for an introductory visualization.
However, in thegeneral case the achievable quality of data summarizations for posteriorapproximations is often limited by the intractability of the original model,hindering the attainable convergence guarantees of coreset constructions.
In an effort to make PSVI available to general black-box models, we make the following contributions:
we derive a principled interpretation of variational coresets as rich variational families,
we modify the objective to obtain a black-box version of it using importance sampling and anested variational inference step,3.
utilizing this objective, we propose BB PSVI and BB Sparse VI, two novel suites of
variational inference algorithms for black-box probabilistic models,
we empirically compare to prior Bayesian and non-Bayesian coreset techniques, and
we showcase these techniques on previously infeasible models, Bayesian neural networks.
Background2.1 Variational Inference
Consider a modeling problem where we are given a standard model:
[Equation]
Given a dataset D =
{ ( x , y ) } ∈
( X × Y ) N , the task of inference and maximization of the marginallikelihood p ( y | x ) entails estimation of the posterior distribution over model parameters θ , p ( θ | x , y ) .
In general, the posterior p ( θ | x , y ) is not tractable.
A common solution is to resort to approximateinference, for instance via variational inference (VI).
In VI, we assume that we can approximatethe true posterior within a variational family q ( θ ; λ ) with free variational parameters λ , by minimizing the Kullback-Leibler (KL) divergence between the approximate and the true posterior
D KL ( q ( θ ; λ ) || p ( θ | x , y )) , or equivalently by maximizing the Evidence Lower Bound (ELBO):
[Equation]
2.2 (Bayesian) coresets and pseudocoresets
Bayesian coreset constructions aim to provide a highly automated data austere approach to theproblem of inference, given a dataset D and a probabilistic model Eq.
Agnostic to the particular inference method used, Bayesian coresets shrink in a model-speciﬁc way the original datasetto a much sparser weighted dataset representation { ( v , x , y )
Nn =1 } ∈ ( R ≥ 0 ×
X × Y ) N , where||
v ||
= M (cid:28) N , that maximizes the pasimony of observations considered over inference, while allowing close approximations of the full data posterior, i.e. p ( θ | x , y , v ) ≈ p ( θ | x , y ) .
Initial Bayesiancoreset constructions formulated the problem as yielding a sparsity-constrained approximation of thetrue data likelihood function uniformly over the parameter space.
This perspective has admitted solutions via convexiﬁcations relying on importance sampling and conditional gradient methods [19, 10],greedy methods for geodesic ascent [9], as well as non-convex approaches employing iterative hardthresholding [43].
Recently, Sparse VI [8] reformulated the problem as performing directly VI withina sparse exponential family of approximate posteriors deﬁned on the coreset datapoints.
The idea of using a set of learnable weighted pseudodata (or inducing points ) ( v , u , z ) :
={ ( v m , u m , z m )
Mm =1 } ∈ ( R ≥ 0 ×
X × Y )
M to summarize a dataset for the purpose of posteriorinference was introduced in PSVI [25], as the following objective:
[Equation]
This objective facilitates learning weighted pseudodata such that the exact posterior over pseudodatap ( θ | u , z , v ) approximates the exact posterior over the true dataset p ( θ | x , y ) , and was demonstratedon a variety of models as a generally applicable objective.
However, in [25] it was assumed that thisobjective is tractable as demonstrated in experiments on models admitting analytical posteriors, orvia the use of heuristics to sample from the coreset posterior, as is the case in a logistic regressionemploying a Laplace approximation.
Sparse VI Sparse VI [8] targets to incrementally minimize the rhs of Eq.
(3) via ﬁxing thecoreset points locations and optimizing their weights v .
In its original construction, the coreset points are initialized to the empty set and v to 0 .
In the general case of an intractablemodel, the coreset posterior is approximately computed via ﬁtting a Laplace approximationq Lap ( θ ; v ) on the weighted datapoints of the coreset ( v , x , y ) .
The next point selection step involves: ( i ) computing approximations of the centered log-likelihood function of each datapoint( x i , y i ) for model parameters θ s sampled from the current coreset posterior, via an MC estimate ˜f (
[Equation]
n ,
[Equation]
n ,
[Equation]
[Equation]
,( ii ) computing correlations with the residual error of the total data log-likelihood r ( v ) =
f T (1 − v ) ,and ( iii ) making a greedy selection of the point that maximizes the correlation between the two:
[Equation]
Corr ( ˜f , r ) v n = 0 , where Corr ( ˜f , r ) :
[Equation]
Subsequently, the vector v gets optimized via minimizing the KL divergence between the coreset and the true posterior.
For the variational family of coresets, under the assumption that θ is sampledfrom the true coreset posterior, the formula of the gradient of the objective in Eq.
(3) can be derived,and approximated via resampling from the coreset over the course of v optimization.
At the end ofthe optimization, the extracted points and weights are given as input to an approximate inferencealgorithm—commonly a Laplace approximation—which provides the ﬁnal estimate of the posterior.
Black-box coreset Variational Inference
In order to derive a general PSVI objective, we observe that we ﬁrst need to posit a variationalfamily involving the pseudodata appropriately.
Let’s assume the (intractable) variational familyq ( θ | u , z ) :
= p ( θ | u , z ) , which has the following particular structure: it is parametrized as theposterior over θ as estimated given inducing points u , z .
In this case, q ( θ | u , z ) can be thought of asa variational program (see [34]) with parameters { u , z } and expressed as:
[Equation]
where the marginal likelihood term p ( z | u ) is typically intractable.
Utilizing this for the formulation of a pseudo-coreset variational inference algorithm leads to thefollowing lower bound on the marginal likelihood of the true data:
[Equation]
Performing VI involves maximizing Eq.
(5) with respect to the parameters { u , z } .
Scrunitizing this objective reveals that the posterior p ( θ | u , z ) is used in two locations: ﬁrst , asthe sampling distribution for θ , second to evaluate log-densities of samples ( scoring ).
As such,without further assumptions on the inference algorithm, this objective can only be evaluated if theposterior p ( θ | u , z ) can be readily evaluated in closed form and sampled from directly, causing thisexpression to remain intractable but for the simplest of cases; second , the same quantity exists insidethe logarithm as a quantity to evaluate , revealing the need to make it computable for any sample θ .
In Section 3.1 we will derive a black-box objective for PSVI to overcome the limitation to tractableposteriors by adressing these two points separately and composing these solutions to a joint blackbox objective, which we will use for deriving black box versions of PSVI and Sparse-VI.
In the appendix we also derive a Gaussian Process version of our algorithm as a special case wherethe posterior is tractable to formally link this work to inducing points for GPs.
3.1 PSVI for intractable posteriors
In Sec.
3 we identiﬁed two key problems to evaluate Eq.
(5) involving p ( θ | u , z ) .
To tackle the ﬁrst, wedevise a strategy to utilize a tractable distribution q ( θ | u , z ; ψ ) paired with importance sampling [30]to be able to sample from that distribution.
We assume that we can draw K samples θ k ∼ q ( θ ; ψ ) , 1however we cannot plug these into Eq.
(5) directly as they do not follow the exact distribution.
Yet,we can leverage a self-normalized importance sampling (IS) [30] correction to obtain the desiredapproximate samples from the coreset posterior.
Assigning the samples unnormalized weights:
[Equation]
for an appropriately selected q such that support ( p ) ⊆ support ( q ) , we denote the correspondingnormalized weights ˜ w k = w k / (cid:80) j w j .
Using q ( θ | u , z ) from Eq.
(4), we can approximate Eq.
(5) as:
ELBO PSVI-
[Equation] [Equation]
After having tackled sampling, we shift our attention to the second problem, which was aboutevaluating p ( θ | u , z ) inside the logarithm .
Applying Jensen’s inequality and Eq.
(4), we get anevidence lower bound of the inducing points (IP):
[Equation]
Thus, substituting q ( θ i | u , z ) , i = 1 , . . . ,
M , per Eq.
(4) in Eq.
(7), we can derive the followingmulti-sample importance sampling lower bound on the evidence of the observed data:
[Equation]
We now realize that by lower bounding log p ( z | u ) using ELBO IP ( ψ ) we are jointly deriving a lowerbound to ELBO PSVI-IS ( u , z ) , which we denote as ELBO PSVI-IS-BB ( u , z , ψ ) .
We remind the reader atthis point that this ensures the following relationship holds rigorously:
[Equation]
which means that the following multi-sample black-box objective, ELBO PSVI-IS-BB , is a valid variational inference objective for the original problem under which we can tractably score u , z :
[Equation]
Maximizing the objective of Eq.
(9) can directly serve the purposes of black-box VI, as it lifts therequirement for sampling from the exact posterior on the summarizing data, proposing instead importance sampling with the use of q ( θ ; ψ ) constrained on the summary.
Crucially, as the ELBO PSVI-IS-BBmixes information from the true and inducing datapoints, the assumed probabilistic structure ofthe variational distribution should be enforced over the optimization of Eq.
(9), namely that ψ isbeing inferred only on the basis of the inducing points via the variational program of Eq.
(8), whilst( u , z ) get optimized according to Eq.
(9), so that the information of the true dataset gets distilledin a way that increases the evidence lower bound.
Moreover, replacing our pseudodata dependentimportance weights with a uniform distribution ˜ w k = 1 /K , k = 1 , . . .
, K , or using a single-sampleobjective (i.e. in the degenerate case when K = 1 ), allows us to recover the classical ELBO of Eq.
(2),cancelling the pseudodata likelihood terms in the variational objective.
2 Iteratively maximizing the
ELBO PSVI-IS-BB enables adapting the importance sampling scheme of Eq.
(6) to the intractable coresetposterior, automating the selection of improved proposals q for the importance weighting throughoutinference, within the conﬁnes of the chosen variational family.
Finally, to compute expectations under the intractable true posterior using samples from the learnedcoreset posterior, assuming q ( θ | u , z )
≈ p ( θ | x , y ) , we can sample from q and correct again usingthe above deﬁned importance sampling scheme as follows:
[Equation]
3.2 Incremental and Batch Sparse BBVI
The estimator we derive can also be used to design a black box version of the Sparse VI algorithmintroduced in [8] (see Sec. 2.2), by the insight that only the variational parameters for locationsare not updated and we can apply various choices to the overall algorithmic ﬂow.
The incrementalconstruction scheme originally can become black-box via modifying how the approximate posterioron the coreset data gets computed, and introducing a generalized objective for the optimizationof model evidence involving the coreset parameters.
In this construction ( Sparse BBVI ), we can posit a standard variational family q ( θ ; ψ ) (e.g. mean-ﬁeld variational distributions) and maximizethe corresponding ELBO computed on the weighted datapoints of the coreset.
Using the extractedvariational approximation, we can draw samples from the coreset and correct for the model.
Next,over the greedy selection step, we can use the samples and importance weights to compute thecentered log-likelihood vectors ˜f ( x n , y n , θ s ) :
= log p ( y n | x n , θ s ) − (cid:80) s (cid:48) w s (cid:48)
[Equation]
p ( y n | x n , θ s (cid:48) ) , andselect the next datapoint via the greedy correlation maximization criterion.
Subsequently, to reﬁnethe weight vector towards the true data posterior, we can maximize wrt v , ψ an extension of the
PSVI ELBO per Eq.
(9), where the coreset data log-likelihood is multiplied by the weight vector v asin Eq.
[Equation]
The importance sampling scheme for the posterior samples is deﬁned identically to Eq.
(6), aftersubstituting the likelihood term with p ( y | v , x , θ i ) for each sample θ i ∼ q ( θ i ; ψ ) .
At each iterationof the algorithm, this variational objective (which, following similar reasoning with the previoussection, can be shown to be a lower bound of the evidence), can be maximized with respect to thecoreset and variational parameters v , ψ .
Moreover, we can omit the incremental inclusion of points tothe coreset, and consider a batch version of this construction, where we initialise at a random subsetof the original dataset, keep ﬁxed coreset point locations, and optimize only the weights attached tothe coreset support using Eq.
(11) with non-negativity constraints.
Finally, we can use the variationalapproximation ﬁtted on the coreset to compute predictive posteriors on unseen datapoints correctingfor our importance sampling.
We also deﬁne a variant of the algorithm without incremental selection through a simple pruningstrategy based on the weighting assigned to the coreset points when optimizing the ELBO of Eq.
We can start from a large coreset size and, after training, reduce the size of the summary by keepingthe top- K datapoints and re-initialising.
3.3 Variational Families: (Un)weighted inducing points
When studying the objective in Eq.
5 it quickly becomes evident that the exact amount of chosenpseudo-data (coreset) deﬁnes the available evidence for that particular posterior construction and assuch is a quantity of interest.
We will consider this quantity pseudo-evidence .
For notational brevity so far we have considered the variational parameters to represent pseudodatadirectly, while rescaling using the ﬁxed data compression ratio to maintain invariance of the totalpseudo-evidence to the coreset size.
We will be calling this setting unweighted pseudodata .
One choice we can make now is to ensure that the available pseudo-evidence is invariant to the chosensize of the coreset set and potentially equal to the regular dataset.
To achieve that, we consider setting:
[Equation]
with p ( z | u ) similarly reweighted via the selected observational data compression ratio N/M , i.e.for continuous parameter spaces p ( z | u ) :
[Equation]
We may also consider a scenario where we have more generally weighted pseudodata , which involvesan additional learnable parameter v i ≥ 0 per datapoint
u i and permits the following interpretation:v i ∼ Mult ( v ) .
Sampling repeatedly from this distribution yields a collection of inducing points withtheir frequencies governed by their respective weights:
[Equation]
[Equation]
This allows us to introduce the following simplifying notation for weighted log-likelihoods, whichwill be used in the continuation:
[Equation]
We can now also combine the ideas of learning the weights of pseudodata and controlling theamount of pseudoevidence by posing parametrizations where v i is non-negative and sums to {+
+}
(e.g. v i := softmax ( β i ) , i = 1 , . . . , M ), and adding global variational parameters to optimize themagnitude of total evidence over the pseudodata (e.g. log p ( z | v , u , θ , α ) :
= α v T log p ( z | u , θ ) ),where α can be set by hand (i.e. to N/M ) or be a learnable quantity.
We will explore various suchchoices for coreset parametrizations in our experiments and the effects they have on learning.
3.4 Nested optimization
When using a nested optimizer for the maximization of the ELBO PSVI-IS-BB , rearranging the terms ofthe objective enables us to place our method in the generic cardinality-constrained bilevel optimizationframework for data distillation and coreset constructions proposed in [21, 5].
VI maximizing theobjective of Eq.
(5) inside a parametric family q ( θ ; ψ ) can be reformulated as a bilevel optimizationproblem in the following way:
[Equation]
where
[Equation]
[Equation]
and
[Equation]
The bilevel structure allow us some freedom in treating our optimization problem.
Two principledapproaches are:
• Iterative differentiation [24]: We trace the gradient computation over the optimization of thevariational parameters ψ when solving the inner problem, and use this information at thecomputation of the outer objective gradient wrt the coreset parameters.• Approximate hypergradients [32, 21, 18]: We can compute approximations of the outergradient using the information of an approximate solution in the inner problem, withoutfully tracing the gradient throughout the inner optimization trajectory.
Our method of choice in the experimental evaluation will be iterative differentiation, while we includesome additional results using hypergradients in the supplementary material.
Experiments
In this section we evaluate the performance of our inference framework in intractable models andcompare against standard variational inference methods and earlier Bayesian coreset constructions, aswell as black-box extensions of existing variational coresets that rely on our generalized ELBO Eq.
As a running baseline with unrestricted access to the training data, we use the standard mean-ﬁeld
VI with a diagonal covariance matrix.
To capture the implications of enforcing data cardinalityconstraints, we also construct a random coreset using a randomly selected data subset with ﬁxed datalocations and ﬁxed likelihood multiplicities correcting for the dataset size reduction, which we thenuse to optimize a mean-ﬁeld approximate posterior.
As earlier work relied on Laplace approximationson coreset data as a black-box approach for sampling for intractable coreset posteriors, we alsoexperiment with Laplace approximations on random data subsets as an extra baseline.
4.1 Logistic regression
In the ﬁrst experiment we perform inference on logistic regression models ﬁtting 3 publicly availablebinary classiﬁcation datasets [41, 14] with sizes ranging between 10 K and 150 K datapoints, and10 and 128 dimensions.
At all models we posit normal priors θ ∼ N (0 , I ) and consider mean-ﬁeldvariational approximations with diagonal covariance.
In this and following experiments, all presentedmetrics are averaged across 3 independent trials for each method, and we show the means over thetrials along with the corresponding standard errors. Impact of variational parameters We show predictive metrics of log-likelihood and accuracy onthe test set inTables 1 and 2. We observe that coreset methods relying on optimizing
[Table]
[Figure] [Table]
have the capacity to approximate quickly inference on the full dataset regardless of dimension, andreach the performance of unrestricted mean-ﬁeld VI with the use of less than 50 weighted points.
Incontrast, the approximations that rely on existing datapoints are limited by data dimensionality andneed to acquire a larger support to reach the performance of full-data VI (with effects being moreevident in webspam, which has a data dimensionality of 128).
Including variational parameterisationsfor optimal scaling of the coreset likelihood aids inference, offering faster convergence.
Omittingto correct for dataset reduction ( v = 1 ) has detrimental effects on predictive performance.
Asopposed to earlier constructions limited by heuristics on coreset posterior approximation throughoutoptimization, our variational framework enables approximating the full data posterior within theassumed variational family.
Evaluation of predictive metrics and computation time requirementsacross a wider range of coreset sizes is included in the Supplement.
Impact of importance weighting In Fig.
2 we plot the normalized effective sample size (ESS) ofour importance weighting scheme for 10 Monte Carlo samples computed on test data.
We observethat the posteriors can beneﬁt from IW, achieving non-trivial ESS, which tend to converge to closevalues for large coreset sizes.
This ﬁnding applies at the testing phase also on ablations of the PSVIvariational training that replace the IW scheme with uniform sampling.
4.2 Bayesian Neural Networks
In this section we present inference results on Bayesian neural networks (BNNs), a model class thatprevious work on Bayesian coresets did not consider due to the absence of a black box variationalestimator.
In the ﬁrst part we perform inference via black-box PSVI on 2-dimensional synthetic datausing single-hidden layer BNNs, while in the latter part we evaluate the performance of our methodson compressing the MNIST dataset using the LeNet architecture [20].
Simulated datasets We generate two synthetic datasets with size 1 k datapoints, corresponding tonoisy samples from a half-moon shaped 2-class dataset, and a mixture of 4 clusters of data each
[Table]
±
[Table]
[Table] [Table] [Table]
[Table] [Table] [Table]
[Table]
[Table] [Table] belonging to a different class, and use a layer with 20 and 50 units respectively.
To evaluate therepresentation ability of the pseudocoresets we consider two initialization schemes: we initialise thepseudo locations on a random subset equally split across categories, and a random initialization usinga Gaussian centered on the means of the empirical distributions.
In Fig.
3 we visualize the inferredpredictive posterior entropy.
Regardless of the initialization, BB PSVI has the capacity to optimizethe locations of the coreset support according to the statistics of the true data distribution, yieldingcorrect decision boundaries.
Importantly, the coreset support consists itself of separable summarizingpseudo points, while the likelihood reweighting variational parameters enable emphasizing on criticalregions of the landscape, assigning larger weights to pseudopoints lying close to the boundariesamong the classes.
Moreover, our pruning scheme which makes use of batch Sparse BBVI is able togradually compress a large coreset of 250 datapoints to a compact performant subset of 20 datapointsarranged in critical locations for learning the Bayesian posterior.
MNIST In this part we assess the approximation quality of large-scale dataset compression for
BNNs via coresets.
We compare the predictive performance of black-box PSVI against standardmean-ﬁeld VI, random coresets and frequentist methods relying on learnable synthetic data, namelydataset distillation w/ and w/o learnable soft labels [42, 37], and data condensation [44].
We can seethat BB PSVI consistently outperforms the random coreset baseline for coreset sizes in the ordersof few tens or hundreds, until the random coreset eventually meets its performance for sizes near
(Table 3).
Moreover, at size 100 our construction provides better test accuracy compared to twofrequentist approaches, enjoying the additional advantage of providing uncertainty estimates.
[Table]
Related work
Coresets and data pruning beyond Bayesian inference.
The paradigm of coresets emerged fromseminal work in computational geometry [1, 16], and has since found broad use in data-intensivemachine learning, including least mean square solvers [23], dimensionality reduction [17], clustering [15, 22, 2] and maximum likelihood estimation [3, 26].
Coreset ideas have also found recentapplications in active learning [33, 6], continual learning [5, 4] and robustness [27].
Recent works inneural networks have investigated several heuristics for pruning datasets, including the misclassiﬁcation events of datapoints over learning [40], or gradient scores in early stages of training [31].
Pseudodata for learning.
Despite the multiple breakthroughs achieved by classical coreset methods, the common conﬁne of selecting the representative data instances among the constituent datapoints might lead to large summary sizes in practice.
Exempt from this limitation are methodsthat synthesize informative pseudo-examples for training that can summarize statistics lying awayfrom the data manifold.
These appear in the literature under various terms, including dataset distillation [24, 42, 21], condensation [44], dataset meta-learning [28] and inducing points [29].
Moreover,recent work [39] has demonstrated that incorporating learnable pseudo-inputs as hyperparametersof the model prior can enhance inference in variational auto-encoders.
Learnable pseudodata thatresemble the sufﬁcient statistics of the original dataset have also been used in herding [11] andscalable GP inference methods [36, 38].
Variational inference, importance weighting, and Monte Carlo objectives.
Importance weighting can provide tighter lower bounds of the data log-likelihood for training variational autoencoders [7], resulting in richer latent representations [12].
In later work, [13] propose an importanceweighting VI scheme for general purpose probabilistic inference.
[35, 34] introduce variationalprograms as programmatic ways to specify a variational distribution.
Conclusion
We have developed a novel black-box construction for variational coresets that can be applied at scaleto large datasets as well as richly structured models, and result in effective model-conditional datasetsummarization.
We overcome the intractability of previous inference algorithms with our proposedblack box objectives, which we develop for both Sparse VI as well as PSVI, and propose a suite ofalgorithms utilizing them.
Through experiments on intractable models, we show that the inferred coresets form a useful summaryof a dataset with respect to the learning problem at hand, and are capable to efﬁciently compressinformation for Bayesian posterior approximations.
In future work, we plan to extend this inferencetoolbox for privacy purposes and non-traditional learning settings, e.g. continual learning.