Instance-Adaptive Parametrization for Amortized Variational Inference

Andrea Pollastro; Andrea Apicella; Francesco Isgrò; Roberto Prevete

cs.LGApr 8, 2026

Instance-Adaptive Parametrization for Amortized Variational Inference

Authors: Andrea Pollastro, Andrea Apicella, Francesco Isgrò, Roberto Prevete

Organizations: Department of Electrical Engineering and Information Technology, University of Naples Federico II, 80125, Naples, Italy · Department of Information Engineering, Electrical Engineering, and Applied Mathematics, University of Salerno, 84084, Fisciano (Salerno), Italy

Abstract

Variational autoencoders (VAEs) rely on amortized variational inference to enable efficient posterior approximation, but this efficiency comes at the cost of a shared parametrization, giving rise to the amortization gap. We propose the instance-adaptive variational autoencoder (IA-VAE), an amortized inference framework in which a hypernetwork generates input-dependent modulations of a shared encoder. This enables input-specific adaptation of the inference model while preserving the efficiency of a single forward pass. From a theoretical perspective, we show that the variational family induced by IA-VAE contains that of standard amortized inference, implying that IA-VAE cannot yield a worse optimal ELBO. By leveraging instance-specific parameter modulations, the proposed approach can achieve performance comparable to standard encoders with substantially fewer parameters, indicating a more efficient use of model capacity. Experiments on synthetic data, where the true posterior is known, show that IA-VAE yields more accurate posterior approximations and reduces the amortization gap. Similarly, on standard image benchmarks, IA-VAE consistently improves held-out ELBO over baseline VAEs, with statistically significant gains across multiple runs. These results suggest that increasing the flexibility of the inference parametrization through instance-adaptive modulation is an effective strategy for mitigating amortization-induced suboptimality in deep generative models.

Figures & tables

Panel of Figure 1 — Figure 1: Comparison between non-amortized and amortized variational inference. In non-amortized inference (left), the variational parameters are optimized independently for each datapoint. In amortized inference (right), a shared inference model directly maps each observation to the parameters of its approximate posterior.

Panel of Figure 2 — Figure 2: Schematic decomposition of the inference error. The approximation gap arises from the limited expressiveness of the variational family, which may not contain the true posterior. The amortization gap captures the additional suboptimality introduced by the inference model, which does not recover the optimal variational parameters within that family.

Panel of Figure 3 — Figure 3: Schematic representation of the proposed approach. A shared inference model provides a global parametrization, which is adapted through instance-specific parameter modulations generated by a hypernetwork. This enables input-dependent inference while preserving the efficiency of amortization, without requiring iterative optimization.

Panel of Figure 4 — Figure 4: Conceptual illustration of the proposed IA-VAE method. Starting from the amortized solution $q(\mathbf{z}|\mathbf{x};\phi_{\mathrm{AVI}})$ , input-dependent parameter modulations adapt the encoder for each observation, yielding an instance-adaptive encoder $q(\mathbf{z}|\mathbf{x};\phi(\mathbf{x};\psi))$ . This induces a shift of the variational distribution toward the optimal approximation $q^{*}(\mathbf{z}|\mathbf{x})$ within the variational family.

Panel of Figure 5 — Figure 5: Architectural illustration of the proposed IA-VAE framework. A frozen encoder with global parameters $\phi_{\mathrm{AVI}}$ is augmented by a hypernetwork $h(\mathbf{x},\mathbf{e};\psi)$ that generates input-dependent parameter modulations $\Delta\phi_{\mathrm{AVI,\ell}}$ for each parameter block $\ell$ , with $1\leq\ell\leq L$ . These modulations are combined with the corresponding base parameters to produce instance-specific encoder parameters $\phi(\mathbf{x},\psi)$ , yielding an adaptive inference encoder with enhanced flexibility in the variational posterior.

