Amortised Inference through One-Step Implicit Sampling

• The sampling problem includes distributions specified by an energy known only up to a normalising constant, while conventional Monte Carlo methods often require long chains.Source: §1 · p. 1
• One-step amortised sampling remains challenging, and one-step normalising flows require invertibility and efficient Jacobian computation.Density-only generative modelling also continues to lag in scalability.Source: §1 · pp. 1–2
• OSIS uses a target-invariant MCMC kernel to define a drifting field and trains its sampler to be invariant to that refinement.Source: §1 · p. 2
• The paper proves OSIS convergence under asymptotic conditions and constraints on the target.The related analysis describes convergence assumptions as similar to those in the MCMC literature.Source: §1 · p. 3
• The empirical study examines how choices such as the MCMC kernel and empirical divergence affect OSIS performance.Benchmarks span synthetic targets to posterior sampling under large-scale image-model priors.Source: §1 · p. 3
• OSIS performs on par with multi-step baselines, is more efficient at sampling time, and remains scalable.Source: §1 · p. 3

2.1 PRELIMINARIES

• A generator maps samples from a simple source distribution to a pushforward distribution that approximates the target, even when its density is intractable.Source: §2.1 · p. 3
• A target-invariant MCMC kernel has the target as its stationary distribution, and common choices include Langevin-based methods and HMC.Source: §2.1 · p. 3

2.2 TRAINING OBJECTIVES FOR ONE-STEP SAMPLING

• OSIS trains the generator so its distribution matches its refinement under a target-invariant kernel, using a divergence-based loss.The kernel and differentiable divergence are design choices that should be efficient to compute.Source: §2.2 · pp. 3–4

• Earlier implicit-sampler work includes MCMC-trained approaches, such as amortising Stein variational gradient descent and fitting an invertible map to short-chain samples.Source: §3 · p. 4
• Compared with Li et al. (2017), OSIS avoids adversarial loss, supports varied MCMC kernels and divergences, and links convergence to the underlying MCMC kernel.The paper also identifies a connection to Wasserstein gradient flows not previously established in the sampling context.Source: §3 · p. 4
• Like drifting models, OSIS defines a local target for evolving an implicit distribution, but OSIS addresses energy-based targets rather than empirical data distributions.Drifting models approximate the vector field; OSIS uses it to define a local optimisation target.Source: §3 · p. 4
• Cao et al. (2026) adapt drifting models for sampling, but their approaches rely on Monte Carlo simulation or second-order energy derivatives that scale poorly with dimension.OSIS instead defines drift implicitly through a divergence between distributions.Source: §3 · p. 4
• OSIS is related to computational Wasserstein gradient flows, while existing neural approximations of those flows rely on unbiased target samples.The paper situates this connection alongside prior work on scalable Wasserstein-gradient-flow methods.Source: §3 · pp. 4–5

4 THEORETICAL ANALYSIS OF OSIS

• The analysis shows that OSIS inherits Langevin dynamics’ convergence properties because Langevin density evolution is the Wasserstein gradient flow of KL divergence.The analysis assumes idealised optimization in Wasserstein space and exact divergence estimation and Monte Carlo updates.Source: §4 · p. 5

4.1 SETTING

• Each idealised OSIS step refines the current distribution with a π-invariant Langevin kernel, then takes an η-weighted Wasserstein gradient step toward that refinement.The step uses the optimal transport map between the current distribution and its kernel-refined distribution.Source: §4.1 · pp. 5–6

4.2 MAIN RESULTS

• OSIS converges geometrically for strongly log-concave targets, and also under an LSI and bounded negative curvature with fixed-step conditions.For the latter case, the stated condition includes κh(1−η) < 2.Source: §4.2 · p. 6

• In practice, OSIS uses discrete-time MCMC and empirical approximations to distributions and their divergence rather than the idealised procedure.The theory’s exact continuous-time Langevin and infinite-population setting cannot be implemented directly.Source: §5 · p. 6
• The practical refinement kernel is MALA, although other MCMC kernels are considered; the theory instead analyzes exact Langevin.Source: §5 · p. 6
• Table 1 compares samplers on GMM, Funnel, and ManyWell targets, reporting five-seed mean ± standard deviation and per-sample NN and energy evaluations.Source: Table 1 · p. 7
• OSIS uses sliced Wasserstein distance for synthetic targets and minibatch optimal transport with squared-Euclidean cost for Bayesian posteriors.The minibatch OT choice converges to exact 2-Wasserstein in the limit.Source: §5 · pp. 6–7
• Training may anneal the kernel-refinement target from a tempered distribution, and uses multiple gradient steps after each MCMC refinement for stability.The starting temperature and schedule balance early exploration with learning the final target.Source: §5 · p. 7

