Amortised Inference through One-Step Implicit Sampling

Vincent Pauline; Kirill Tamogashev; Arran Carter; Sanghyeok Choi; Stefan Bauer; Esmeralda S. Whitammer

cs.LGOct 8, 2026

Amortised Inference through One-Step Implicit Sampling

Authors: Vincent Pauline, Kirill Tamogashev, Arran Carter, Sanghyeok Choi, Stefan Bauer, Esmeralda S. Whitammer

Organizations: Technical University of Munich · Munich Center for Machine Learning · University of Edinburgh · CIFAR Fellow

Abstract

Efficient sampling from Bayesian posterior distributions in high dimensions remains a hard problem, despite advances in theory and algorithms. Multiple steps are usually required to draw unbiased samples, for example, iterations of an MCMC algorithm or integration steps in a diffusion-based amortised sampler. However, scalable one-step amortised inference remains a challenging open problem. In this work, we build upon recent progress in data-driven generative modelling to offer a possible solution for one-step amortised sampling. Our one-step implicit sampler (OSIS) uses short-time MCMC dynamics to define a local target for the evolution of an implicitly modelled distribution. In the asymptotic limit, the modelled distribution is a fixed point of the MCMC kernel and thus samples from the target. We analyse the convergence properties of OSIS and show connections with other generative modelling and sampling algorithms. We study various design choices and show that OSIS performs well compared to multi-step amortised samplers and non-amortised MCMC baselines while being far more efficient at sampling time.

Figures & tables

Panel of Figure 1 — Figure 1: OSIS evolves an implicit modelled distribution $\rho_{\theta}$ towards a target $\pi$ available through its unnormalised density (green contours). $G_{\theta}$ maps samples $z$ the source distribution $q$ to samples $x=G_{\theta}(z)$ from the modelled distribution $\rho_{\theta}\coloneqq(G_{\theta})_{\#}q$ (blue KDE contours). This pushforward distribution is evolved by a $\pi$ -invariant MCMC kernel $K$ , giving $\rho_{h}\coloneqq K\rho_{{\theta}}$ . The parameters $\theta$ are updated to minimise an empirical divergence between $\rho_{\theta}$ and $\rho_{h}$ . Upon convergence, $\rho_{\theta}\approx\pi$ .

Panel of Figure 2 — Figure 2: One idealised OSIS iteration. KL and Wasserstein distance share the minimiser $\pi$ . However, KL decreases along the Langevin path (red), whereas Wasserstein decreases along the geodesic (grey). OSIS takes a step along the Wasserstein gradient of $\frac{1}{2}\mathcal{W}_{2}^{2}(\cdot,\rho_{h})$ towards $\rho_{k+1}$ (blue).

Target density $\rightarrow$		40GMM ( $d=2$ )		40GMM ( $d=10$ )		Funnel ( $d=10$ )		ManyWell ( $d=32$ )		ManyWell ( $d=64$ )
Method $\downarrow$ Metric $\rightarrow$	NN / $\mathcal{E}$	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )
Reference		$7.68{\scriptstyle\pm 3.22}$	$7.17{\scriptstyle\pm 1.51}$	$202.81{\scriptstyle\pm 16.77}$	$7.98{\scriptstyle\pm 0.78}$	$87.60{\scriptstyle\pm 1.20}$	$9.22{\scriptstyle\pm 1.54}$	$23.61{\scriptstyle\pm 0.07}$	$8.75{\scriptstyle\pm 0.72}$	$78.36{\scriptstyle\pm 0.11}$	$9.58{\scriptstyle\pm 0.65}$
OSIS	$1$ / $0$	$103.26{\scriptstyle\pm 41.41}$	$64.37{\scriptstyle\pm 13.38}$	$3086.96{\scriptstyle\pm 439.59}$	$139.26{\scriptstyle\pm 36.09}$	$93.99{\scriptstyle\pm 3.33}$	$28.06{\scriptstyle\pm 3.06}$	$26.79{\scriptstyle\pm 0.41}$	$55.92{\scriptstyle\pm 6.00}$	$75.88{\scriptstyle\pm 0.34}$	$91.05{\scriptstyle\pm 3.98}$
FAB (flow only)	$1$ / $0$	$296.01{\scriptstyle\pm 333.93}$	$81.05{\scriptstyle\pm 38.73}$	$6217.94{\scriptstyle\pm 472.86}$	$382.64{\scriptstyle\pm 44.65}$	$583.48{\scriptstyle\pm 233.71}$	$12.94{\scriptstyle\pm 2.34}$	$27.10{\scriptstyle\pm 0.46}$	$47.95{\scriptstyle\pm 4.56}$	$125.58{\scriptstyle\pm 0.04}$	$401.28{\scriptstyle\pm 0.64}$
FAB	$1$ / $120$	$72.03{\scriptstyle\pm 19.29}$	$46.42{\scriptstyle\pm 6.32}$	$6257.04{\scriptstyle\pm 455.10}$	$380.98{\scriptstyle\pm 45.19}$	$582.97{\scriptstyle\pm 233.17}$	$11.87{\scriptstyle\pm 1.74}$	$25.52{\scriptstyle\pm 0.30}$	$36.63{\scriptstyle\pm 5.56}$	$132.55{\scriptstyle\pm 0.48}$	$299.11{\scriptstyle\pm 1.51}$
SMC	$0$ / $3825$	$14.67{\scriptstyle\pm 1.79}$	$16.85{\scriptstyle\pm 2.11}$	$4362.81{\scriptstyle\pm 139.06}$	$208.62{\scriptstyle\pm 11.73}$	$149.56{\scriptstyle\pm 1.37}$	$125.72{\scriptstyle\pm 8.99}$	$26.04{\scriptstyle\pm 0.16}$	$22.48{\scriptstyle\pm 1.10}$	$85.76{\scriptstyle\pm 0.29}$	$42.63{\scriptstyle\pm 1.36}$