Component	Synthetic Data	Image Data
Dataset
Dataset	Synthetic ( $N=5000$ )	OMNIGLOT, MNIST, Fashion MNIST
Latent dim	2	32
Observation dim	3	$28\times 28$ (binarized)
Likelihood	Gaussian ( $\sigma=0.1$ )	Bernoulli
Train / Test split	–	Standard protocol

Table 1: Summary of hyperparameters and experimental settings for synthetic and image data experiments.

VAE	IA-VAE
$-7.95\;(5.47)$	$\mathbf{-6.43\pm 0.08\;(5.55\pm 0.01)}$
$-7.86\;(5.48)$	$\mathbf{-6.42\pm 0.04\;(5.54\pm 0.01)}$
$-7.91\;(5.48)$	$\mathbf{-6.43\pm 0.03\;(5.53\pm 0.01)}$
$-7.83\;(5.48)$	$\mathbf{-6.42\pm 0.04\;(5.61\pm 0.02)}$
$-7.84\;(5.46)$	$\mathbf{-6.42\pm 0.02\;(5.58\pm 0.02)}$
$-8.60\;(5.50)$	$\mathbf{-6.45\pm 0.06\;(5.54\pm 0.04)}$

Table 2: Robustness to random initialization on the synthetic dataset. Each row corresponds to a base VAE trained with a different random initialization. For each VAE, 10 IA-VAE instances are trained using different random seeds. The table reports ELBO values (with KL divergence in parentheses), and for IA-VAE we report mean and standard deviation across runs. Comparisons should be made row-wise between VAE and IA-VAE.

Panel of Figure 6 — Figure 6: Posterior density in the latent space for two randomly selected observations from the synthetic dataset. For each observation, the posterior is evaluated on a dense grid over the two-dimensional latent space, and the resulting density is shown as a heatmap; brighter regions correspond to higher posterior density. The red triangle denotes the true latent variable $\mathbf{z}_{true}$ used to generate the observation, the white star denotes the MAP estimate $\mathbf{z}_{\mathrm{MAP}}$ , and the orange and blue crosses denote the posterior means inferred by the VAE ( $\boldsymbol{\mu}_{\mathrm{VAE}}$ ) and the proposed IA-VAE ( $\boldsymbol{\mu}_{\mathrm{IA-VAE}}$ ) model, respectively. The overlaid contours correspond to level sets of a Gaussian approximation of the posterior. The MAP estimate $\mathbf{z}_{\mathrm{MAP}}$ serves as a reference point, as it identifies the most probable latent configuration under the true posterior.

Method	$\mathrm{ELBO}\,\uparrow$	$d_{\mathrm{MAP}}\,\downarrow$	$r_{\mathrm{MAP}}\,\uparrow$
VAE	-7.83 (5.48)	1.44	0.53
IA-VAE	-6.36 (5.55)	0.56	0.84

Table 3: Quantitative comparison of inference accuracy on the synthetic dataset in the oracle setting. The table reports the ELBO (with the corresponding KL divergence in parentheses), the Mahalanobis distance from the MAP estimate

d_{\mathrm{MAP}}

, and the posterior density ratio

r_{\mathrm{MAP}}=p(\boldsymbol{\mu}|\mathbf{x}))/p(\mathbf{z}_{\mathrm{MAP}}|\mathbf{x}))

, where

\mu

denotes the inferred posterior mean and

\mathbf{z}_{\mathrm{MAP}}

the MAP estimate of the posterior. Arrows indicate the preferred direction for each metric.

Figure 7: Parameter efficiency comparison between IA-VAE and standard VAEs with increasing encoder capacity in the oracle setting. The curve shows the ELBO achieved by baseline VAEs as a function of the number of encoder parameters obtained by progressively increasing the hidden layer size. The horizontal red line indicates the ELBO achieved by IA-VAE using a total of 68 parameters for the inference mechanism.