6.1 SYNTHETIC BENCHMARKS

• OSIS is not consistently outperformed by MCMC, normalising-flow, or diffusion baselines on synthetic targets, despite using one neural function evaluation and no energy evaluations per sample.Off-policy diffusion samplers are strongest overall, but require many sampling steps to perform well.Source: §6.1 · p. 7
• On NICE targets, OSIS is comparable to baselines while sampling up to 1000× faster.Table 2 reports sampler comparisons on NICE targets.Source: Table 2 · p. 7
• Visualized samples show comparable quality with significantly faster OSIS inference across 40GMM, Funnel, ManyWell, and MNIST.Funnel and ManyWell samples are projected onto R2.Source: §6.1 · p. 8, Figure 4 · p. 8
• On ManyWell, OSIS consistently outperforms DIS at strictly lower sampling time, even as DIS performance improves with more denoising steps.Figure 3 compares sampling time per batch of 2,000 samples; DIS uses 4, 16, 64, or 128 denoising steps.Source: Figure 3 · p. 7, §6.1 · p. 8

6.2 BAYESIAN POSTERIOR SAMPLING UNDER DEEP GENERATIVE MODEL PRIORS

• For Bayesian posterior sampling, the experiments train OSIS on latent-space targets induced by a VAE classifier reward and an R3GAN ImageReward conditioned on prompts.The latent-space posterior has density proportional to q(z)r(g(z)), and its pushforward by g gives the posterior in image space.Source: §6.2 · p. 9
• OSIS performs similarly to TB+MALA on the ImageReward-conditioned FFHQ-64 task, with higher mean reward and lower FID using one function evaluation per sample.MCMC refinement further improves OSIS.Source: §6.2 · p. 9

7 CONCLUSION

• OSIS matches the generator-induced distribution to its MCMC refinement and performs well against strong samplers that require many neural-network and energy evaluations.Practical techniques make the refinement and divergence optimisation feasible.Source: §7 · p. 9
• Highly multimodal targets remain an immediate sampling-quality challenge; tempering alone lacks the geometric mixing from noise that has helped diffusion samplers.The authors also call for theory covering temperature annealing and inexact integration within refinement kernels.Source: §7 · p. 9

A.1 VECTOR-FIELD-BASED LOSS

• The vector-field alternative estimates a drift field from model samples and trains the generator by regression toward displaced samples, with targets fixed during differentiation.For OT-based objectives, these targets still require the corresponding transport information.Source: §A.1 · p. 16

A.2 FULL ALGORITHM

• The practical OSIS algorithm adds tempering, approximate chains and divergences, a replay buffer, and inner iterations to the idealized version.The paper focuses on Algorithms 1 and 3; Algorithms 2 and 4 are included for completeness.Source: §A.2 · p. 16

B.1 ADDITIONAL NOTATION AND TEMPERATURE SCALING

• The supplementary notation assumes finite KL divergence and Fisher information, and notes that 2-Wasserstein distance satisfies the triangle inequality.The appendix also presents Kantorovich, Monge, and dynamic formulations of 2-Wasserstein distance.Source: §B.1 · p. 18

B.2 PROOFS

• The KL convexity proof separates the energy term, which is λ-strongly convex, from a second term that is convex.The appendix also states that setting λ = 0 recovers geodesic convexity.Source: §B.2 · p. 19
• Under λ-strong log-concavity, one OSIS step contracts KL with factor 1 − η_k(1 − e^(−2λh_k)).The proof also establishes λ-strong convexity of KL along the interpolation path.Source: §B.2 · pp. 19–21
• Under an LSI and bounded negative curvature, OSIS converges in both KL and W2; the fixed-step corollary requires κh(1 − η) < 2.The convergence result assumes positive step sizes and η_k in (0, 1].Source: §B.2 · pp. 22–23

C.1 CHOICE OF MCMC KERNEL

• The kernel discussion distinguishes MALA, whose proposal is Metropolis-adjusted, from ULA, which need not preserve the target at finite step size.Neither is the exact Langevin semigroup transition used in the theoretical analysis.Source: §C.1 · pp. 23–24
• Higher-order splitting discretises underdamped Langevin by composing exact updates for transport, force-driven velocity changes, and friction with Gaussian noise.The appendix notes that splitting schemes can reduce finite-step sampling bias.Source: §C.1 · p. 24