Table 1: Sampler comparison on GMM, Funnel, and ManyWell targets. Entries are mean

\pm

standard deviation over five seeds. Second column shows number of NN and energy evaluations per sample

Target density $\rightarrow$		NICE-MNIST ( $d=196$ )		NICE-FMNIST ( $d=784$ )
Method $\downarrow$ Metric $\rightarrow$	NN / $\mathcal{E}$	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )
Reference		$129.56{\scriptstyle\pm 0.22}$	$8.75{\scriptstyle\pm 0.08}$	$549.42{\scriptstyle\pm 1.35}$	$9.24{\scriptstyle\pm 1.09}$
OSIS	$1$ / $0$	$189.59{\scriptstyle\pm 20.34}$	$184.55{\scriptstyle\pm 68.71}$	$1743.96{\scriptstyle\pm 180.58}$	$382.13{\scriptstyle\pm 59.92}$
FAB (flow only)	$1$ / $0$	$162.20{\scriptstyle\pm 1.91}$	$57.94{\scriptstyle\pm 11.09}$	$1768.75{\scriptstyle\pm 335.78}$	$430.25{\scriptstyle\pm 72.05}$
FAB	$120$ / $120$	$143.83{\scriptstyle\pm 3.10}$	$80.83{\scriptstyle\pm 14.39}$	$1746.47{\scriptstyle\pm 335.35}$	$438.79{\scriptstyle\pm 71.94}$
SMC	$0$ / $3825$	$268.32{\scriptstyle\pm 4.16}$	$246.17{\scriptstyle\pm 13.54}$	$2856.36{\scriptstyle\pm 15.26}$	$629.18{\scriptstyle\pm 3.24}$

Table 2: Sampler comparison on NICE targets. Conventions follow Table 1 .

Figure 3: Sampling time comparison on ManyWell (median of

250

runs). DIS is trained for different numbers of denoising steps. OSIS achieves

\mathcal{W}_{2}^{2}

comparable to the baselines while sampling up to

1000\times

faster.

Panel of Figure 4 — Figure 4: Visualisation of samples generated by OSIS and the baselines for 40GMM ( $d=2$ ), Funnel ( $d=10$ ), ManyWell ( $d=32$ ), MNIST ( $d=196$ ). Funnel and ManyWell are projected onto $\mathbb{R}^{2}$ .

