Meta-Amortization: Sharing Inference Across Tasks

Updated 27 April 2026

Meta-amortization is a statistical learning paradigm that extends classical amortized inference by using a global network to share computation across related tasks.
It enables rapid adaptation in few-shot learning, self-supervised representation, and optimal transport by bypassing costly per-task optimization.
The approach reduces computational overhead and improves generalization through meta-training and regularization techniques tailored for diverse problem settings.

Meta-amortization is a paradigm in statistical learning and optimization that extends classical amortized inference beyond sharing computation across data points to sharing it across entire sets of related tasks, probabilistic models, or optimization problems. Unlike standard amortization—which trains a single inference network for a fixed generative model—meta-amortization leverages a global inference (or optimization) network that can rapidly adapt or generalize to novel tasks or models after meta-training, often without further per-task optimization. This concept underpins algorithms in probabilistic meta-learning, fast adaptation for few-shot learning, self-supervised representation learning, and high-throughput optimization, combining the efficiency of amortized inference with the flexibility and generalization of meta-learning (Gordon et al., 2018, Wu et al., 2019, Jang et al., 2023, Chang et al., 2024, Truong et al., 16 Apr 2026, Bae et al., 2022, Hayashi et al., 2020).

1. Foundational Principles and Formalization

Meta-amortization arises in settings where a family of related probabilistic models or tasks is indexed by latent parameters or data distributions. Classical amortization, as in variational autoencoders (VAEs), trains an encoder $q_\phi(z|x)$ to approximate the posterior for a single generative model $p_\theta(x,z)$ . Meta-amortization generalizes this by learning a joint inference network $q_\phi(\psi|D)$ or $g_\phi(p_{D_i}, x)$ , which takes as input not just individual datapoints, but an entire support set (or a summary/context of a task, or a marginal distribution over data), and outputs task-specific or model-specific posteriors or parameterizations.

Key instantiations include:

VERSA and MetaVAE: $q_\phi(\psi|D)$ is trained to approximate $p(\psi|D)$ for tasks $\tau$ with dataset $D^\tau$ , amortizing inference over both datapoints and tasks (“doubly-amortized inference”) (Gordon et al., 2018, Wu et al., 2019).
MetaMAE: Interprets each masked autoencoder reconstruction, determined by a random mask, as a separate task. The Transformer encoder produces an initial latent representing the support set (unmasked tokens), followed by gradient-based adaptation to the reconstruction target (Jang et al., 2023).
Amortized OT: Solves a meta-collection of optimal transport problems by learning a mapping from problem descriptors to Kantorovich potentials, allowing for rapid inference on new problem instances (Truong et al., 16 Apr 2026).
Amortized Conditioning Engine (ACE): Learns a single transformer-based conditioning engine for arbitrary probabilistic conditioning and prediction, directly ingesting both observed data and interpretable latent variables (Chang et al., 2024).

Mathematically, the meta-amortized objective for variational inference can be formalized as the MetaELBO: $\max_\phi\, \mathbb{E}_{p_{D_i}\sim p_\mathcal{M}} \mathbb{E}_{x\sim p_{D_i}(x)} \mathbb{E}_{z\sim g_\phi(z|p_{D_i}, x)} [\log p_{\theta_i}(x, z) - \log g_\phi(z|p_{D_i}, x)]$ where $p_\mathcal{M}$ is a meta-distribution over data marginals (Wu et al., 2019).

2. Meta-Amortization in Probabilistic Meta-Learning

In probabilistic meta-learning for few-shot adaptation, meta-amortization replaces per-task gradient-based adaptation with a global, task-conditional inference network. ML-PIP and VERSA minimize the negative expected log-posterior-predictive across tasks using an inference network $p_\theta(x,z)$ 0, achieving rapid task adaptation: $p_\theta(x,z)$ 1 This approach sidesteps inner-loop optimization and second-derivative computations required by MAML, offering a single forward-pass per task, with learned amortization of inference across both datapoints and tasks (Gordon et al., 2018, Wu et al., 2019).

The ACE architecture generalizes this further by allowing arbitrary probabilistic conditioning and prediction queries, accepting as context a mixture of observed data, latent-variable observations, and arbitrary priors. Prediction is performed via cross-attention over the context, producing output distributions for any target variable in one pass (Chang et al., 2024).

3. Meta-Amortized Self-Supervision and Representation Learning

