Resampled Priors for Variational Autoencoders

1 Introduction

• The aggregate posterior minimizes the ELBO’s KL term, but using it directly can overfit by memorizing the training set [30].VampPrior models it with a fixed-size mixture of learned virtual-observation posteriors; flow-based and autoregressive priors offer alternative expressivity-efficiency tradeoffs [30],,.Source: §1 · p. 1
• LARS transforms a tractable proposal using a learned acceptance function, which can be trained jointly with the model or separately on aggregate-posterior samples.It can also use expressive existing priors as proposals, making it complementary to many prior-flexibility approaches.Source: §1 · p. 1
• LARS is a generic method for making tractable base distributions more expressive and is not limited to variational inference or continuous variables.It modulates the base density while keeping learning or exact sampling tractable.Source: §1 · p. 1
• A truncated LARS sampler trades approximation accuracy against sampling efficiency by accepting the next proposal after a predetermined number of consecutive rejections.The truncation is introduced to improve acceptance and reduce inefficient repeated rejection.Source: §1 · pp. 1–2
• Replacing standard Normal VAE priors with LARS priors consistently improves ELBOs, with further gains from expressive flow-based proposals.The introduction also demonstrates LARS on a VAE’s high-dimensional discrete output distribution.Source: §1 · p. 2

• VAEs [17] [25] model observations with latent variables and are trained by maximizing an ELBO because computing the marginal likelihood is intractable.The encoder parameterizes the variational posterior and the decoder parameterizes the likelihood; ELBO gradients can be estimated using the reparameterization trick [17] [25].Source: §2 · p. 2
• The ELBO’s reconstruction term measures decoding fidelity, while its KL term regularizes latent information and supports generation rather than reconstruction alone.The KL term is between the variational posterior and the prior.Source: §2 · p. 2
• Trained VAEs can retain a mismatch between the prior and aggregate posterior, even though maximizing the ELBO minimizes their KL divergence.The aggregate posterior is q(z) = E_q(x)[q(z|x)].Source: §2 · p. 2
• Prior-heavy regions with little aggregate-posterior density are “holes” that are rarely encountered in training and can decode to observations off the data manifold.These regions have high prior density but very low aggregate-posterior density.Source: §2 · p. 2
• The paper aims to make the prior more expressive and closer to the aggregate posterior without making the distributions identical, which could cause overfitting.Its proposed approach learns to reject proposal samples with low aggregate-posterior density.Source: §2 · p. 2

3 Learned Accept/Reject Sampling

• LARS transforms a tractable proposal into a more expressive density by multiplying it by a learned acceptance function and renormalizing; unlimited acceptance-function capacity can match the target exactly.Rejection sampling draws from the proposal and accepts each candidate with probability a(z), repeating after rejection.Source: §3 · pp. 2–3
• Truncating rejection sampling at T candidates yields a mixture of the resampled density and proposal, guarantees acceptance of at least 1/T, and trades approximation accuracy against sampling efficiency.T = 1 recovers the proposal, while T → ∞ recovers untruncated sampling.Source: §3 · p. 3
• For training, LARS estimates Z once per batch, smooths its forward-pass value exponentially, computes backward-pass gradients from current-batch samples, and adds a reweighted target sample to reduce gradient variance.For evaluation, it estimates Z once with many samples after training.Source: §3 · p. 3

3.1 Illustrative examples in 2D

• A two-layer 10-unit acceptance network nearly separates the Gaussian-mixture modes, while a one-layer network improves on the proposal; a jointly trained RealNVP proposal isolates all modes and raises average acceptance probability.The higher-capacity network fits the target more closely but has lower sampling efficiency (smaller Z).Source: Figure 2 · p. 3, §3.1 · p. 4

3.2 Relationship to rejection sampling

• Unlike classical rejection sampling, which assumes an evaluable target density, LARS addresses targets that are sampleable but infeasible to evaluate; with unlimited acceptance-function capacity it can match the target, while prioritizing a compact, fast-to-evaluate approximation.In this setting Z plays the role of the classical sampler's acceptance rate 1/M.Source: §3.2 · p. 4

4 VAEs with resampled priors

• The original VAE prior serves as proposal π(z), while a flexible neural network defines acceptance probability a(z) and the resulting resampled prior p(z).LARS can enrich several distributions, but this section focuses on the prior; “LARS prior” and “resampled prior” are used interchangeably.Source: §4 · p. 4
• Joint training replaces only the standard Normal prior term in the ELBO, and requires no accept/reject sampling during training.Training evaluates the resampled prior density on variational-posterior samples and uses proposal samples to update the moving-average estimate of Z; the maximum resampling steps T can be chosen for sampling-time compute.Source: §4 · p. 4