Panel of Figure 5 — Figure 5: Uncurated samples obtained using outsourced sampling ( Section 6.2 ). The samplers correspond to the ImageReward-based target conditioned on the given prompts. More examples in Section E.6 .

Target density $\rightarrow$	40GMM ( $d=2$ )		ManyWell ( $d=32$ )
Kernel $\downarrow$ Metric $\rightarrow$	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )
Reference	$7.68{\scriptstyle\pm 3.22}$	$7.17{\scriptstyle\pm 1.51}$	$23.61{\scriptstyle\pm 0.07}$	$8.75{\scriptstyle\pm 0.72}$
RWMH	$90.29{\scriptstyle\pm 1.08}$	$49.88{\scriptstyle\pm 0.56}$	$30.08{\scriptstyle\pm 0.62}$	$112.31{\scriptstyle\pm 6.76}$
ULA	$85.42{\scriptstyle\pm 25.65}$	$60.98{\scriptstyle\pm 13.75}$	$25.07{\scriptstyle\pm 0.29}$	$39.54{\scriptstyle\pm 6.08}$
MALA	$103.26{\scriptstyle\pm 41.41}$	$64.37{\scriptstyle\pm 13.38}$	$26.79{\scriptstyle\pm 0.41}$	$55.92{\scriptstyle\pm 6.00}$
BAOAB	$40.24{\scriptstyle\pm 16.81}$	$41.66{\scriptstyle\pm 15.91}$	$43.88{\scriptstyle\pm 3.70}$	$227.46{\scriptstyle\pm 52.12}$

Table 3: OSIS kernel ablation with

M=1

inner iterations ( Algorithm 3 ); details in Appendix D .

Algorithm $\downarrow$ Metric $\rightarrow$	$\mathbb{E}[\log r(x)]$ ( $\uparrow$ )	FID $(\downarrow)$
Rejection sampling	$2.164$	$2.345$
MALA	$2.127$	$3.070$
OSIS	$1.387$	$8.965$
OSIS + refinement	$1.854$	$7.510$
TB $+$ MALA	$0.992$	$4.267$

Table 4: Comparison on the outsourced sampling problem.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 10

Method	Parameter	GMM ( $d=10$ )	ManyWell ( $d=32,64$ )
FAB / PIS / DIS	Learning rate	$\{10^{-5},10^{-4},5\cdot 10^{-4},10^{-3}\}$
FAB	Base scale	$\{20,40,80\}$	$\{1,2,4\}$
DIS	Initial std.	$\{20,40,80\}$	$\{1,2,4\}$
PIS	Max. diffusion amplitude	$\{80,160,320\}$	$\{2,4,8\}$
SMC	Initial std.	$\{20,40,80\}$	$\{1,2,4\}$
	HMC step size	$\{0.1,0.2\}$	$\{0.05,0.1\}$

Appendix

Table 5: Baseline hyperparameter search grids.

Target	Loss	Model	lr	$\tau_{\mathrm{start}}$	target accept.
40GMM ( $d=2$ )	SWD	ResNet (256, 3)	$3\cdot 10^{-3}$	128	0.3
40GMM ( $d=10$ )	SWD	ResNet (256, 3)	$3\cdot 10^{-3}$	256	0.3
ManyWell ( $d=32$ )	SWD	ResNet (256, 3)	$3\cdot 10^{-3}$	256	0.90
ManyWell ( $d=64$ )	SWD	ResNet (256, 3)	$3\cdot 10^{-3}$	16	0.574
Funnel ( $d=10$ )	SWD	ResNet (128, 2)	$3\cdot 10^{-5}$	1	0.90
NICE-MNIST ( $d=196$ )	SWD	UNet (16, 4)	$1\cdot 10^{-3}$	32	0.574

Appendix

Table 6: OSIS hyperparameters. ResNet (hidden, blocks) and UNet (base channels, attention heads). The temperature schedule is as defined in Appendix D .

