Iterative Amortized Inference: Unifying In-Context Learning and Learned Optimizers (2510.11471v1)

Abstract: Modern learning systems increasingly rely on amortized learning - the idea of reusing computation or inductive biases shared across tasks to enable rapid generalization to novel problems. This principle spans a range of approaches, including meta-learning, in-context learning, prompt tuning, learned optimizers and more. While motivated by similar goals, these approaches differ in how they encode and leverage task-specific information, often provided as in-context examples. In this work, we propose a unified framework which describes how such methods differ primarily in the aspects of learning they amortize - such as initializations, learned updates, or predictive mappings - and how they incorporate task data at inference. We introduce a taxonomy that categorizes amortized models into parametric, implicit, and explicit regimes, based on whether task adaptation is externalized, internalized, or jointly modeled. Building on this view, we identify a key limitation in current approaches: most methods struggle to scale to large datasets because their capacity to process task data at inference (e.g., context length) is often limited. To address this, we propose iterative amortized inference, a class of models that refine solutions step-by-step over mini-batches, drawing inspiration from stochastic optimization. Our formulation bridges optimization-based meta-learning with forward-pass amortization in models like LLMs, offering a scalable and extensible foundation for general-purpose task adaptation.

• The paper proposes iterative amortized inference that refines task adaptation over mini-batches, unifying meta-learning, in-context learning, and learned optimizers.
• It introduces a taxonomy of parametric, implicit, and explicit amortization to clarify trade-offs in expressivity, scalability, and efficiency.
• Extensive experiments show that iterative refinement enhances performance on predictive and generative tasks while significantly improving runtime efficiency.

Iterative Amortized Inference: A Unified Framework for In-Context Learning and Learned Optimizers

Overview

This paper introduces a unified formalism for amortized learning, encompassing meta-learning, in-context learning (ICL), prompt tuning, and learned optimizers. The authors propose a taxonomy—parametric, explicit, and implicit amortization—based on how task adaptation is encoded and leveraged. The central contribution is the formulation of iterative amortized inference, which enables scalable task adaptation by refining solutions over mini-batches, bridging the gap between optimization-based meta-learning and forward-pass amortization as in LLMs. The work is supported by extensive empirical analysis across predictive and generative tasks, with ablations on architecture and information modalities.

Unified Framework and Taxonomy

The paper formalizes amortized learning as a two-component system: (1) a shared mechanism encoding task-invariant inductive biases, and (2) a task adaptation function utilizing data from novel tasks. The general objective is:

$$\min_{\gamma, \phi} \mathbb{E}_{\mathcal{T}} \mathbb{E}_{(x, y), \mathcal{D}} \left[ \ell\left(y, f_\gamma\left(x, g_\phi(\mathcal{D})\right)\right) \right]$$

where $f_\gamma$ is the prediction function and $g_\phi$ is the adaptation mechanism mapping task data $\mathcal{D}$ to a task-specific representation.

Three amortization regimes are defined:

• Parametric: $f$ is fixed, $g_\phi$ is learned (e.g., hypernetworks, learned optimizers). Task adaptation is externalized via explicit parameter inference.
• Implicit: $f_\gamma$ is learned, $g$ is fixed (e.g., ICL, PFNs). Task adaptation is internalized via forward-pass conditioning.
• Explicit: Both $f_\gamma$ and $g_\phi$ are learned (e.g., neural processes). Task adaptation is jointly modeled, with a learned likelihood and a dataset-level embedding.

This taxonomy clarifies the trade-offs in expressivity, scalability, and efficiency. For example, parametric methods can leverage gradients for efficient optimization but may be limited in representational richness, while implicit methods can model richer task structures but are constrained by context length and computational cost.

Figure 1: Iterative Amortized Inference for parametric, explicit, and implicit parameterizations, highlighting the distinction in how task-specific information is incorporated and refined.

Iterative Amortized Inference

A key limitation of existing amortized methods is their inability to scale to large datasets due to context length or reliance on low-dimensional summaries. The authors propose iterative amortized inference, where adaptation is performed via stepwise refinement over mini-batches, inspired by stochastic optimization. This approach generalizes learned optimizers and hypernetworks by allowing the adaptation function to process streams of mini-batches, incorporating both observations and gradients.

For parametric and explicit models, the adaptation function $g_\phi$ is implemented as a recurrent sequence model $h_\phi$:

$$z^{(0)} \xrightarrow{h_\phi(\cdot, \mathcal{B}_0)} z^{(1)} \xrightarrow{h_\phi(\cdot, \mathcal{B}_1)} \cdots \xrightarrow{h_\phi(\cdot, \mathcal{B}_{k-1})} z^{(k)}$$

where each $\mathcal{B}_i$ is a mini-batch from the task data. For implicit models, the prediction state $\hat{y}^{(i)}$ is iteratively refined conditioned on the query and mini-batch.

This iterative process enables scalable adaptation, efficient memory usage, and improved generalization, as demonstrated empirically.

Empirical Evaluation

The framework is evaluated on a suite of predictive and generative tasks, including linear regression, MNIST/FashionMNIST/ImageNet classification, topological order prediction in SCMs, and generative modeling with GMMs and alphabets. The models are parameterized as transformers, with careful design of masking and state representations to enable efficient parallelization and variable context lengths.

Predictive Tasks

• Parametric Amortization: Iterative refinement consistently improves performance across tasks and modalities (gradients, observations, or both). Notably, combining gradients and observations yields the best results, especially in low-dimensional settings. The approach generalizes out-of-distribution (e.g., ImageNet pretraining transfers to CIFAR-10).
• Explicit Amortization: Performance improves with more refinement steps, but reliance on gradients becomes critical for complex tasks. Explicit models struggle to scale to high-dimensional tasks like ImageNet, often performing at chance.
• Implicit Amortization: Iterative refinement leads to substantial gains, outperforming single-step ICL baselines. The approach is robust across in-distribution and out-of-distribution settings, including causal structure learning.

Generative Tasks

• Implicit Models: Iterative amortization improves sample quality (measured by Wasserstein distances) for both in-distribution and out-of-distribution tasks in GMM and alphabet generation.

Figure 2: Samples generated from the implicit generative model for GMM and Alphabets task, showing improved fidelity with more refinement steps.

Ablations and Analysis

• Gradient vs. Observation Signal: Sole reliance on gradients is suboptimal in low dimensions; observations provide richer information. In high dimensions, gradients are more reliable due to the larger hypothesis space.
• History Length: Including multiple past states and gradients does not significantly improve performance, except in specific settings.
• State Parameterization: For implicit models, carrying over logits as the recurrent state yields the best performance, with softmax outputs performing comparably.

  Figure 3: Analysis of leveraging multiple past states and gradients for parametric amortization; additional history provides limited benefit.

  

  Figure 4: Carrying over logits as the recurrent state in implicit models outperforms other state representations.

• Runtime Efficiency: Iterative amortization is $K$ times more efficient than one-step models when processing the same amount of data, both in runtime and memory.

Implementation Considerations

• Transformer Architecture: All adaptation functions are implemented as transformers with variable context lengths, using causal or non-causal masking depending on the amortization regime.
• State Management: For parametric/explicit models, the state is a task-specific parameter or latent; for implicit models, the state is tied to the query and prediction.
• Training Objective: Models are trained greedily, optimizing for single-step improvements without backpropagation through time, which simplifies optimization and reduces memory requirements.
• Scalability: The iterative approach enables efficient handling of large datasets by processing mini-batches sequentially, avoiding the quadratic scaling of context length in standard ICL.

Implications and Future Directions

The unified framework clarifies the relationships and trade-offs among meta-learning, ICL, and learned optimizers. Iterative amortized inference provides a scalable, extensible foundation for general-purpose task adaptation, with demonstrated benefits in both predictive and generative settings. The empirical results highlight that:

• Iterative refinement is consistently beneficial across amortization regimes and tasks.
• Explicit models are often suboptimal compared to parametric models, despite their greater expressivity, likely due to optimization challenges.
• Gradient information is not always sufficient; leveraging observations is critical, especially in low-dimensional or data-rich regimes.

Future work should explore richer learning frameworks for optimizing the unified objective beyond greedy refinement, such as evolutionary algorithms or reinforcement learning, and investigate more sophisticated forms of persistent memory for robust adaptation in non-iid environments.

Conclusion

This work provides a principled, scalable approach to amortized learning, unifying disparate methods under a common formalism and demonstrating the practical advantages of iterative amortized inference. The framework enables efficient, flexible task adaptation and opens avenues for further research in scalable meta-learning, in-context adaptation, and learned optimization.