Hierarchical models

• In hierarchical VAEs following [30], resampling is applied to the top-most latent prior pλ(z2).The model has two stochastic layers, with generative distribution pθ(x|z1, z2)pλ(z1|z2)pλ(z2) and variational posterior qψ(z1|x, z2)qφ(z2|x); their means and standard deviations are neural-network parameterized.Source: §4 · p. 4

5.1 Quantitative Results

• Across all tested architectures and datasets, the resampled prior outperforms the standard Normal prior, usually improving test NLL by more than 1 nat.The gains are notably larger for simpler architectures.Source: Table 1 · p. 5, §5.1 · p. 5
• Post-hoc LARS nearly matches joint training; its ELBO gain comes entirely from lower KL, while joint training also reduces reconstruction error.Joint training achieves a better ELBO, although its KL reduction is smaller than with post-hoc fitting.Source: Table 2 · p. 5, §5.1 · p. 5
• On dynamic MNIST and Omniglot, LARS is competitive with approaches using more expressive priors or posteriors, and its gains over the Normal prior are generally comparable to, but slightly smaller than, VampPrior [30].VampPrior [30] sampling is non-iterative, but the cited high-performing models use hundreds of thousands of parameters; LARS has over an order of magnitude fewer and trains more efficiently in latent space.Source: §5.1 · pp. 5–6, Table 3 · p. 5, Table 4 · p. 5
• More expressive acceptance networks improve NLL, but gains are marginal beyond the chosen MLP[100−100]; increasing truncation steps improves NLL while lowering Z.The truncation-related NLL gains come from both KL and reconstruction, more from KL; for a VAE, LARS adds 25% multiply-adds per iteration and about 40% wall-clock time in the unoptimized implementation.Source: §5.1 · p. 6, Table 5 · p. 6, Table 6 · p. 6
• Applying LARS to a RealNVP proposal improves NLL by more than 0.5 nats on both MNIST and Omniglot, while combining it with MAF is impractical because MAF sampling is slow.RealNVP alone outperforms resampling a standard Normal prior, but resampled RealNVP performs better than MAF prior models.Source: §5.1 · p. 6, Table 7 · p. 6

5.2 Qualitative Results

• For jointly trained and post-hoc resampled priors, samples with a(z) ≈ 1 look pleasing with few artifacts, while samples with a(z) ≤ 10^-10 look poor and are likely rejected.The learned acceptance function can therefore rank samples and reject low-quality ones.Source: §5.2 · p. 7, Figure 3 · p. 7
• Dimension-wise resampling did not improve over a factorized Normal prior; in 2D, the resampled prior better matches the aggregate posterior and joint training broadens its tails.The paper suggests dimension-wise resampling may work better when the target has strong factorial structure, but leaves this for future work; the 2D case showed no substantial NLL difference.Source: §5.2 · p. 7

5.3 Resampling discrete outputs of a VAE

• On dynamic MNIST, the VAE with resampled outputs outperforms the original VAE by over 1 nat in NLL, and is only about 0.6 nats worse than the VAE with resampled prior.Resampling outputs is slower because it requires decoding all prior samples, applying a larger a(x) on a higher-dimensional space, and using more samples to estimate Z accurately. Because x is discrete, the contribution of log Z to encoder and decoder gradients is hard to compute with the reparameterization trick; the authors ignore it, effectively fitting a(x) post hoc.Source: §5.3 · p. 7

Expressive prior distribution for VAEs

• Unlike VampPrior [30], Gaussian mixtures, autoregressive priors [8, 10, 23], and flows, LARS permits unrestricted acceptance networks and parallel candidate sampling, but needs Monte Carlo normalization.VampPrior density evaluation can be costly, Gaussian mixtures are harder to optimize and perform worse, autoregressive sampling is sequential, and flows require invertible transformations with tractable Jacobians.Source: §6 · p. 8

Rejection Sampling (RS)

• Variational Rejection Sampling and Naesseth et al. [21] use rejection sampling but require target-density evaluation, making them inapplicable to LARS’s setting.Variational Rejection Sampling makes the variational posterior closer to the exact posterior; Naesseth et al. derive reparameterization gradients using a reparameterizable proposal.Source: §6 · p. 8

Density Ratio Estimation (DRE)