Kernel	Learning rate d	Initial g temperature d	Initial proposal standard g deviation	Target p acceptance d g rate d	ULA g drift step	Integration step size	Friction	Leapfrog steps d	SVGD g update step	RBF kernel g variance
RWMH	$10^{-3}$	4	1	0.5	—	—	—	—	—	—
MALA	$10^{-3}$	16	—	0.90	—	—	—	—	—	—
ULA	$10^{-3}$	16	—	—	0.008	—	—	—	—	—
BAOAB	$1.9\cdot 10^{-5}$	2	—	—	—	0.0025	2	—	—	—
OBABO	$10^{-3}$	16	—	—	—	0.006	2	—	—	—
HMC	$1.7\cdot 10^{-3}$	16	—	—	—	0.019	—	10	—	—

Appendix

Table 7: Selected hyperparameters for ManyWell (

d=32

Target $\rightarrow$	40GMM		ManyWell		Funnel	NICE
Method $\downarrow$	$d=2$	$d=10$	$d=32$	$d=64$	$d=10$	MNIST	FMNIST
SMC	—	—	—	—	—	—	—
FAB	70.7	83.1	117.2	166.9	83.1	371.8	1284.4
PIS / DIS	37.8	39.3	43.6	49.8	39.3	75.4	189.5
OSIS	198.7	202.8	214.0	230.5	35.7	227.9	689.9

Appendix

Table 8: Number of trainable parameters in the benchmark study (

\times 10^{3}

Target $\rightarrow$	40GMM ( $d=2$ )		40GMM ( $d=10$ )		Funnel ( $d=10$ )		ManyWell ( $d=32$ )		ManyWell ( $d=64$ )
Method $\downarrow$ Metric $\rightarrow$	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )
Reference	$7.68{\scriptstyle\pm 3.22}$	$7.17{\scriptstyle\pm 1.51}$	$202.81{\scriptstyle\pm 16.77}$	$7.98{\scriptstyle\pm 0.78}$	$87.60{\scriptstyle\pm 1.20}$	$9.22{\scriptstyle\pm 1.54}$	$23.61{\scriptstyle\pm 0.07}$	$8.75{\scriptstyle\pm 0.72}$	$78.36{\scriptstyle\pm 0.11}$	$9.58{\scriptstyle\pm 0.65}$
OSIS (OBABO, OT)	$81.85{\scriptstyle\pm 31.72}$	$58.06{\scriptstyle\pm 15.95}$	$2359.76{\scriptstyle\pm 535.07}$	$123.10{\scriptstyle\pm 39.78}$	$189.86{\scriptstyle\pm 0.33}$	$246.80{\scriptstyle\pm 0.72}$	$35.54{\scriptstyle\pm 0.01}$	$255.33{\scriptstyle\pm 0.34}$	$78.70{\scriptstyle\pm 0.04}$	$274.45{\scriptstyle\pm 0.43}$
OSIS (OBABO, SWD)	$848.79{\scriptstyle\pm 388.69}$	$378.69{\scriptstyle\pm 132.12}$	$7465.66{\scriptstyle\pm 7.91}$	$477.83{\scriptstyle\pm 0.51}$	$99.37{\scriptstyle\pm 2.09}$	$57.69{\scriptstyle\pm 8.07}$	$28.77{\scriptstyle\pm 0.53}$	$85.55{\scriptstyle\pm 8.86}$	$91.24{\scriptstyle\pm 2.01}$	$170.13{\scriptstyle\pm 9.52}$
OSIS (MALA, OT)	$111.24{\scriptstyle\pm 94.65}$	$83.77{\scriptstyle\pm 39.73}$	$2913.22{\scriptstyle\pm 489.84}$	$140.34{\scriptstyle\pm 31.18}$	$243.15{\scriptstyle\pm 3.27}$	$474.04{\scriptstyle\pm 16.40}$	$36.60{\scriptstyle\pm 0.02}$	$276.63{\scriptstyle\pm 0.29}$	$81.37{\scriptstyle\pm 0.10}$	$293.66{\scriptstyle\pm 0.61}$
OSIS (MALA, SWD)	$553.89{\scriptstyle\pm 305.14}$	$277.34{\scriptstyle\pm 116.75}$	$9854.75{\scriptstyle\pm 1328.11}$	$482.27{\scriptstyle\pm 132.16}$	$89.41{\scriptstyle\pm 0.53}$	$45.32{\scriptstyle\pm 3.21}$	$25.51{\scriptstyle\pm 0.43}$	$48.57{\scriptstyle\pm 7.55}$	$89.93{\scriptstyle\pm 14.74}$	$159.27{\scriptstyle\pm 20.76}$

Appendix

Table 9: Sampler comparison on GMM, Funnel, and ManyWell targets. OSIS rows specify the training kernel and loss. This table was computed using

M=1

inner iteration (see Algorithm 3 ).

Target $\rightarrow$	NICE-MNIST ( $d=196$ )		NICE-FMNIST ( $d=784$ )
Method $\downarrow$ Metric $\rightarrow$	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )
Reference	$129.56{\scriptstyle\pm 0.22}$	$8.75{\scriptstyle\pm 0.08}$	$549.42{\scriptstyle\pm 1.35}$	$9.24{\scriptstyle\pm 1.09}$
OSIS (OBABO, OT)	$270.02{\scriptstyle\pm 4.60}$	$428.07{\scriptstyle\pm 10.29}$	$1631.93{\scriptstyle\pm 63.53}$	$412.35{\scriptstyle\pm 14.51}$
OSIS (OBABO, SWD)	$183.18{\scriptstyle\pm 5.45}$	$122.34{\scriptstyle\pm 27.80}$	$2269.09{\scriptstyle\pm 475.05}$	$410.57{\scriptstyle\pm 89.89}$
OSIS (MALA, OT)	$277.40{\scriptstyle\pm 8.45}$	$462.69{\scriptstyle\pm 18.73}$	$1776.79{\scriptstyle\pm 82.07}$	$450.47{\scriptstyle\pm 21.16}$
OSIS (MALA, SWD)	$209.74{\scriptstyle\pm 39.21}$	$248.96{\scriptstyle\pm 125.42}$	$2506.97{\scriptstyle\pm 29.95}$	$507.93{\scriptstyle\pm 31.28}$