MetaMAE demonstrates meta-amortization in self-supervised learning by reframing the masked-token reconstruction problem as meta-learning over randomly masked reconstruction tasks. For each input $p_\theta(x,z)$ 2 tokenized into $p_\theta(x,z)$ 3 units, a support set $p_\theta(x,z)$ 4 (unmasked tokens) and query set $p_\theta(x,z)$ 5 (masked tokens) are defined. The transformer encoder computes an amortized latent $p_\theta(x,z)$ 6, which is further adapted via a single gradient step to an adapted latent $p_\theta(x,z)$ 7: $p_\theta(x,z)$ 8 A contrastive alignment loss encourages the amortized and adapted latents to be close for the same task and dissimilar across tasks, facilitating fast adaptation and robust feature extraction across diverse modalities. This scheme achieves state-of-the-art results on modality-agnostic self-supervised benchmarks (DABS), outperforming conventional MAE and other baselines (Jang et al., 2023).

4. Meta-Amortization in Optimization and Control

Amortized Proximal Optimization (APO) frames meta-optimization as meta-amortized adaptation of optimization parameters, such as global learning rates or structured preconditioners. APO amortizes the inner minimization of a stochastic proximal-point objective by meta-learning a parametric update rule $p_\theta(x,z)$ 9: $q_\phi(\psi|D)$ 0 $q_\phi(\psi|D)$ 1 (e.g., learning rate, preconditioner) is updated through a meta-objective that evaluates the improvement on a one-step lookahead and penalizes function- and weight-space divergence. Classical optimizers such as natural gradient or KFAC are recovered as special cases under certain assumptions (Bae et al., 2022).

This approach transfers well to new tasks and dynamically adapts optimizer behavior online, amortizing over the space of encountered optimization tasks.

5. Meta-Amortization Error and Regularization

Meta-amortization, while efficient, is susceptible to error stemming from (i) the variational approximation gap due to restricted posterior families, and (ii) the amortization gap when the inference network cannot exactly solve every per-task inference problem. Under small support sets (few-shot learning), this frequently results in posterior collapse, where the variational posterior degenerates to a single point or ignores the task-specific latent entirely (Hayashi et al., 2020).

To address this, meta-regularization techniques such as cyclical annealing schedules for KL or divergence penalties and Maximum Mean Discrepancy (MMD) regularization are introduced. The cyclical annealing schedule forces the model to carry information in the latent by periodically relaxing and reintroducing the regularizer, while replacing KL with MMD ensures tractable alignment between inferred distributions over tasks (Hayashi et al., 2020). These methods demonstrably reduce meta-amortization error and achieve superior few-shot performance compared to standard meta-learning algorithms.

6. Applications, Limitations, and Extensions

Meta-amortization is applied in:

Few-shot learning: Rapid probabilistic adaptation to new classification or regression tasks (Gordon et al., 2018, Wu et al., 2019).
Modality-agnostic SSL: Representation learning across image, audio, and text modalities using architectures such as MetaMAE (Jang et al., 2023).
Optimal Transport: Solving a distribution of OT problems using regression- or objective-based meta-amortized mappings from problem descriptors to Kantorovich potentials, far outpacing conventional solvers in evaluation cost (Truong et al., 16 Apr 2026).
Autonomous Conditioning and Inference: Unified probabilistic conditional inference and simulation tasks in one-pass transformer models such as ACE (Chang et al., 2024).
Meta-Optimization: Adaptive optimizer schemes for deep learning models without per-task tuning or re-derivation of updates (Bae et al., 2022).

Limitations arise in capacity constraints of the meta-inference network, potential for meta-overfitting if the meta-training set is insufficiently representative, and approximation errors when true task posteriors or optima deviate significantly from the learned parametric family. Extensions include nonlinear amortization architectures, learned projections or context encodings, and broadening to other problem structures such as unbalanced or Gromov–Wasserstein OT (Truong et al., 16 Apr 2026).

Meta-amortization is distinguished from classical amortized inference by its sharing across both datapoints and tasks/models, in contrast to standard $q_\phi(\psi|D)$ 2 which is retrained when the task distribution changes. Unlike MAML-style meta-learners, which require per-task adaptation via optimization steps at deployment, meta-amortized schemes yield immediate inference or optimization results for new tasks in a single network pass, offering both practical speed and improved cross-task generalization (Gordon et al., 2018, Wu et al., 2019, Chang et al., 2024).

Empirical evidence from benchmarks in few-shot reasoning, representation learning, optimization, and optimal transport confirm significant accuracy and efficiency gains over non-amortized or semi-amortized baselines, especially in settings requiring flexible generalization across highly heterogeneous task collections (Gordon et al., 2018, Jang et al., 2023, Truong et al., 16 Apr 2026, Chang et al., 2024, Bae et al., 2022, Hayashi et al., 2020).