• DRE has been used to estimate the ELBO’s KL term for implicit-prior VAEs [20], but tends to overestimate the ELBO, complicating model comparison [26].Source: §6 · p. 8
• LARS’s acceptance function estimates a rescaled density ratio end-to-end, prioritizing a compact model that balances aggregate-posterior fit with generalization over exact ratio estimation.Its explicit density allows reliable model comparison by estimating the normalizing constant with many samples.Source: §6 · p. 8

Products of Experts

• LARS’s induced density is a two-expert product-of-experts, yet unlike typical PoE models requiring MCMC, it yields exact independent samples through rejection sampling.Source: §6 · p. 8

• LARS approximates target density q by reweighting proposal samples with a learned acceptance function, and applies to both continuous and discrete variables.Source: §7 · p. 8
• LARS priors outperform a standard Normal prior by a considerable margin across several VAE architectures and datasets, and are competitive with approaches enriching the posterior or prior.Source: §7 · p. 8
• LARS can be applied post hoc to trained models, and its learned acceptance function effectively ranks proposal samples by visual quality.Source: §7 · p. 8
• LARS can be combined efficiently with approaches that support efficient sampling, such as flows.Source: §7 · p. 8
• Inference remains tractable despite sampling-efficiency and normalization-constant challenges; truncated resampling provides an interpretable trade-off between sampling efficiency and approximation quality.Source: §7 · p. 8
• Beyond variational inference, discrete VAE-output resampling is a first step; future directions include multiple acceptance functions for independent dimensions or different schemes for early and later attempts.Source: §7 · p. 8

A.1 Resampling ad infinitum

• Indefinite resampling yields a density by summing the geometric series of rejected proposals followed by an accepted sample, and its expected number of steps is 1/Z.At each draw, the proposal density π(z) is weighted by acceptance probability a(z).Source: §A.1 · p. 11

A.2 Truncated Resampling

• Truncated resampling accepts the Tth proposal unconditionally after T − 1 rejections, interpolates between the proposal at T = 1 and indefinite resampling as T →∞, and takes 1 to T expected steps as acceptance falls.The supplement describes the truncated density as a mixture of the indefinite-resampling density and proposal density.Source: §A.2 · pp. 12–13

A.3 Illustration of standard/classical Rejection Sampling

• The rejection-sampling illustration shows sampling from a simpler proposal and accepting candidates with probability a(z) = q(z) Mπ(z).The figure distinguishes relative acceptance and rejection probabilities.Source: Figure A.1 · p. 13

6 end

• For both resampling schemes, the smoothed Z estimate is used in the forward pass, while the current estimate supplies gradients; the estimator also adapts to truncated densities.The gradient computation avoids backpropagating through all terms in the moving average.Source: §A.4 · p. 14, §6 · p. 15

B.1 Datasets

• The supplement reports TensorFlow implementations and dataset setups for MNIST, Omniglot, and FashionMNIST, including dynamically binarized images and statically binarized MNIST.The MNIST setup uses both dynamic and static binarization; Omniglot and FashionMNIST use dynamic binarization.Source: B Experimental Details · p. 16, §B.1 · p. 16

A.4 Estimation of the normalization constant Z

• During training, LARS evaluates the resampled prior density without accept/reject sampling; estimating its normalizer uses exponential moving averages and encoder-posterior samples to reduce noise and log-normalizer bias.For model evaluation, the supplement uses S = 1010 proposal samples; training smooths Z with ϵ = 0.1 or ϵ = 0.01 and includes the q(z|x) sample in its estimator.Source: §A.4 · pp. 13–14, §B.3 · p. 17

B.2 Network architectures

• The acceptance function uses an MLP with a logistic output, while resampled-prior training changes only the KL prior-density evaluation and leaves the encoder and decoder unchanged.The logistic output constrains acceptance probabilities to [0, 1].Source: §B.3 · p. 17, §B.2 · p. 17

B.3 Training procedure for a VAE with resampled prior

• For VAE training, the supplement specifies estimating the normalizer from proposal samples without performing accept/reject steps or decoding those samples into image space.The presented pseudocode covers indefinite resampling, with an extension to truncated resampling described as straightforward.Source: §B.3 · p. 17

B.4 Training a VAE with resampled outputs