Appendix

Table 10: Sampler comparison on high-dimensional targets. OSIS rows specify the training kernel and loss. This table was computed using

M=1

inner iteration (see Algorithm 3 ).

Target $\rightarrow$	40GMM ( $d=2$ )		40GMM ( $d=10$ )		Funnel ( $d=10$ )		ManyWell ( $d=32$ )		ManyWell ( $d=64$ )
Method $\downarrow$ Metric $\rightarrow$	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )
Reference	$7.68{\scriptstyle\pm 3.22}$	$7.17{\scriptstyle\pm 1.51}$	$202.81{\scriptstyle\pm 16.77}$	$7.98{\scriptstyle\pm 0.78}$	$87.60{\scriptstyle\pm 1.20}$	$9.22{\scriptstyle\pm 1.54}$	$23.61{\scriptstyle\pm 0.07}$	$8.75{\scriptstyle\pm 0.72}$	$78.36{\scriptstyle\pm 0.11}$	$9.58{\scriptstyle\pm 0.65}$
OSIS (MALA, OT)	$111.24{\scriptstyle\pm 94.65}$	$83.77{\scriptstyle\pm 39.73}$	$2913.22{\scriptstyle\pm 489.84}$	$140.34{\scriptstyle\pm 31.18}$	$243.15{\scriptstyle\pm 3.27}$	$474.04{\scriptstyle\pm 16.40}$	$36.60{\scriptstyle\pm 0.02}$	$276.63{\scriptstyle\pm 0.29}$	$81.37{\scriptstyle\pm 0.10}$	$293.66{\scriptstyle\pm 0.61}$
OSIS (MALA, SWD)	$553.89{\scriptstyle\pm 305.14}$	$277.34{\scriptstyle\pm 116.75}$	$9854.75{\scriptstyle\pm 1328.11}$	$482.27{\scriptstyle\pm 132.16}$	$89.41{\scriptstyle\pm 0.53}$	$45.32{\scriptstyle\pm 3.21}$	$25.51{\scriptstyle\pm 0.43}$	$48.57{\scriptstyle\pm 7.55}$	$89.93{\scriptstyle\pm 14.74}$	$159.27{\scriptstyle\pm 20.76}$
MALA only	$82.82{\scriptstyle\pm 2.28}$	$52.04{\scriptstyle\pm 1.67}$	$200.14{\scriptstyle\pm 12.26}$	$8.58{\scriptstyle\pm 0.62}$	$90.71{\scriptstyle\pm 3.41}$	$22.47{\scriptstyle\pm 3.52}$	$36.58{\scriptstyle\pm 0.05}$	$149.82{\scriptstyle\pm 0.74}$	$104.71{\scriptstyle\pm 0.11}$	$166.18{\scriptstyle\pm 0.86}$