C.2 CHOICE OF DIVERGENCE

• The divergence options include MMD, Wasserstein and sliced Wasserstein distances, and Sinkhorn divergence, which uses entropic regularisation for smoother optimisation.For empirical measures with squared-Euclidean cost, debiased Sinkhorn approaches W2 as ε decreases to zero.Source: §C.2 · p. 25

D.1 SYNTHETIC TARGETS

• Synthetic-target evaluations use five training seeds and five independent 10,000-sample batches per method at the selected checkpoint.OSIS uses its final iteration; baselines use their highest-ELBO checkpoint.Source: §D.1 · p. 25
• The synthetic baselines include retrained FAB, PIS, and DIS, reference-implementation SMC, and TB+MALA and SubTB+SMC implementations.The reimplemented TB+MALA is reported to perform at least as well as, or better than, the original on synthetic targets.Source: §D.1 · p. 26
• The main synthetic benchmark trains OSIS with squared SWD-2048, 4,096-sample batches, 10,000 updates, 20 MCMC-refinement steps, and 50 inner gradient iterations.Source: §D.1 · p. 26

D.2 OUTSOURCED SAMPLING

• For MNIST outsourced sampling, the target uses a 16-dimensional VAE latent with Bernoulli likelihood and a CNN classifier reward.Source: §D.2 · p. 27
• For FFHQ-64, the target combines a frozen pretrained R3GAN generator, ImageReward, and a quadratic latent penalty.Reward and gradient passes use bfloat16.Source: §D.2 · p. 27
• Outsourced-sampling OSIS uses 1,200 outer and 50 inner steps, 20-step adaptive MALA refinement, and a cosine temperature schedule from 2.0 to 1.0.Each method is evaluated with 10,000 samples against a rejection-sampling reference.Source: §D.2 · p. 27

E.1 EXTENDED EVALUATION OF OSIS

• With OT, OBABO has lower mean squared W2 and MMD than MALA on every evaluated target, while SWD results are mixed.These comparisons use equal OSIS training-step counts, not necessarily equal computation; OBABO uses two energy-gradient evaluations per step.Source: §E.1 · p. 28

E.2 OSIS AGAINST ITS LONG RUN REFINEMENT

• Against equal-length MALA refinement, OSIS has higher discrepancies on 40GMM, but SWD-trained OSIS outperforms refinement on ManyWell.The authors describe effective amortisation of the refinement signal on specific targets.Source: §E.2 · p. 28, Table 11 · p. 29

E.3 DRIFT PARAMETERISATION OF OSIS

• Drift regression lowers both metric means on ManyWell (d = 32), but raises mean Sinkhorn divergence and MMD on 40GMM (d = 2), within reported variability.Source: §E.3 · p. 30, Table 12 · p. 30

E.4 ABLATIONS

• The ManyWell (d = 32) ablation uses 20-step MALA, batch size 4096, SWD-2048, and 50,000 training updates.Results are reported as means and standard deviations over five training seeds, averaging five sampling-seed evaluations per checkpoint.Source: Table 13 · p. 31, §E.4 · p. 31

E.5 ADDITIONAL SAMPLES FROM BAYESIAN POSTERIORS ON MNIST

• Figure 6 compares OSIS, 20-step MALA-refined OSIS, and rejection sampling for a classifier-based MNIST target with reward 0.5p(0|x) + 0.5p(1|x).Source: Figure 6 · p. 32
• Figure 7 shows samples obtained using outsourced sampling for a classifier-based target, with OSIS, OSIS with 20-step MALA refinement, and rejection sampling shown from left to right.The reward is r(x) = 1p(2|x) + 0.5p(5|x) + 2p(7|x).Source: Figure 7 · p. 32

E.6 ADDITIONAL SAMPLES FROM BAYESIAN POSTERIORS ON FFHQ-64

• Figure 8 compares samples from OSIS, OSIS with 20-step MALA refinement, and rejection sampling on an outsourced sampling task using an ImageReward-based target.The four prompts describe a young blond girl, a smiling toddler, a man with a beard, and a person wearing glasses.Source: Figure 8 · p. 33
• Figure 9 visualizes samples from OSIS, rejection sampling, TB+MALA, and MALA, matching log-density evaluation counts for all but rejection sampling.TB+MALA is attributed to Sendera et al. (2024).Source: Figure 9 · p. 34