• The normalizer is estimated using decoded samples from the regular prior, with xs sampled from the decoder given zs sampled from p(z).For binary MNIST, observations and model samples are discrete, but learning acceptance-function parameters does not require propagating gradients through samples.Source: §B.4 · p. 18
• For the decoder, the normalizer estimate depends on discrete model samples, so the authors omit its gradient contribution rather than use high-variance REINFORCE estimates.A reparameterization trick would be straightforward for continuous samples, but the MNIST samples are discrete.Source: §B.4 · p. 18
• The encoder and decoder optimize the first two ELBO terms, while the acceptance function optimizes the last, making its fitting effectively post-hoc because it does not affect the rest of model training.The authors describe this as essentially equivalent to the post-hoc setting in Table 2.Source: §B.4 · p. 18

C.1 Illustrative examples in 2D

• In the 2D examples, RealNVP alone does not isolate all modes, a shortcoming also observed for Masked Autoregressive Flows by Huang et al..Figure C.4 contrasts a RealNVP-only prior with a RealNVP proposal combined with a learned rejection sampler.Source: §C.1 · p. 19, Figure C.4 · p. 19
• Figure C.3 illustrates learned acceptance functions reweighting a N(0, 1) proposal to approximate a fixed target density.The figure depicts the acceptance function in red, target in blue, and approximate density in green.Source: Figure C.3 · p. 19

C.2 Lars in combination with a non-factorial prior on Omniglot

• On dynamic Omniglot, a RealNVP prior alone outperforms a VAE with resampled prior, and combining the RealNVP proposal with LARS further improves test NLL.The passage provides no table-cell values, so the comparison is qualitative.Source: §C.2 · p. 19, Table C.1 · p. 19

C.3 Samples from jointly trained VAE models

• High-acceptance samples from the jointly trained VAE look good, random proposal samples vary in quality, and low-acceptance samples—typically below a(z) ≈10^-10—look very poor.The samples are selected by acceptance value from 104 draws; the reported average acceptance probability is Z ≈0.014, and acceptance values correlate with sample quality.Source: §C.3 · p. 20, Figure C.5 · p. 20, Figure C.6 · p. 21

C.4 Samples from applying Lars to a pretrained VAE model

• For post-hoc LARS, the encoder and decoder remain fixed while only the acceptance function is trained using the usual ELBO objective.The starting VAE was trained with the standard Normal prior.Source: §C.4 · p. 20

C.3 Samples from jointly trained VAE models

• The jointly trained model’s acceptance values correlate with sample quality, with an average acceptance probability of Z ≈0.014.Figure C.6 shows the acceptance-value distribution and representative samples with their values.Source: Figure C.6 · p. 21

C.4 Samples from applying Lars to a pretrained VAE model

• The post-hoc sample comparison groups 104 draws into highest-acceptance, random, and lowest-acceptance samples.The figure shows samples from a pretrained MLP encoder/decoder VAE with a resampled LARS prior.Source: Figure C.7 · p. 21

C.5 Samples from a VAE with Lars output distribution

• Applying LARS to the discrete output space defines a resampled marginal log likelihood rather than resampling in the continuous latent space.The figure depicts a jointly trained acceptance function on the output.Source: §C.5 · p. 22, Figure C.8 · p. 22
• Even in the high-dimensional output space, the acceptance function reliably ranks images.The ranked samples are grouped by acceptance function value as best, random, and worst.Source: §C.5 · p. 22, Figure C.8 · p. 22

C.6 VAE with resampled prior and 2D latent space

• KL(q(z)∥p(z)) ≈0.4 for the regular VAE, indicating mismatch between its aggregate posterior and fixed standard Normal prior.Source: §C.6 · p. 22, Figure C.9 · p. 23
• KL(q(z)∥p(z)) ≈0.15 for the jointly trained high-capacity LARS prior, which matches the aggregate posterior very well.The aggregate posterior is more spread out than in the regular VAE; the acceptance function slices out parts of the prior and redistributes weight toward the sides.Source: §C.6 · p. 22, Figure C.10 · p. 23
• KL(q(z)∥p(z)) ≈0.20 with the low-capacity acceptance function, whose latent-space structure is similar but more coarsely partitioned than the high-capacity case.The authors attribute the slightly larger final KL to the lower flexibility of the acceptance function.Source: §C.6 · p. 22, Figure C.11 · p. 24
• Post-hoc LARS on the pretrained regular VAE fits its unchanged aggregate posterior very well, even slightly better than the jointly trained case.Only the acceptance function is learned; the encoder and decoder remain fixed, and the aggregate posterior is identical to that of the regular VAE.Source: §C.6 · pp. 22–23, Figure C.12 · p. 24
• Figure C.13 shows encoded-mean scatter plots for the regular and large- and small-network LARS VAEs, with colors indicating MNIST classes.Source: §C.6 · p. 23, Figure C.13 · p. 24