Appendix

Table 11: Comparison of OSIS +MALA with long-run MALA only. OSIS rows specify the training kernel and loss. MALA-only runs start from Gaussian noise. We compare one set of MALA samples against five reference sets and report the average and standard deviation. This table was computed using

M=1

inner iteration (see Algorithm 3 ).

Target $\rightarrow$	40GMM ( $d=2$ ) + OT		ManyWell ( $d=32$ ) + SWD
Method $\downarrow$ Metric $\rightarrow$	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )	$\mathcal{W}_{2}^{2}$	MMD ( $\times 10^{-3}$ )
Reference	$7.68{\scriptstyle\pm 3.22}$	$7.17{\scriptstyle\pm 1.51}$	$23.61{\scriptstyle\pm 0.07}$	$8.75{\scriptstyle\pm 0.72}$
Discrepancy	$81.85{\scriptstyle\pm 31.72}$	$58.06{\scriptstyle\pm 15.95}$	$28.77{\scriptstyle\pm 0.53}$	$85.55{\scriptstyle\pm 8.86}$
Drift	$95.93{\scriptstyle\pm 46.34}$	$77.97{\scriptstyle\pm 35.87}$	$27.11{\scriptstyle\pm 1.29}$	$69.34{\scriptstyle\pm 13.05}$

Appendix

Table 12: Comparison of discrepancy and drift objectives for OSIS trained with the OBABO refinement kernel. Results are mean

\pm

standard deviation over five training seeds, after averaging five sampling-seed evaluations per checkpoint. This table was computed using

M=1

inner iteration (see Algorithm 3 ).

MCMC steps (default: $20$ )
Configuration	$\mathcal{W}_{2}^{2}$ ( $\downarrow$ )	MMD ( $\times 10^{-3}$ )
Default	$25.51{\scriptstyle\pm 0.43}$	$48.57{\scriptstyle\pm 7.55}$
$1$	$28.40{\scriptstyle\pm 1.19}$	$77.78{\scriptstyle\pm 20.34}$
$5$	$30.17{\scriptstyle\pm 1.32}$	$109.10{\scriptstyle\pm 29.69}$
$10$	$27.98{\scriptstyle\pm 1.87}$	$85.43{\scriptstyle\pm 41.68}$
$40$	$25.30{\scriptstyle\pm 0.55}$	$51.75{\scriptstyle\pm 12.59}$

Appendix

Table 13: Ablation on ManyWell (

d=32

). The default OSIS training configuration uses 20-step MALA, batch size

4096

, SWD-

2048

, and

50,000

training updates. Results are mean

\pm

standard deviation over five training seeds, after averaging five sampling-seed evaluations per checkpoint. This table was computed using

M=1

inner iteration (see Algorithm 3 ).

Panel of Figure 6 — Figure 6: Samples obtained using outsourced sampling. Samples correspond to the classifier-based target with the reward $r(x)=0.5p(0|x)+0.5p(1|x)$ . From left to right, samples are obtained with OSIS , OSIS with 20-step MALA refinement, and rejection sampling.

Panel of Figure 7 — Figure 7: Samples obtained using outsourced sampling. Samples correspond to the classifier-based target with the reward $r(x)=1p(2\mid x)+0.5p(5\mid x)+2p(7\mid x)$ . From left to right, samples are obtained with OSIS , OSIS with 20-step MALA refinement, and rejection sampling.

