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
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
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.
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.
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.
Figure 4: Conceptual illustration of the proposed IA-VAE method. Starting from the amortized solution q(z∣x;ϕAVI) , input-dependent parameter modulations adapt the encoder for each observation, yielding an instance-adaptive encoder q(z∣x;ϕ(x;ψ)) . This induces a shift of the variational distribution toward the optimal approximation q∗(z∣x) within the variational family.
Figure 5: Architectural illustration of the proposed IA-VAE framework. A frozen encoder with global parameters ϕAVI is augmented by a hypernetwork h(x,e;ψ) that generates input-dependent parameter modulations ΔϕAVI,ℓ for each parameter block ℓ , with 1≤ℓ≤L . These modulations are combined with the corresponding base parameters to produce instance-specific encoder parameters ϕ(x,ψ) , 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×28 (binarized)
Likelihood
Gaussian ( σ=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)
−6.43±0.08(5.55±0.01)
−7.86(5.48)
−6.42±0.04(5.54±0.01)
−7.91(5.48)
−6.43±0.03(5.53±0.01)
−7.83(5.48)
−6.42±0.04(5.61±0.02)
−7.84(5.46)
−6.42±0.02(5.58±0.02)
−8.60(5.50)
−6.45±0.06(5.54±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.
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 ztrue used to generate the observation, the white star denotes the MAP estimate zMAP , and the orange and blue crosses denote the posterior means inferred by the VAE ( μVAE ) and the proposed IA-VAE ( μIA−VAE ) model, respectively. The overlaid contours correspond to level sets of a Gaussian approximation of the posterior. The MAP estimate zMAP serves as a reference point, as it identifies the most probable latent configuration under the true posterior.
Method
ELBO↑
dMAP↓
rMAP↑
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 dMAP , and the posterior density ratio rMAP=p(μ∣x))/p(zMAP∣x)) , where μ denotes the inferred posterior mean and zMAP 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±0.03
−89.73±0.01
<0.01∗
MNIST
−79.42±0.03
−78.68±0.01
<0.01∗
Fashion MNIST
−225.87±0.07
−224.79±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 ( α=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.
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.
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 promises fast test-time Bayesian inference, but existing methods are inherently tied to fixed models. Extending amortization to unseen models typically requires retraining or costly test-time finetuning. In this paper, we ask: is it possible to build a single inference network capable of generalizing across varying priors, likelihoods, and dimensionality? We introduce Amortized Factor Inference Networks (AFINs), a family of encode-merge-decode inference networks built on dimension-independent modules that map a model specification and its observations to the parameters of a variational posterior. Experimentally, a single trained AFIN achieves posterior accuracy comparable to NUTS and several variational inference methods, while requiring 2 to 4 orders of magnitude less test-time compute. Code is available at https://github.com/joohwanko/AFINs.
Joohwan Ko, Justin Domke
Manning College of Information and Computer Sciences University of Massachusetts Amherst
We present Empirical Variational Autoencoder, a general generative framework for continuous-valued (i.e., non-vector-quantized) sequences. EVA is based on the evidence lower bound of the Variational Autoencoder (VAE) but learns autoregressive latent priors empirically from training data, which can be implemented only by an additional single linear layer on top of VAEs. By replacing the conventional standard-Gaussian constraint with the self-predicted priors, EVA significantly alleviates the latent distribution gap between prior and posterior which is typically observed in conventional VAEs, and leads to high-fidelity ancestral sampling for sequential data generation. Extensive experiments on image and sound synthesis demonstrate that EVA achieves competitive generation quality with autoregressive diffusion baselines despite its much faster inference time.