Dataset	VAE	IA-VAE	$p$ -value (ELBO)
OMNIGLOT	$-90.30\pm 0.03$	$\mathbf{-89.73\pm 0.01}$	$<0.01^{*}$
MNIST	$-79.42\pm 0.03$	$\mathbf{-78.68\pm 0.01}$	$<0.01^{*}$
Fashion MNIST	$-225.87\pm 0.07$	$\mathbf{-224.79\pm 0.01}$	$<0.01^{*}$

Table 4: Test ELBO on OMNIGLOT, MNIST, and Fashion MNIST for both VAE and IA-VAE. Results are reported as mean and standard deviation over 10 random seeds. The last column reports the

p

-value of a one-sided paired hypothesis test (

\alpha=0.05

). Entries marked with ∗ indicate statistically significant improvements of IA-VAE over VAE.

Figure 8: Reconstruction examples on MNIST, Fashion-MNIST, and OMNIGLOT. For each dataset, the top row shows the original inputs and the bottom row the corresponding reconstructions obtained using IA-VAE.

Explore similar work

Variational objectives for amortized Bayesian inference in inverse problems: The role of posterior conditioning
Sep 21, 2026Abhishek Srivastava, Arijit Hazra, Rajesh DubbakuAmortized Inference Bayesian Inverse Problems
Department of Mechanical Engineering, Indian Institute of Technology Palakkad · Department of Mathematics, Indian Institute of Technology Palakkad
Amortized Factor Inference Networks for Posterior Inference
May 26, 2026Joohwan Ko, Justin DomkeAmortized Inference Bayesian Inference
Manning College of Information and Computer Sciences University of Massachusetts Amherst
Empirical Variational Autoencoder
Oct 5, 2026Kaede ShioharaVariational Autoencoders Latent Variable Models
The University of Tokyo

Sep 21, 2026J/K move · Enter open · S save

Variational objectives for amortized Bayesian inference in inverse problems: The role of posterior conditioning

Abhishek Srivastava, Arijit Hazra, Rajesh Dubbaku

Department of Mechanical Engineering, Indian Institute of Technology Palakkad · Department of Mathematics, Indian Institute of Technology Palakkad

Amortized Inference Bayesian Inverse Problems Variational Inference PDE Inverse Problems

Open paper AskPDF arXiv

Variational autoencoders (VAEs) offer an efficient approach to amortized Bayesian inference for inverse problems, but posterior accuracy can depend strongly on the choice of variational regularization, particularly when the inverse problem contains weakly identified parameter directions. This study investigates three objectives: a reverse Kullback--Leibler formulation (VAE-KL), an asymmetric Jensen--Shannon formulation (VAE-JS), and a Jensen--Shannon--Wasserstein formulation (VAE-JSWA), which replaces the reverse Kullback--Leibler regularizer with the squared 2-Wasserstein distance while retaining forward-Kullback--Leibler posterior supervision. A full-covariance Gaussian encoder and a pre-trained physics-based surrogate are used for amortized posterior inference. A local linear--Gaussian analysis in the generalized Fisher basis is developed to characterize the variance-dependent gradients of the three objectives. The formulations are first evaluated using linear--Gaussian benchmarks with known posterior solutions and subsequently tested on nonlinear physics-based inverse problems, including an inverse problem governed by a linear ODE and two PDE-constrained problems. VAE-KL performs slightly better than the other formulations in the well-conditioned benchmark, where all three approaches yield comparable posterior approximations, whereas VAE-JSWA provides substantially lower posterior errors in the strongly ill-conditioned benchmark. The nonlinear physics-based problems exhibit a similar conditioning-dependent trend, with JS-based formulations providing greater benefit as posterior ill-conditioning increases. These results indicate that posterior conditioning is an important factor in selecting variational objectives and motivate geometry-adaptive variational inference for Bayesian inverse problems.

Abstract

Explore similar work

Variational objectives for amortized Bayesian inference in inverse problems: The role of posterior conditioning

Amortized Factor Inference Networks for Posterior Inference

Empirical Variational Autoencoder