Panel of Figure 8 — Figure 8: Comparison of samples obtained with OSIS , OSIS with 20-step MALA refinement, and rejection sampling on outsourced sampling task with ImageReward-based target. The following 4 prompts are considered (from to to bottom): a young blond girl , smiling toddler , a man with a beard , person wearing glasses .

Panel of Figure 9 — Figure 9: Visualisation of samples from OSIS , rejection sampling, TB $+$ MALA ( Sendera et al., 2024 ) and MALA. All algorithms, except rejection sampling, are matched in terms of the number of log-density evaluations.

Explore similar work

Oct 12, 2026

Amortised adaptive importance sampling

Monte Carlo sampling algorithms aim to produce samples from a target distribution with a given unnormalised probability density function. The high cost of running Monte Carlo algorithms, with many simulation steps or particles needed to obtain a good empirical approximation to the target, can be reduced by introducing learnable components into the sampler -- a form of amortisation. Once these components are trained, amortised samplers can produce high-quality samples with less computation at sampling time than training-free algorithms would require. We propose an amortised sampler built upon the framework of adaptive importance sampling -- a hierarchical importance sampling procedure where the proposal is iteratively adapted. Our amortised adaptive importance sampling algorithm fits a lightweight neural network that adapts the proposal using information obtained over the course of sampling. By training our sampler over many different targets taken from the same family of distributions, we obtain a sampler that can adapt to a specific target at sampling time without a priori access to the target's parameters. We show that our procedure can efficiently sample from targets within the family the sampler was trained on, giving improvements over target-specific amortised samplers and training-free adaptive importance sampling algorithms. We also explore the potential for generalisation, where the trained sampler is used to sample from targets outwith its training family.

Arran Carter, Víctor Elvira, Esmeralda S. Whitammer

University of Edinburgh

Importance Sampling Adaptive Sampling Amortized Inference

Ask PDF

Sep 30, 2026

Amortized Bayesian Inference on Multilevel Models of Arbitrary Structure

We develop a general method for amortized Bayesian inference on multilevel models of arbitrary structure. Given a generative model specified as a directed acyclic graph, our method automatically derives valid factorizations of the joint posterior and matching neural network architectures. The key steps, graph expansion and graph inversion, yield an inverse graph that determines how inference networks are stacked and conditioned, producing factorizations that amortize over the number of groups and the number of observations within each group. Unlike approaches that simplify the dependency structure to speed up learning or inference, our method preserves all conditional independence and exchangeability assumptions of the generative model. Across three case studies, it closely matches gold-standard samplers on models with more than 6,500 parameters while reducing inference to a near-instant forward pass once trained.

Daniel Habermann, Andreas Bulling, Stefan T. Radev +1

Department of Statistics TU Dortmund University, Germany · Institute for Visualisation and Interactive Systems University of Stuttgart, Germany · Department of Cognitive Science Rensselaer Polytechnic Institute, USA

Amortized Inference Hierarchical Bayesian Modeling Probabilistic Graphical Models

Ask PDF

Aug 9, 2026

Path-dependent Discrete Amortized Inference

We consider the problem of sampling compositional and discrete objects from a given unnormalized posterior distribution. Notably, recent studies have shown that this problem can be efficiently solved by learning a deterministic Markov Decision Process (MDP) that progressively builds each object in proportion to the posterior. In this work, however, we demonstrate that the Markovian assumption can both hamper signal propagation during training and catastrophically reduce the learned sampler's expressivity due to state aliasing. To address these issues, we propose lifting the MDP with a learnable latent dynamical system that allows the underlying policy to depend on the entire past trajectory---and not only on the current state. In view of this, we refer to the resulting method as path-dependent discrete amortized inference. Importantly, we provably extend existing learning algorithms for discrete amortized samplers to our setting. In experiments on standard benchmark problems, we also show that our approach often leads to faster learning convergence and improved state space exploration relatively to prior techniques.

Tiago da Silva, Esmeralda S. Whitammer, Salem Lahlou

1MBZUAI · University of Edinburgh · 3CIFAR Fellow

Markov Decision Processes Latent Dynamics Modeling Amortized Inference Posterior Sampling

Ask PDF