Iterative Amortized Inference: Unifying In-Context Learning and Learned Optimizers

## Abstract

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.

## 1 Introduction

Consider the problem of modeling the motion of an object on various planets, under different gravitational conditions. While each planet has its own gravitational constant, the underlying dynamics, governed by Newtonian physics, remain invariant. One could train a new model from scratch for each planet, but this would ignore the substantial structure shared across tasks. A more sample-efficient approach is to reuse knowledge acquired from other planets and adapt only a small set of task-specific parameters, such as the gravitational constant. In this setting, the shared inductive bias (e.g., the equations of motion) is amortized across tasks, enabling fast, local, and sample-efficient adaptation.

This principle serves as the basis for several methods which avoid solving each task from scratch by learning mechanisms that captures the shared information between tasks, enabling rapid adaptation to new ones.
For instance, gradient-based meta-learners (Finn et al., 2017; Rusu et al., 2018) encode inductive biases through learned meta initialization
to facilitate downstream task training.
In contrast, amortization is implicitly incentivized in large language models (LLMs) by training on diverse contexts, allowing them to solve new tasks at inference by conditioning on observations (ICL; Brown et al., 2020; Dong et al., 2022) or instructions (Lester et al., 2021; Liu et al., 2021).
Learned optimizers (Andrychowicz et al., 2016; Metz et al., 2022a, b; Knyazev et al., 2024) instead amortize optimization itself, learning to predict parameter updates of fixed models conditioned on gradients.

**Table 1: : Functional decomposition of amortized learners. We express each method in terms of $g_{\bm{\varphi}}$ that maps task observations $\mathcal{D}_{\mathcal{T}}^{\text{train}}$ to some representation ${\bm{\theta}}_{\mathcal{T}}$ $-$ weights, prompts, dataset itself $-$ which is then fed along with query to $f_{\gamma}({\mathbf{x}},{\bm{\theta}}_{\mathcal{T}})$ for prediction. Differences between methods arise from the aspects of learning they amortize, and our proposed taxonomy offers a categorization.**
| Method | $g_{\bm{\varphi}}({\mathcal{D}}_{\mathcal{T}}^{\text{train}})$ | ${\bm{\theta}}_{\mathcal{T}}$ | $f_{\gamma}({\mathbf{x}},{\bm{\theta}}_{\mathcal{T}})$ | Amortization |
| --- | --- | --- | --- | --- |
| Standard training | SGD | weights | Fixed architecture ($\gamma=\varnothing$) | $-$ |
| MAML | SGD, learned ${\bm{\theta}}_{0}$ | weights | Fixed architecture ($\gamma=\varnothing$) | ${\bm{\theta}}_{0}$ |
| Learned optimizers | Learned updates | weights | Fixed architecture ($\gamma=\varnothing$) | $g_{\bm{\varphi}}$ |
| ICL | Identity Map | Tokens | Transformer (weights $\gamma$) | $f_{\gamma}$ |
| Parametric | Learned updates | weights | Fixed architecture ($\gamma=\varnothing$) | $g_{\bm{\varphi}}$ |
| Implicit | Subsampling | Batch ${\mathcal{B}}_{\mathcal{T}}^{\text{train}}$ | Transformer (weights $\gamma$) | $f_{\gamma}$ |
| Explicit | Learned updates | latents | Transformer (weights $\gamma$) | $g_{\bm{\varphi}},f_{\gamma}$ |

We introduce a unified framework that encapsulates a broad spectrum of amortized learning methods $-$ meta-learners, learned optimizers, ICL $-$ as particular instances within a common structure, characterized by two primary components: (1) a shared mechanism to encode task-invariant inductive biases, and (2) a task adaptation function that utilizes data from novel tasks to model task-specific behavior. Our framework identifies three distinct amortization regimes distinguished by the inductive biases involved within each component $-$ (a) parametric amortization where a learned function maps task-specific data to corresponding parameters of a fixed model, e.g. learned optimizers (Andrychowicz et al., 2016) and hypernetworks (Ha et al., 2016), (b) implicit amortization where a single model jointly internalizes task-invariant mechanisms and task adaptation through forward-pass conditioning (Brown et al., 2020; Müller et al., 2021), e.g. ICL, and (c) explicit amortization which disentangles generalization and local adaptation by learning both a dataset-level embedding and a task-conditioned prediction function (Garnelo et al., 2018a; Mittal et al., 2024) via architectural inductive biases.

These paradigms reflects trade-offs in expressivity, scalability, and efficiency $-$ using gradients, like in learned optimizers, provides direct access to task-specific loss landscapes, making optimization efficient but potentially limited in the richness of task information encoded, while conditioning on observations, like in in-context learning and hypernetworks, allows richer task representations but can be computationally expensive, especially as the conditioning dataset grows. In contrast, pooling-based methods offer compact representations but struggle with fine-grained task variability.

We also propose an iterative amortization framework bridging these perspectives by embracing the stochastic optimization viewpoint. Rather than collapsing task information into a single pass or summary, we model amortization as an iterative refinement process over stochastic mini-batches of task data, leveraging either observations directly or through gradients. This mirrors the success of stochastic gradient descent in scaling optimization to large datasets. By doing so, we generalize the notion of learned optimizers and hypernetworks: instead of operating solely on gradients or generating parameters one-shot, the adaptation function now incorporates streams of mini-batches of observations and corresponding gradients from the task data, enabling greater flexibility and scalability.
Our contributions are $-$

- •
Unified Framework: A general formulation of amortized learning that connects meta-learning, in-context learning, prompt tuning, and learned optimizers as special cases ([Section˜2](https://arxiv.org/html/2510.11471v1#S2)).
- •
Amortization Taxonomy: A clear categorization of amortization methods into parametric, implicit, and explicit regimes, based on how they encode inductive bias and adapt to new tasks ([Section˜3](https://arxiv.org/html/2510.11471v1#S3)).
- •
Stochastic Iterative Amortization: A scalable approach through iterative refinement over mini-batches, overcoming limitations of scaling to large datasets ([Section˜4](https://arxiv.org/html/2510.11471v1#S4)).

## 2 Unified Perspective

Given a task ${\mathcal{T}}$ and a corresponding data distribution $p_{\mathcal{T}}$, a fundamental problem of machine learning is to learn a model $f(\cdot,{\bm{\theta}}^{*}_{\mathcal{T}})$
that minimizes the true risk based on a loss function ${\mathcal{L}}$, i.e.

$$ $\displaystyle{\bm{\theta}}^{*}_{\mathcal{T}}$ $\displaystyle=\arg\min_{{\bm{\theta}}}\mathbb{E}_{{\mathbf{x}},{\mathbf{y}}\sim p_{\mathcal{T}}}\left[{\mathcal{L}}\left({\mathbf{y}},f({\mathbf{x}},{\bm{\theta}})\right)\right],$ (1) $$

which is often accomplished by empirical risk minimization where the task ${\mathcal{T}}$ is described through a set of i.i.d observations ${\mathcal{D}}_{\mathcal{T}}=\{({\mathbf{x}},{\mathbf{y}})\sim p_{\mathcal{T}}\}_{i}$ and we learn a proxy to the optimal model as

$$ $\displaystyle\hat{{\bm{\theta}}}_{\mathcal{T}}$ $\displaystyle=\arg\min_{{\bm{\theta}}}\sum_{({\mathbf{x}},{\mathbf{y}})\in{\mathcal{D}}_{\mathcal{T}}}{\mathcal{L}}\left({\mathbf{y}},f({\mathbf{x}},{\bm{\theta}})\right),$ (2) $$

in the hope that $\hat{{\bm{\theta}}}_{\mathcal{T}}$ generalizes well to new samples from the same distribution, and thus leads to low true risk modeling similar solutions to the optimal one ${\bm{\theta}}^{*}_{\mathcal{T}}$(^1^11The same extends to generative modeling with ${\mathbf{y}}$ denoting images / sentences and ${\mathbf{x}}$ class conditioning.). Owing to theoretical advances (Shalev-Shwartz and Ben-David, 2014; Zhang et al., 2016; Neal et al., 2018), the scale of data available (Brown et al., 2020; Achiam et al., 2023), and the architectural models and inductive biases (Bahdanau et al., 2014; Vaswani et al., 2017), we have evidence that this paradigm leads to generalizable models. But what happens when we look at a new task?

While effective for single-task learning, the standard empirical risk minimization approach in [Eq.˜2](https://arxiv.org/html/2510.11471v1#S2.E2) learns a separate model $\hat{{\bm{\theta}}}_{{\mathcal{T}}}$ for each task ${\mathcal{T}}$, using only its associated dataset ${\mathcal{D}}_{{\mathcal{T}}}$. This setup suffers from two key limitations: (a) it lacks the ability to systematically adapt to new tasks, even if they are similar, and (b) it fails to leverage cross-task information when tasks ${{\mathcal{T}}_{1}}$ and ${\mathcal{T}}_{2}$ share underlying structure. Data from ${\mathcal{T}}_{2}$ could provide valuable inductive signal for learning $\hat{{\bm{\theta}}}_{{\mathcal{T}}_{1}}$, but is ignored.

Traditional meta-learning approaches e.g. Model-Agnostic Meta-Learning (MAML; Finn et al., 2017) alleviate this problem by learning global initialization parameters $\hat{{\bm{\theta}}}_{0}$ which encapsulate all the shareable knowledge across ${\mathcal{T}}$s such that small amounts of finetuning leads to good models ${\bm{\theta}}_{\mathcal{T}}$ for any particular task (Nichol and Schulman, 2018; Nichol et al., 2018; Rajeswaran et al., 2019).

$$ $\displaystyle\hat{{\bm{\theta}}}_{0}=\arg\min_{{\bm{\theta}}_{0}}\mathbb{E}_{{\mathcal{T}}}\mathbb{E}_{{\mathbf{x}},{\mathbf{y}},{\mathcal{D}}_{\mathcal{T}}\sim p_{\mathcal{T}}}\left[{\mathcal{L}}\left({\mathbf{y}},f({\mathbf{x}},{\bm{\theta}}_{\mathcal{T}})\right)\right],\quad\text{where}\quad{\bm{\theta}}_{\mathcal{T}}=g_{{\bm{\theta}}_{0}}({\mathcal{D}}_{\mathcal{T}}).$ (3) $$

where $g_{{\bm{\theta}}_{0}}$ is an optimization routine, most commonly stochastic gradient descent, with learnable initialization ${\bm{\theta}}_{0}$. Alternatively, hypernetworks (Li and Liang, 2021; Ha et al., 2016; Gaier and Ha, 2019; Jia et al., 2016; Munkhdalai and Yu, 2017) and learned optimizers (Andrychowicz et al., 2016; Li et al., 2017; Metz et al., 2019; Wichrowska et al., 2017; Metz et al., 2022b; Li and Malik, 2017) instead directly model task-specific parameters as outputs of a learned process, often without a notion of “good initialization”. Such methods are also naturally described by Equation [3](https://arxiv.org/html/2510.11471v1#S2.E3) where ${\bm{\theta}}_{0}$ is replaced by ${\bm{\varphi}}$ and $g_{\bm{\varphi}}$ describes a sequence model which either takes observations (e.g. hypernetworks) or gradients (e.g. learned optimizers) as input signal.

Certain meta-learning methods like example-based in-context learning or prior fitted networks instead directly model the conditional predictive distribution (Garg et al., 2022; Müller et al., 2021), i.e.

$$ $\displaystyle\hat{\gamma}$ $\displaystyle=\arg\min_{\gamma}\mathbb{E}_{{\mathcal{T}}}\mathbb{E}_{{\mathbf{x}},{\mathbf{y}},{\mathcal{D}}_{\mathcal{T}}}\Big[{\mathcal{L}}\Big({\mathbf{y}},f_{\gamma}({\mathbf{x}},{\mathcal{D}}_{\mathcal{T}})\Big)\Big],$ (4) $$

where the model does not expose explicit task parameters ${\bm{\theta}}_{\mathcal{T}}$ anymore, and conditioning on the task is achieved either through direct conditioning on observations ${\mathcal{D}}_{\mathcal{T}}$ $-$ Prior Fitted Networks (PFNs) or ICL (Brown et al., 2020; Dong et al., 2022; Garg et al., 2022; Von Oswald et al., 2023; Hollmann et al., 2022) $-$ or its compressed natural language description when using LLMs (Mishra et al., 2021; Efrat and Levy, 2020; Sanh et al., 2021; Wei et al., 2022b, a; Min et al., 2022).

The fundamental approach across all these methods is to learn an amortized model that leverages shared knowledge across tasks for each individual task, and thus can generalize to new ones with similar underlying mechanisms. All the above problems can then be unified as

$$ $\displaystyle\min_{\gamma,{\bm{\varphi}}}\mathbb{E}_{{\mathcal{T}}}\mathbb{E}_{{\mathbf{x}},{\mathbf{y}},{\mathcal{D}}_{\mathcal{T}}}\left[{\mathcal{L}}\left({\mathbf{y}},f_{\gamma}\left({\mathbf{x}},g_{\bm{\varphi}}\left({\mathcal{D}}_{\mathcal{T}}\right)\right)\right)\right]$ (5) $$

where $f_{\gamma}({\mathbf{x}},g_{\bm{\varphi}}({\mathcal{D}}_{\mathcal{T}}))$(^2^22$g_{\bm{\varphi}}({\mathcal{T}})$ describes the general form, while we look at cases where ${\mathcal{D}}_{\mathcal{T}}$ provides information about ${\mathcal{T}}$.) is an inference procedure encapsulating our modeling assumptions about obtaining predictions for query ${\mathbf{x}}$ and a set of observations ${\mathcal{D}}_{\mathcal{T}}$.
This general framework unifies supervised learning, meta-learning, in-context learning, and learned optimizers under a common formalism by varying which components $-$ $f_{\gamma}$ or $g_{\bm{\varphi}}$ $-$ are learned and how.

### 2.1 Specific Cases of Amortization

In this section, we describe how individual methods are special cases of our unified framework. Given a set of tasks ${\mathcal{T}}$ along with their corresponding training ${\mathcal{D}}_{\mathcal{T}}^{\text{train}}$ and validation ${\mathcal{D}}_{\mathcal{T}}^{\text{valid}}$ set, the empirical counterpart of our proposed framework can be described as

$$ $\displaystyle\min_{\gamma,{\bm{\varphi}}}\mathbb{E}_{\mathcal{T}}\left[\sum_{({\mathbf{x}},{\mathbf{y}})\in{\mathcal{D}}_{\mathcal{T}}^{\text{valid}}}{\mathcal{L}}\left({\mathbf{y}},f_{\gamma}\left({\mathbf{x}},{\bm{\theta}}_{\mathcal{T}}\right)\right)\right],\quad\text{where }{\bm{\theta}}_{\mathcal{T}}=g_{\bm{\varphi}}({\mathcal{D}}_{\mathcal{T}}^{\text{train}})$ (6) $$

where $f_{\gamma}(\cdot,\ {\bm{\theta}}_{\mathcal{T}})$ denotes a model with shared parameters $\gamma$ and task-specific parameters ${\bm{\theta}}_{\mathcal{T}}$ (e.g. weights, soft prompts, latent states, or context tokens). The function $g_{\bm{\varphi}}$, parameterized by ${\bm{\varphi}}$, represents the inner optimization mechanism that maps the support set ${\mathcal{D}}_{\mathcal{T}}^{\text{train}}$ for task ${\mathcal{T}}$ to ${\bm{\theta}}_{\mathcal{T}}$. The loss ${\mathcal{L}}$ is then computed on the query set ${\mathcal{D}}_{\mathcal{T}}^{\text{valid}}$. By selecting different forms for $f$, partitioning parameters between ${\bm{\varphi}}$ and $\gamma$, and varying the adaptation function $g_{\bm{\varphi}}$, this framework subsumes a wide spectrum of learning paradigms, as outlined in [Tables˜1](https://arxiv.org/html/2510.11471v1#S1.T1) and [7](https://arxiv.org/html/2510.11471v1#S7). In particular, we obtain

- •
Gradient-based supervised learning when $f$ is fixed $-$ e.g. a neural network architecture $-$ and $g$ an optimization procedure, like stochastic gradient descent, solving [Eq.˜2](https://arxiv.org/html/2510.11471v1#S2.E2).
- •
MAML when $f$ is fixed and $g_{\bm{\varphi}}$ is an optimization procedure with learnable initialization ${\bm{\theta}}_{0}\in{\bm{\varphi}}$.
- •
Learned Optimizers / Hypernetworks when $f$ is fixed and $g_{\bm{\varphi}}$ is a neural network with learnable parameters ${\bm{\varphi}}$. We obtain hypernetworks when $g_{\bm{\varphi}}$ maps observations to parameters, and learned optimizers when $g_{\bm{\varphi}}$ is an iterative procedure, i.e. $g_{\bm{\varphi}}:=h_{\bm{\varphi}}^{(k)}\circ\ldots\circ h_{\bm{\varphi}}^{(1)}$ where $h_{\bm{\varphi}}^{(i)}$ is a sequence model taking ${\bm{\theta}}^{(0)},\ldots,{\bm{\theta}}^{(i-1)}$ and their gradients $\nabla^{(j)}:=\nabla_{\bm{\theta}}\sum_{({\mathbf{x}},{\mathbf{y}})\in{\mathcal{B}}^{(j)}}{\mathcal{L}}({\mathbf{y}},f({\mathbf{x}},{\bm{\theta}}))\big\rvert_{{\bm{\theta}}^{(j)}}$ as input, such that ${\bm{\theta}}_{\mathcal{T}}={\bm{\theta}}^{(k)}$ and ${\mathcal{B}}^{(j)}\subseteq{\mathcal{D}}_{\mathcal{T}}$.(^3^33A system is Markovian if it only relies on the current $({\bm{\theta}}^{(j)},\nabla^{(j)})$ or non Markovian if on more past states.)
- •
In-Context Learning when $g$ samples a set of observations and $f_{\gamma}$ is modeled as a predictive sequence model with the context as observations. If $g$ describes ${\mathcal{T}}$ in natural language, then we recover the prompt-based ICL $-$ an emergent phenomena from pretraining of LLMs (Brown et al., 2020).

Non-unique decomposition.
Importantly, the functional form $f_{\gamma}({\mathbf{x}},g_{\bm{\varphi}}({\mathcal{D}}_{\mathcal{T}}))$ does not define a unique way to split the computation between $f_{\gamma}$ and $g_{\bm{\varphi}}$. In principle, any computation performed by $g_{\bm{\varphi}}$, along with its parameters, can be folded into $f_{\gamma}$, making the decomposition arbitrary. However, the converse is not true as $g_{\bm{\varphi}}$ does not have access to ${\mathbf{x}}$. Thus, this separation is still meaningful from an algorithmic perspective with the following convention: assign all computations independent of the query ${\mathbf{x}}$ to $g_{\bm{\varphi}}$, and leave the query-dependent computations to $f_{\gamma}$. In this view, $f_{\gamma}$ processes the query while accessing task-specific information ${\mathcal{D}}_{\mathcal{T}}$ only through the structure or constraints imposed by $g_{\bm{\varphi}}$, for example, pooled representations, gradient information, or natural language.

Learning the learner.
So far, we only focused on the modeling assumptions $-$ the form of $f_{\gamma}({\mathbf{x}},g_{\bm{\varphi}}({\mathcal{D}}_{\mathcal{T}}))$ $-$ behind different amortized learners did not address how the optimization problem in [Eq.˜6](https://arxiv.org/html/2510.11471v1#S2.E6) is actually performed. The main complexity comes from iterative natures of $f_{\gamma}$ and $g_{\bm{\varphi}}$ and the presence of higher order gradients if $g_{\bm{\varphi}}$ describes a gradient-based procedure. Multiple approaches tackle this problem through first-order approximations (Finn et al., 2017; Nichol and Schulman, 2018), back-propagation through time (Ha et al., 2016), reinforcement learning (Andrychowicz et al., 2016) or evolutionary strategies (Metz et al., 2022b) depending on the form of $g_{\bm{\varphi}}$.

Figure: Figure 1: Iterative Amortized Inference for parametric, explicit and implicit parameterizations. While the parametric and explicit setup provide task specific information ${\bm{\theta}}$ invariant to the query, the implicit setup instead continually refines a prediction tied to the query.
Refer to caption: x2.png

## 3 Amortized Learning Systems: A Taxonomy

Motivated by the unified formulation of general-purpose amortized models and their ability to efficiently generalize to novel tasks, we introduce a taxonomy categorizing such models into three classes: parametric, explicit, and implicit. While we focus our attention to tasks described through a set of observations ${\mathcal{D}}_{\mathcal{T}}$ rather than language or task descriptors, the same categorization holds in general.

Parametric. We define the class of amortized models with a fixed $f$ and learnable $g_{\bm{\varphi}}$ as parametric amortization. This includes hypernetworks, learned optimizers, and other approaches to parametric inference using ICL. Here, the functional form of the likelihood is fixed, for example a known simulator (Cranmer et al., 2020) or a categorical likelihood with linear mapping, while $g_{\bm{\varphi}}$ is an inference or estimation procedure that discovers the optimal parameters (e.g., linear coefficients) from the observed data. Several works (Mittal et al., 2025b; Reuter et al., 2025; Chang et al., 2024) approximate the posterior distribution over parameters of the likelihood as $g_{\bm{\varphi}}$ while others (Mittal et al., 2025a) investigate the interplay between point-based and distribution-based modeling of $g_{\bm{\varphi}}$(^4^44In general $g_{\bm{\varphi}}$ can represent probability distributions or sampling operations.). Relying on a fixed parametric assumption enables us to leverage gradient information in addition to ${\mathcal{D}}_{\mathcal{T}}$ to infer the parameters of $f$. These inferred parameters can also offer a low-dimensional representation of the task and provide interpretability benefits, especially when $f$ possesses specific, interpretable structure.

Implicit. In contrast, we refer to amortized models with a trainable $f_{\gamma}$ and a fixed $g$ as implicit amortization, which subsumes in-context learning and specific cases of prior fitted networks and conditional neural processes (Nguyen and Grover, 2022; Hollmann et al., 2022). In this setting, a single trained model takes both the query and the set of observations as input and directly learns to model predictions. The function $g$ is typically either the identity mapping or a subsampling mechanism used to manage large datasets. Unlike parametric approaches, the functional form of the output is learned directly, bypassing explicitly inferring any task-specific parameters. It is non-trivial to pinpoint what part of network activations correspond to dataset-specific vs prediction based activations. Some works (Von Oswald et al., 2023; Nichani et al., 2024) demonstrate that, under certain assumptions, implicit models recover algorithmic behaviors such as gradient descent or causal discovery. However, it remains unclear how these findings generalize beyond specific setups, apart from the broader perspective of learning the posterior predictive distribution (Müller et al., 2021; Garg et al., 2022).

Explicit. It is possible to leverage the dimensionality reduction benefits offered by the parametric approach while still utilizing a learnable form for the likelihood which the implicit model affords. This is accomplished with a trainable $f_{\gamma}$, which learns the likelihood form shared across tasks like the laws of physics, as well as $g_{\bm{\varphi}}$, which provides a low-dimensional encapsulation of the entire dataset like the gravitational constant. The explicit models considered in (Mittal et al., 2024; Elmoznino et al., 2024) and neural processes (Garnelo et al., 2018a, b) that first learn an embedding of the dataset can be seen as specific cases of this approach. Similar to parametric methods, it allows using gradient information and provides dimensionality reduction and interpretability benefits; but at the cost of inherent non-stationarity during learning $-$ $f_{\gamma}$ and $g_{\bm{\varphi}}$ are dependent on each other but have to be learned together.

**Table 2: Our experiments on parametric amortization reveal benefits of multiple steps of iterative refinement across ID and OoD (gray columns) evaluation. Top row describes pre-training tasks and the second row evaluation tasks, with error metric. We also see that sole reliance on gradients is often insufficient and leveraging observations directly can lead to improved performance with fewer iterations.**
| Training Tasks $\rightarrow$ | MNIST | FMNIST | ImageNet |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Signal $\downarrow$ | Steps | Lin Reg | MNIST | \columncoloroodFMNIST | FMNIST | \columncoloroodMNIST | ImageNet | \columncoloroodCIFAR10 | \columncoloroodCIFAR100 |
| Grad | 1 | 51.3$\pm 0.3$ | 76.2$\pm 2.0$ | \columncolorood66.7$\pm 1.0$ | 63.8$\pm 1.7$ | \columncolorood79.9$\pm 1.6$ | 88.5$\pm 0.1$ | \columncolorood44.9$\pm 4.5$ | \columncolorood92.5$\pm 1.1$ |
| 5 | 4.1$\pm 0.1$ | 48.6$\pm 1.2$ | \columncolorood51.8$\pm 0.5$ | 41.6$\pm 1.5$ | \columncolorood49.5$\pm 3.0$ | \cellcolorcream 83.1$\pm 0.2$ | \columncolorood17.9$\pm 1.0$ | \columncolorood\cellcolorcream 88.1$\pm 0.8$ |  |
| 10 | 0.5$\pm 0.0$ | 39.7$\pm 0.7$ | \columncolorood43.5$\pm 0.9$ | 38.1$\pm 0.4$ | \columncolorood40.7$\pm 0.8$ | 83.4$\pm 0.1$ | \columncolorood\cellcolorcream 14.7$\pm 0.2$ | \columncolorood\cellcolorcream 88.1$\pm 0.6$ |  |
| Data | 1 | 16.3$\pm 0.2$ | 43.9$\pm 1.3$ | \columncolorood40.5$\pm 1.2$ | 35.9$\pm 1.7$ | \columncolorood38.8$\pm 1.4$ | 90.8$\pm 0.2$ | \columncolorood60.9$\pm 3.5$ | \columncolorood93.9$\pm 1.1$ |
| 5 | 0.5$\pm 0.0$ | 35.5$\pm 0.5$ | \columncolorood34.3$\pm 0.6$ | 30.8$\pm 0.3$ | \columncolorood31.3$\pm 0.5$ | 95.2$\pm 0.1$ | \columncolorood46.9$\pm 1.8$ | \columncolorood96.3$\pm 0.6$ |  |
| 10 | \cellcolorcream 0.3$\pm 0.0$ | 36.6$\pm 0.3$ | \columncolorood35.3$\pm 0.3$ | \cellcolorcream 29.8$\pm 0.3$ | \columncolorood\cellcolorcream 29.1$\pm 0.4$ | 93.2$\pm 0.1$ | \columncolorood28.8$\pm 0.8$ | \columncolorood95.0$\pm 0.8$ |  |
| Grad | 1 | 25.1$\pm 0.3$ | 52.3$\pm 1.5$ | \columncolorood48.2$\pm 0.8$ | 37.7$\pm 1.9$ | \columncolorood41.9$\pm 1.3$ | 96.7$\pm 0.1$ | \columncolorood70.2$\pm 4.2$ | \columncolorood97.2$\pm 0.8$ |
| + | 5 | 0.6$\pm 0.0$ | 39.2$\pm 1.0$ | \columncolorood37.5$\pm 0.3$ | 32.0$\pm 0.7$ | \columncolorood33.4$\pm 0.5$ | 93.8$\pm 0.1$ | \columncolorood42.1$\pm 2.4$ | \columncolorood94.8$\pm 0.7$ |
| Data | 10 | 0.4$\pm 0.0$ | \cellcolorcream 32.0$\pm 0.5$ | \columncolorood\cellcolorcream 32.0$\pm 0.2$ | 30.5$\pm 0.2$ | \columncolorood30.3$\pm 0.3$ | 92.5$\pm 0.1$ | \columncolorood24.8$\pm 1.0$ | \columncolorood94.1$\pm 0.4$ |

## 4 Iterative Amortized Inference

A fundamental limitation of the existing approaches is their inability to leverage large-scale dataset as conditioning $-$ they are either restricted by context length or rely solely on low-dimensional pooling operations or gradients which can be quite restrictive. Analogously, non-amortized learners like gradient descent resolve this scalability issue through a stochastic mini-batch framework which iteratively refines an approximate solution based on a subset of observations, not the whole dataset. Note that this iterative refinement is Markovian(^5^55Certain second-order gradient-based learners are non Markovian, or Markovian in an augmented state space.) and greedy, i.e. the update given current state is independent of history and leads to immediate single-step improvement which could be sub-optimal over multiple steps.

Taking inspiration from SGD and the connection between amortized meta-learners and optimization routines (Elmoznino et al., 2024), we extend existing approaches using an iterative refinement approach; instead of directly modeling predictions, embeddings or parameters, we iteratively refine the output from the previous step greedily using mini-batches as input to a trained sequence model.

For parametric and explicit amortization, we achieve this by considering $g_{\bm{\varphi}}$ as an iterative application of a learned sequence model $h_{\bm{\varphi}}$ on different subsampled training batches ${\mathcal{B}}^{(i)}_{\text{train}}\subset{\mathcal{D}}^{\text{train}}_{\mathcal{T}}$, starting from a learned initialization ${\bm{\theta}}^{(0)}$.(^6^66In theory, this should learn the appropriate prior for the class of tasks for the parametric setup.)
This model $h_{\bm{\varphi}}$ takes the current state ${\bm{\theta}}^{(i)}$ and a mini-batch ${\mathcal{B}}^{(i)}_{\text{train}}$ as input and returns a refined state ${\bm{\theta}}^{(i+1)}$ which decreases validation loss.

$$ $\displaystyle{\bm{\theta}}^{(0)}\overset{h_{\bm{\varphi}}(\cdot,{\mathcal{B}}_{\text{train}}^{(0)})}{\xrightarrow{\hskip 20.00003pt}}{\bm{\theta}}^{(1)}\overset{h_{\bm{\varphi}}(\cdot,{\mathcal{B}}_{\text{train}}^{(1)})}{\xrightarrow{\hskip 20.00003pt}}\ldots\overset{h_{\bm{\varphi}}(\cdot,{\mathcal{B}}_{\text{train}}^{(k-1)})}{\xrightarrow{\hskip 20.00003pt}}{\bm{\theta}}^{(k)}\eqqcolon{\bm{\theta}}_{\mathcal{T}}$ (7) $$

denotes a $k$-step iterative refinement procedure for parametric and explicit models. While learned optimizers already utilize a mini-batch approach, they only consider parametric cases with fixed $f$ and the information about observations is fed to $h_{\bm{\varphi}}$ solely through gradients. In particular, they model
$h_{\bm{\varphi}}({\bm{\theta}},{\mathcal{B}})\coloneqq\text{Transformer}_{{\bm{\varphi}}}\Big({\bm{\theta}},\nabla_{\bm{\theta}}\sum_{({\mathbf{x}},{\mathbf{y}})\in{\mathcal{B}}}{\mathcal{L}}\left({\mathbf{y}},f_{\gamma}\left({\mathbf{x}},{\bm{\theta}}\right)\right)\big\rvert_{{\bm{\theta}}}\Big)$
where the learned model takes a sequence of parameter updates and gradients. Such methods can only recover gradient-based inference routines and cannot model more general optimization schemes (e.g. closed-form linear regression solution) as the mapping from observations to gradient is non-invertible.

In parametric and explicit methods, ${\bm{\theta}}$s can be seen as recurrent memory compressing useful information of the current task from ${\mathcal{D}}_{\mathcal{T}}$ to make predictions for new ${\mathbf{x}}$, i.e. prediction for ${\mathbf{x}}$ given ${\bm{\theta}}_{\mathcal{T}}$ is conditionally independent of ${\mathcal{D}}_{\mathcal{T}}$. In contrast, the implicit formulation does not consider task-specific memory as its recurrent state is always tied to the query. In this case, $g_{\bm{\varphi}}$ simply subsamples mini-batches while $f_{\gamma}$ is a recurrent application of a transformer $r_{\gamma}$ with weights $\gamma$

$$ $\displaystyle\hat{{\mathbf{y}}}^{(0)}\overset{r_{\gamma}([{\mathbf{x}},\hat{{\mathbf{y}}}^{(0)}],{\mathcal{B}}^{(0)}_{\text{train}})}{\xrightarrow{\hskip 40.00006pt}}\hat{{\mathbf{y}}}^{(1)}\overset{r_{\gamma}([{\mathbf{x}},\hat{{\mathbf{y}}}^{(1)}],{\mathcal{B}}^{(1)}_{\text{train}})}{\xrightarrow{\hskip 40.00006pt}}\ldots\overset{r_{\gamma}([{\mathbf{x}},\hat{{\mathbf{y}}}^{(k-1)}],{\mathcal{B}}^{(k-1)}_{\text{train}})}{\xrightarrow{\hskip 40.00006pt}}\hat{{\mathbf{y}}}^{(k)},$ (8) $$

where $k$ is the number of iterative refinement steps for implicit models, $[\cdot,\cdot]$ describes a token and $\hat{{\mathbf{y}}}$ the query-specific predictions which form the recurrent states in this framework and are sequentially updated with $\hat{{\mathbf{y}}}^{(0)}$ being a learned initialization(^7^77In theory, this initialization should learn the marginal (unconditional) distribution over predictions.). Here, $f_{\gamma}$ gets the current prediction state $\hat{{\mathbf{y}}}^{(i)}$ and a mini-batch ${\mathcal{B}}^{(i)}_{\text{train}}$ as input and provides a refined prediction $\hat{{\mathbf{y}}}^{(i+1)}$. Since the implicit model never exposes task specific parameters ${\bm{\theta}}_{\mathcal{T}}$, it is non-trivial to leverage gradient information.

**Table 3: On explicit amortization we consistently see improved performance with increasing number of steps (OoD evaluation in gray columns) with error as metric. Here, we see that reliance on gradients is essential for more complex problems and in particular, explicit models struggled to scale to ImageNet and could only obtain random chance performance.**
| Training Tasks $\rightarrow$ | MNIST | FMNIST |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Signal $\downarrow$ | Steps | Lin Reg | MNIST | \columncoloroodFMNIST | FMNIST | \columncoloroodMNIST |
| Grad | 1 | 55.5$\pm 0.5$ | 90.4$\pm 0.7$ | \columncolorood90.8$\pm 0.4$ | 88.2$\pm 0.3$ | \columncolorood89.3$\pm 0.3$ |
| 5 | 9.4$\pm 0.1$ | 86.9$\pm 0.1$ | \columncolorood84.4$\pm 0.5$ | 70.8$\pm 0.5$ | \columncolorood76.4$\pm 0.4$ |  |
| 10 | 2.9$\pm 0.0$ | 81.4$\pm 0.5$ | \columncolorood78.9$\pm 0.9$ | 56.4$\pm 0.7$ | \columncolorood65.7$\pm 0.4$ |  |
| Data | 1 | 19.7$\pm 0.4$ | 89.2$\pm 1.0$ | \columncolorood90.0$\pm 0.1$ | 90.0$\pm 0.1$ | \columncolorood90.3$\pm 0.2$ |
| 5 | 2.6$\pm 0.0$ | 88.7$\pm 0.0$ | \columncolorood90.0$\pm 0.0$ | 90.0$\pm 0.0$ | \columncolorood90.2$\pm 0.0$ |  |
| 10 | 2.4$\pm 0.0$ | 88.8$\pm 0.2$ | \columncolorood90.0$\pm 0.0$ | 90.1$\pm 0.1$ | \columncolorood90.2$\pm 0.0$ |  |
| Grad | 1 | 24.4$\pm 0.4$ | 72.5$\pm 0.7$ | \columncolorood62.9$\pm 0.8$ | 51.1$\pm 1.2$ | \columncolorood61.2$\pm 0.6$ |
| + | 5 | 2.6$\pm 0.0$ | \cellcolorcream 67.1$\pm 0.2$ | \columncolorood\cellcolorcream 58.3$\pm 0.3$ | 47.5$\pm 0.3$ | \columncolorood57.7$\pm 0.2$ |
| Data | 10 | \cellcolorcream 2.1$\pm 0.0$ | 75.1$\pm 0.1$ | \columncolorood69.5$\pm 0.1$ | \cellcolorcream 45.1$\pm 0.3$ | \columncolorood\cellcolorcream 54.4$\pm 0.1$ |

Motivated by SGD, we learn ${\bm{\varphi}}$ and $\gamma$ through a greedy strategy where training is done solely on single-step improvements for all three paradigms. The respective computational graphs are provided below, where we prevent gradients to flow back from ${\bm{\theta}}^{(i)}$ or $\hat{{\mathbf{y}}}^{(i)}$

$$ $\displaystyle{\bm{\theta}}^{(i)}\overset{h_{\bm{\varphi}}(\cdot,{\mathcal{B}})}{\xrightarrow{\hskip 30.00005pt}}{\bm{\theta}}^{(i+1)}\to{\mathcal{L}}({\mathbf{y}},f_{\gamma}({\mathbf{x}},{\bm{\theta}}^{(i+1)}))\quad\text{or}\quad\hat{{\mathbf{y}}}^{(i)}\overset{r_{\gamma}([{\mathbf{x}},\cdot],{\mathcal{B}})}{\xrightarrow{\hskip 30.00005pt}}\hat{{\mathbf{y}}}^{(i+1)}\to{\mathcal{L}}({\mathbf{y}},\hat{{\mathbf{y}}}^{(i+1)})$ (9) $$

Here, an outer loop is used to obtain the states at the $i^{th}$ step $-$ ${\bm{\theta}}^{(i)}$ and $\hat{{\mathbf{y}}}^{(i)}$ respectively. Note that in this setup, $({\mathbf{x}},{\mathbf{y}})\in{\mathcal{D}}_{\mathcal{T}}^{\text{valid}}$ while ${\mathcal{B}}\subseteq{\mathcal{D}}_{\mathcal{T}}^{\text{train}}$, i.e. we are amortizing a learner to minimize validation loss. We refer to [Figure˜1](https://arxiv.org/html/2510.11471v1#S2.F1) for an illustration of the proposed framework and to [Section˜8](https://arxiv.org/html/2510.11471v1#S8) for a walk-through with an example. For evaluation, we use training and evaluation splits from completely new tasks ${\mathcal{T}}$ and thus test for generalization to new problems. All three frameworks $-$ parametric, explicit and implicit $-$ modify causal masking to provide efficient and parallelizable computation of ${\bm{\theta}}$ and $\hat{{\mathbf{y}}}$ with mini-batches having varied number of observations, detailed in [Section˜11](https://arxiv.org/html/2510.11471v1#S11).

## 5 Experiments

We evaluate our proposed greedy iterative refinement approach as well as the pros and cons of using the two modalities $-$ observations or gradients $-$ through a suite of experiments. We consider a range of predictive and generative modeling tasks listed below and detailed in [Section˜9](https://arxiv.org/html/2510.11471v1#S9) to assess both in-distribution (ID) and out-of-distribution (OoD) performance under challenging task variations. Learned $g_{\bm{\varphi}},f_{\gamma}$ are parameterized as transformers with maximum sequence length of $100$ observations, following Kirsch et al. (2022) (details in [Section˜11](https://arxiv.org/html/2510.11471v1#S11)).

Tasks. Our tasks include

- •
Linear Regression. ${\mathcal{T}}$ is defined by ground-truth linear coefficients ${\mathbf{w}}_{\mathcal{T}}\in\mathbb{R}^{100}$ which control the parameters of $p_{\mathcal{T}}({\mathbf{y}}|{\mathbf{x}})$.
- •
MNIST / FashionMNIST Classification. ${\mathcal{T}}$ involves a random projection $W_{\mathcal{T}}\in\mathbb{R}^{784\times 100}$ applied to image pixels and a random label mapping (e.g. digits $4$ are mapped to label $7$; even digits grouped together, etc.) which preserves the semantic structure (Kirsch et al., 2022). We also consider OoD settings where models are trained on MNIST and evaluated on FashionMNIST, and vice versa.
- •
ImageNet Classification. ${\mathcal{T}}$ samples images from $100$ different ImageNet classes and randomly regrouping them into a maximum $100$-way classification task, similar to above. Evaluation is done on ImageNet validation data as well as OoD tasks like CIFAR-10 and CIFAR-100. The images are embedded using Dino-v2 (Oquab et al., 2023) before feeding to the transformer.
- •
Topological Order Prediction. Following (Scetbon et al., 2024), each task involves sampling a structural causal model (SCM) and inferring its topological order using only observational data.
- •
Mixture of Gaussian (MoG) Generation. Each task is defined by a MoG with the number of mixture components and corresponding means sampled randomly per task. We apply the flow-matching framework (Lipman et al., 2022; Tong et al., 2023; Albergo et al., 2023) to learn the conditional vector field given task data ${\mathcal{D}}_{\mathcal{T}}$, and evaluate samples generated conditioned on novel densities as context unseen during training.
- •
Alphabet Generation. Following (Atanackovic et al., 2024), we use a dataset of alphabets with different scales and rotations. The task is to leverage the conditioning information to draw samples from the underlying density, with OoD settings corresponding to alphabets unseen during training.

Metrics. For predictive tasks, we look at $l_{2}$ loss for regression and error for classification, while for generative modeling problems we consider the $2$-Wasserstein $({\mathcal{W}}_{2})$ and $1$-Wasserstein $({\mathcal{W}}_{1})$ distance between samples from the true and generated distribution. We refer to [Section˜10](https://arxiv.org/html/2510.11471v1#S10) for details.

**Table 4: Our experiments on implicit amortization show consistent improvements in performance with increasing number of steps across a wide range of predictive tasks, with error as evaluation metric.**
| Training Tasks $\rightarrow$ | MNIST | FMNIST | ImageNet |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Steps | Lin Reg | MNIST | \columncoloroodFMNIST | FMNIST | \columncoloroodMNIST | ImageNet | \columncoloroodCIFAR10 | \columncoloroodCIFAR100 |
| 1 | 14.8$\pm 0.1$ | 24.9$\pm 0.4$ | \columncolorood31.9$\pm 0.4$ | 22.8$\pm 0.1$ | \columncolorood29.1$\pm 0.4$ | 43.0$\pm 0.7$ | \columncolorood19.4$\pm 0.5$ | \columncolorood73.9$\pm 0.1$ |
| 5 | 6.2$\pm 0.0$ | 12.4$\pm 0.1$ | \columncolorood26.1$\pm 0.2$ | 18.8$\pm 0.1$ | \columncolorood16.9$\pm 0.2$ | 13.3$\pm 0.2$ | \columncolorood\cellcolorcream 15.0$\pm 0.1$ | \columncolorood59.4$\pm 0.2$ |
| 10 | \cellcolorcream 4.7$\pm 0.0$ | \cellcolorcream 9.7$\pm 0.1$ | \columncolorood\cellcolorcream 24.2$\pm 0.2$ | \cellcolorcream 17.8$\pm 0.1$ | \columncolorood\cellcolorcream 15.7$\pm 0.2$ | \cellcolorcream 12.3$\pm 0.1$ | \columncolorood15.9$\pm 0.3$ | \columncolorood\cellcolorcream 54.2$\pm 0.4$ |

**Table 5: We evaluate implicit models on topological order prediction, where ID uses SCMs with linear (LIN IN) / non-linear (RFF IN) mechanisms for 20 node graphs and OoD (gray columns) changes function parameters (LIN OUT / RFF OUT) or graph size (50 nodes), with classification error as metric.**
|  | 20 Nodes | 50 Nodes |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Steps | LIN IN | RFF IN | \columncoloroodLIN OUT | \columncoloroodRFF OUT | \columncoloroodLIN IN | \columncoloroodRFF IN | \columncoloroodLIN OUT | \columncoloroodRFF OUT |
| 1 | 0.52 $\pm 0.05$ | 0.43 $\pm 0.04$ | \columncolorood\cellcolorcream 0.64 $\pm 0.04$ | \columncolorood0.51 $\pm 0.06$ | \columncolorood0.7 $\pm 0.05$ | \columncolorood0.66 $\pm 0.04$ | \columncolorood0.78 $\pm 0.03$ | \columncolorood0.69 $\pm 0.08$ |
| 5 | 0.52 $\pm 0.06$ | \cellcolorcream 0.35 $\pm 0.05$ | \columncolorood0.68 $\pm 0.08$ | \columncolorood\cellcolorcream 0.33 $\pm 0.06$ | \columncolorood0.66 $\pm 0.05$ | \columncolorood0.56 $\pm 0.04$ | \columncolorood\cellcolorcream 0.74 $\pm 0.04$ | \columncolorood\cellcolorcream 0.67 $\pm 0.05$ |
| 10 | \cellcolorcream 0.44 $\pm 0.11$ | 0.37 $\pm 0.07$ | \columncolorood0.66 $\pm 0.08$ | \columncolorood0.61 $\pm 0.05$ | \columncolorood\cellcolorcream 0.64 $\pm 0.06$ | \columncolorood\cellcolorcream 0.52 $\pm 0.04$ | \columncolorood0.77 $\pm 0.03$ | \columncolorood0.73 $\pm 0.03$ |

Baselines. For parametric amortization, we consider transformer-based hypernetworks and learned optimizers as baselines. The former provides single-step amortization using only data, while the latter is a multi-step system that relies solely on gradients as input. For the explicit models, the $1$-step approach solely based on data is equivalent to the explicit model defined in Mittal et al. (2024) as well as to conditional neural processes (Garnelo et al., 2018a). Finally, for the implicit setup, the $1$-step approach boils down to ICL
(Mittal et al., 2024; Garg et al., 2022; Müller et al., 2021; Von Oswald et al., 2023).

### 5.1 Iterative amortized inference provides consistent improvements

Parametric. We first test our framework on parametric models $-$ $f$ is a linear predictor and $g_{\bm{\varphi}}$ is trained to infer its parameters ${\bm{\theta}}_{\mathcal{T}}$. We use a transformer $g_{\bm{\varphi}}$ which amortizes by conditioning on ${\mathcal{D}}_{\mathcal{T}}$, either through gradient signal or observations, or both. [Table˜2](https://arxiv.org/html/2510.11471v1#S3.T2) demonstrates that iterative amortization consistent improves performance with increasing number of steps across various tasks and input signals $-$ i.e. gradients, observations, or both, consistently outperforming the baselines. The amortized model also generalizes OoD: for e.g. pre-training on ImageNet classification can lead to good performance on CIFAR-10 at inference, based solely on in-context examples.

Explicit. Next, we consider the explicit model $-$ $f_{\gamma}$ is a trained MLP taking a query ${\mathbf{x}}$ and latent ${\bm{\theta}}_{\mathcal{T}}$ as input, while $g_{\bm{\varphi}}$ is a transformer inferring the latent from ${\mathcal{D}}_{\mathcal{T}}$. Similar to the parametric case, [Table˜3](https://arxiv.org/html/2510.11471v1#S4.T3) showcases consistent improvements of our proposed approach. In contrast to the parametric experiments, we see an increased importance of gradient information in inferring the right latents. Since there is inherent non-stationarity in the optimization procedure we hypothesize that gradients provide a clearer signal for the inference mechanism $g_{\bm{\varphi}}$.

Figure: Ground Truth
Refer to caption: x3.png

**Table 6: We see improved sample quality (under Wasserstein metrics) using implicit amortization with increasing number of steps for generative modeling over OoD alphabets and new GMMs.**
|  | Alphabets | GMM |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Steps | ID | \cellcolorood OoD | 2-dimensional | 5-dimensional |  |  |  |  |
|  | ${\mathcal{W}}_{2}$ | ${\mathcal{W}}_{1}$ | \columncolorood${\mathcal{W}}_{2}$ | \columncolorood${\mathcal{W}}_{1}$ | ${\mathcal{W}}_{2}$ | ${\mathcal{W}}_{1}$ | ${\mathcal{W}}_{2}$ | ${\mathcal{W}}_{1}$ |
| 1 | \cellcolorcream 0.28$\pm 0.00$ | \cellcolorcream 0.20$\pm 0.00$ | \columncolorood0.78$\pm 0.01$ | \columncolorood0.64$\pm 0.00$ | 1.12$\pm 0.01$ | 0.58$\pm 0.00$ | 3.27$\pm 0.01$ | 1.82$\pm 0.01$ |
| 5 | 0.31$\pm 0.00$ | 0.22$\pm 0.00$ | \columncolorood0.72$\pm 0.01$ | \columncolorood0.58$\pm 0.00$ | 0.90$\pm 0.01$ | 0.38$\pm 0.00$ | 2.50$\pm 0.02$ | 1.05$\pm 0.01$ |
| 10 | 0.29$\pm 0.00$ | 0.21$\pm 0.00$ | \columncolorood\cellcolorcream 0.67$\pm 0.01$ | \columncolorood\cellcolorcream 0.54$\pm 0.01$ | \cellcolorcream 0.81$\pm 0.02$ | \cellcolorcream 0.32$\pm 0.01$ | \cellcolorcream 2.37$\pm 0.02$ | \cellcolorcream 0.95$\pm 0.01$ |

Implicit. Finally, we consider implicit parameterization where $f_{\gamma}$ jointly leverages both training dataset ${\mathcal{D}}_{\mathcal{T}}$ as well as the query ${\mathbf{x}}$ to model predictions ${\mathbf{y}}$. Our experiments in [Tables˜4](https://arxiv.org/html/2510.11471v1#S5.T4) and [5](https://arxiv.org/html/2510.11471v1#S5.T5) indicate that iterative amortization substantially outperforms competitive baselines and leads to improved performance with increasing number of steps. In addition to predictive tasks, we see similar trends on generative modeling where the task is to model the underlying distribution described through the context examples. The amortized model is trained to infer the conditional vector field which interpolates a path between standard normal and the observed distribution, conditioned on ${\mathcal{D}}_{\mathcal{T}}$ (Atanackovic et al., 2024; Chang et al., 2024). [Tables˜6](https://arxiv.org/html/2510.11471v1#S5.T6) and [2](https://arxiv.org/html/2510.11471v1#S5.F2) showcase improved ability of modeling the underlying distributions defined by the context with more steps(^8^88Here, more steps imply steps to infer the vector field $v_{t}(\cdot|{\mathcal{D}}_{\mathcal{T}})$ for each $t$, and not more integration steps.).

We refer to [Section˜12](https://arxiv.org/html/2510.11471v1#S12) for additional comparisons to transformers without causal masking.

### 5.2 Analysis

Gradient signal is often sub-optimal. When $g_{\bm{\varphi}}$ solely leverages gradient information, we recover the space of learned optimizers. For relatively lower dimensional problems, we see that the space of sole gradient-based learned optimizers is suboptimal and additionally leveraging the observations can lead to substantial improvements. Alternatively, in higher dimensions since the hypothesis space is larger, it is harder for the model to infer the right parameters. In such cases, gradients provide more reliable signals to explore this space, especially when there are fewer observations.

Effect of larger history information. We relax the Markovian assumption from the earlier experiments and study the impact of including multiple past states and gradient information on generalization to novel tasks. The augmented system is now defined as $h_{\bm{\varphi}}:{\bm{\theta}}_{t-k:t},\nabla_{t-k:t},{\mathcal{B}}_{t}\to{\bm{\theta}}_{t+1}$, where $k$ is a hyperparameter for the number of past states. Our results in [Fig.˜3](https://arxiv.org/html/2510.11471v1#S5.F3) showcase more history information does not impact results much, except in the case of pre-training on FashionMNIST tasks.

Explicit parameterization is sub-optimal. Importantly, parametric modeling outperforms the explicit setup highlighting the difficulty of jointly learning $f_{\gamma}$ and $g_{\bm{\varphi}}$. This is surprising as the class of solutions expressed by the former form a subset of the latter. We attribute this suboptimality to the optimization process as $f_{\gamma}$ and $g_{\bm{\varphi}}$ are linked in a complicated, non-stationary manner, i.e. if the prediction function $f_{\gamma}$ changes then its inference mechanism $g_{\bm{\varphi}}$ also needs to adapt, and vice versa.

State parameterization in implicit amortization. We analyze the design choices associated with the state that persists across iterations to be $-$ (a) Pre-MLP: high-level latent variables tied to the query, (b) Logits: current prediction before undergoing parameter-free normalization like softmax, or (c) Softmax: current predictions after softmax. [Fig.˜4](https://arxiv.org/html/2510.11471v1#S5.F4) shows that modeling the state at logits leads to best performance, and in general shows the importance of being close to prediction for greedy refinement.

Runtime Efficiency of Amortization. We note that for a $K$-step iterative method with mini-batch size $B$ leads to a runtime cost of $O(KB^{2})$. In contrast, a one-step model runs in ${\mathcal{O}}(B^{2})$ but uses $K$ times less data. When given access to the same amount of data, the one-step model incurs runtime of ${\mathcal{O}}((KB)^{2})$. Thus, iterative amortization is $K$ times more efficient than one-step models when processing the same amount of data. A similar argument applies to memory usage, making the iterative approach more scalable for large tasks. We next compare the amortized parametric model to an off-the-shelf optimizer (Adam) with standard normal initialization as well as MAML (Finn et al., 2017): Adam achieves a loss of $112$ while MAML achieves $44.7$ in $100$ steps, both taking approximately $0.396$ seconds. In contrast, a single-step amortized parametric model can achieve the performance of $\approx 16$ with inference taking only 0.05 seconds, thus illustrating the importance of amortized learners which do not rely on off-the-shelf optimizers at inference.

Figure: Figure 3: We analyze the benefits, or lack thereof, of leveraging multiple past states and gradients for parametric amortization, when we use both observations and gradients as conditional inputs.
Refer to caption: x9.png

Figure: Figure 4: Our ablations reveal that carrying over logits as the recurrent state across iterations in implicit models outperforms other state representations, with softmax outputs performing comparably.
Refer to caption: x10.png

## 6 Conclusion

Amortized learning serves as a powerful paradigm for rapid generalization through reuse of computations. By introducing a unified framework and taxonomy, we clarify how different approaches internalize or externalize task adaptation. This perspective not only highlights shared principles but also exposes common bottlenecks: the limited scalability of existing methods in handling large task datasets as well as suboptimality when solely leveraging gradients with fixed functional mappings. To overcome this, we propose iterative amortization, a novel strategy that incrementally refines solutions through mini-batch updates, effectively merging the strengths of optimization-based and forward-pass approaches consistently across parametric, explicit and implicit frameworks for predictive and generative ID and OoD tasks. This framework paves the way for scalably exploring more sophisticated forms of persistent memory across mini-batches for robust task adaptation in increasingly complex, non-iid environments. Additionally, future work involves incorporating richer learning frameworks to optimize [Eq.˜5](https://arxiv.org/html/2510.11471v1#S2.E5) beyond greedy refinement, such as evolutionary algorithms or reinforcement learning.

## Acknowledgements

GL acknowledges support from NSERC Discovery Grant RGPIN-2018-04821, the Canada Research Chair in Neural Computations and Interfacing, and a Canada-CIFAR AI Chair.
DM acknowledges support via
FRQNT doctoral scholarship ([https://doi.org/10.69777/354785](https://doi.org/10.69777/354785)) for his graduate studies. SM acknowledges funding from FRQNT Doctoral Fellowship ([https://doi.org/10.69777/372208](https://doi.org/10.69777/372208)).

## References

- Achiam et al. (2023)
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al.
Gpt-4 technical report.
*arXiv preprint arXiv:2303.08774*, 2023.
- Albergo et al. (2023)
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden.
Stochastic interpolants: A unifying framework for flows and diffusions.
*arXiv preprint arXiv:2303.08797*, 2023.
- Amos et al. (2023)
Brandon Amos et al.
Tutorial on amortized optimization.
*Foundations and Trends® in Machine Learning*, 16(5):592–732, 2023.
- Andrychowicz et al. (2016)
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas.
Learning to learn by gradient descent by gradient descent.
*Advances in neural information processing systems*, 29, 2016.
- Antoniou et al. (2018)
Antreas Antoniou, Harrison Edwards, and Amos Storkey.
How to train your maml.
*arXiv preprint arXiv:1810.09502*, 2018.
- Atanackovic et al. (2024)
Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo J Lee, Yoshua Bengio, Alexander Tong, and Kirill Neklyudov.
Meta flow matching: Integrating vector fields on the wasserstein manifold.
*arXiv preprint arXiv:2408.14608*, 2024.
- Bahdanau et al. (2014)
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio.
Neural machine translation by jointly learning to align and translate.
*arXiv preprint arXiv:1409.0473*, 2014.
- Bai et al. (2023)
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei.
Transformers as statisticians: Provable in-context learning with in-context algorithm selection.
*Advances in neural information processing systems*, 36:57125–57211, 2023.
- Barabási and Albert (1999)
Albert-László Barabási and Réka Albert.
Emergence of scaling in random networks.
*science*, 286(5439):509–512, 1999.
- Brown et al. (2020)
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.
Language models are few-shot learners.
*Advances in neural information processing systems*, 33:1877–1901, 2020.
- Chang et al. (2024)
Paul E Chang, Nasrulloh Loka, Daolang Huang, Ulpu Remes, Samuel Kaski, and Luigi Acerbi.
Amortized probabilistic conditioning for optimization, simulation and inference.
*arXiv preprint arXiv:2410.15320*, 2024.
- Chauhan et al. (2024)
Vinod Kumar Chauhan, Jiandong Zhou, Ping Lu, Soheila Molaei, and David A Clifton.
A brief review of hypernetworks in deep learning.
*Artificial Intelligence Review*, 57(9):250, 2024.
- Cranmer et al. (2020)
Kyle Cranmer, Johann Brehmer, and Gilles Louppe.
The frontier of simulation-based inference.
*Proceedings of the National Academy of Sciences*, 117(48):30055–30062, 2020.
- Dong et al. (2022)
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Jingyuan Ma, Rui Li, Heming Xia, Jingjing Xu, Zhiyong Wu, Tianyu Liu, et al.
A survey on in-context learning.
*arXiv preprint arXiv:2301.00234*, 2022.
- Efrat and Levy (2020)
Avia Efrat and Omer Levy.
The turking test: Can language models understand instructions?
*arXiv preprint arXiv:2010.11982*, 2020.
- Elmoznino et al. (2024)
Eric Elmoznino, Tom Marty, Tejas Kasetty, Leo Gagnon, Sarthak Mittal, Mahan Fathi, Dhanya Sridhar, and Guillaume Lajoie.
In-context learning and occam’s razor.
*arXiv preprint arXiv:2410.14086*, 2024.
- Erdos and Renyi (1959)
P Erdos and A Renyi.
On random graphs i.
*Publ. math. debrecen*, 6(290-297):18, 1959.
- Finn et al. (2017)
Chelsea Finn, Pieter Abbeel, and Sergey Levine.
Model-agnostic meta-learning for fast adaptation of deep networks.
In *International conference on machine learning*, pages 1126–1135. PMLR, 2017.
- Gaier and Ha (2019)
Adam Gaier and David Ha.
Weight agnostic neural networks.
*Advances in neural information processing systems*, 32, 2019.
- Garg et al. (2022)
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant.
What can transformers learn in-context? a case study of simple function classes.
*Advances in Neural Information Processing Systems*, 35:30583–30598, 2022.
- Garnelo et al. (2018a)
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Ali Eslami.
Conditional neural processes.
In *International conference on machine learning*, pages 1704–1713. PMLR, 2018a.
- Garnelo et al. (2018b)
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh.
Neural processes.
*arXiv preprint arXiv:1807.01622*, 2018b.
- Ha et al. (2016)
David Ha, Andrew Dai, and Quoc V Le.
Hypernetworks.
*arXiv preprint arXiv:1609.09106*, 2016.
- Holland et al. (1983)
Paul W Holland, Kathryn Blackmond Laskey, and Samuel Leinhardt.
Stochastic blockmodels: First steps.
*Social networks*, 5(2):109–137, 1983.
- Hollmann et al. (2022)
Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter.
Tabpfn: A transformer that solves small tabular classification problems in a second.
*arXiv preprint arXiv:2207.01848*, 2022.
- Jia et al. (2016)
Xu Jia, Bert De Brabandere, Tinne Tuytelaars, and Luc V Gool.
Dynamic filter networks.
*Advances in neural information processing systems*, 29, 2016.
- Kingma et al. (2013)
Diederik P Kingma, Max Welling, et al.
Auto-encoding variational bayes, 2013.
- Kirsch et al. (2022)
Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz.
General-purpose in-context learning by meta-learning transformers.
*arXiv preprint arXiv:2212.04458*, 2022.
- Knyazev et al. (2024)
Boris Knyazev, Abhinav Moudgil, Guillaume Lajoie, Eugene Belilovsky, and Simon Lacoste-Julien.
Accelerating training with neuron interaction and nowcasting networks.
*arXiv preprint arXiv:2409.04434*, 2024.
- Krueger et al. (2017)
David Krueger, Chin-Wei Huang, Riashat Islam, Ryan Turner, Alexandre Lacoste, and Aaron Courville.
Bayesian hypernetworks.
*arXiv preprint arXiv:1710.04759*, 2017.
- Lester et al. (2021)
Brian Lester, Rami Al-Rfou, and Noah Constant.
The power of scale for parameter-efficient prompt tuning.
*arXiv preprint arXiv:2104.08691*, 2021.
- Li and Malik (2017)
Ke Li and Jitendra Malik.
Learning to optimize neural nets.
*arXiv preprint arXiv:1703.00441*, 2017.
- Li and Liang (2021)
Xiang Lisa Li and Percy Liang.
Prefix-tuning: Optimizing continuous prompts for generation.
*arXiv preprint arXiv:2101.00190*, 2021.
- Li et al. (2023)
Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak.
Transformers as algorithms: Generalization and stability in in-context learning.
In *International conference on machine learning*, pages 19565–19594. PMLR, 2023.
- Li et al. (2017)
Zhenguo Li, Fengwei Zhou, Fei Chen, and Hang Li.
Meta-sgd: Learning to learn quickly for few-shot learning.
*arXiv preprint arXiv:1707.09835*, 2017.
- Lipman et al. (2022)
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le.
Flow matching for generative modeling.
*arXiv preprint arXiv:2210.02747*, 2022.
- Liu et al. (2021)
Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Lam Tam, Zhengxiao Du, Zhilin Yang, and Jie Tang.
P-tuning v2: Prompt tuning can be comparable to fine-tuning universally across scales and tasks.
*arXiv preprint arXiv:2110.07602*, 2021.
- Lorch et al. (2022)
Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, and Bernhard Schölkopf.
Amortized inference for causal structure learning.
*Advances in Neural Information Processing Systems*, 35:13104–13118, 2022.
- Metz et al. (2019)
Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein.
Understanding and correcting pathologies in the training of learned optimizers.
In *International Conference on Machine Learning*, pages 4556–4565. PMLR, 2019.
- Metz et al. (2022a)
Luke Metz, C Daniel Freeman, James Harrison, Niru Maheswaranathan, and Jascha Sohl-Dickstein.
Practical tradeoffs between memory, compute, and performance in learned optimizers.
In *Conference on Lifelong Learning Agents*, pages 142–164. PMLR, 2022a.
- Metz et al. (2022b)
Luke Metz, James Harrison, C Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, et al.
Velo: Training versatile learned optimizers by scaling up.
*arXiv preprint arXiv:2211.09760*, 2022b.
- Min et al. (2022)
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer.
Rethinking the role of demonstrations: What makes in-context learning work?
*arXiv preprint arXiv:2202.12837*, 2022.
- Mishra et al. (2021)
Swaroop Mishra, Daniel Khashabi, Chitta Baral, Yejin Choi, and Hannaneh Hajishirzi.
Reframing instructional prompts to gptk’s language.
*arXiv preprint arXiv:2109.07830*, 2021.
- Mittal et al. (2024)
Sarthak Mittal, Eric Elmoznino, Leo Gagnon, Sangnie Bhardwaj, Dhanya Sridhar, and Guillaume Lajoie.
Does learning the right latent variables necessarily improve in-context learning?
*arXiv preprint arXiv:2405.19162*, 2024.
- Mittal et al. (2025a)
Sarthak Mittal, Yoshua Bengio, Nikolay Malkin, and Guillaume Lajoie.
In-context parametric inference: Point or distribution estimators?
*arXiv preprint arXiv:2502.11617*, 2025a.
- Mittal et al. (2025b)
Sarthak Mittal, Niels Leif Bracher, Guillaume Lajoie, Priyank Jaini, and Marcus Brubaker.
Amortized in-context bayesian posterior estimation.
*arXiv preprint arXiv:2502.06601*, 2025b.
- Müller et al. (2021)
Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter.
Transformers can do bayesian inference.
*arXiv preprint arXiv:2112.10510*, 2021.
- Munkhdalai and Yu (2017)
Tsendsuren Munkhdalai and Hong Yu.
Meta networks.
In *International conference on machine learning*, pages 2554–2563. PMLR, 2017.
- Neal et al. (2018)
Brady Neal, Sarthak Mittal, Aristide Baratin, Vinayak Tantia, Matthew Scicluna, Simon Lacoste-Julien, and Ioannis Mitliagkas.
A modern take on the bias-variance tradeoff in neural networks.
*arXiv preprint arXiv:1810.08591*, 2018.
- Nguyen and Grover (2022)
Tung Nguyen and Aditya Grover.
Transformer neural processes: Uncertainty-aware meta learning via sequence modeling.
*arXiv preprint arXiv:2207.04179*, 2022.
- Nichani et al. (2024)
Eshaan Nichani, Alex Damian, and Jason D Lee.
How transformers learn causal structure with gradient descent.
*arXiv preprint arXiv:2402.14735*, 2024.
- Nichol and Schulman (2018)
Alex Nichol and John Schulman.
Reptile: a scalable metalearning algorithm.
*arXiv preprint arXiv:1803.02999*, 2(3):4, 2018.
- Nichol et al. (2018)
Alex Nichol, Joshua Achiam, and John Schulman.
On first-order meta-learning algorithms.
*arXiv preprint arXiv:1803.02999*, 2018.
- Oquab et al. (2023)
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al.
Dinov2: Learning robust visual features without supervision.
*arXiv preprint arXiv:2304.07193*, 2023.
- Raghu et al. (2019)
Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals.
Rapid learning or feature reuse? towards understanding the effectiveness of maml.
*arXiv preprint arXiv:1909.09157*, 2019.
- Rajeswaran et al. (2019)
Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine.
Meta-learning with implicit gradients.
*Advances in neural information processing systems*, 32, 2019.
- Reuter et al. (2025)
Arik Reuter, Tim GJ Rudner, Vincent Fortuin, and David Rügamer.
Can transformers learn full bayesian inference in context?
*arXiv preprint arXiv:2501.16825*, 2025.
- Rezende et al. (2014)
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra.
Stochastic backpropagation and approximate inference in deep generative models.
In *International conference on machine learning*, pages 1278–1286. PMLR, 2014.
- Rusu et al. (2018)
Andrei A Rusu, Dushyant Rao, Jakub Sygnowski, Oriol Vinyals, Razvan Pascanu, Simon Osindero, and Raia Hadsell.
Meta-learning with latent embedding optimization.
*arXiv preprint arXiv:1807.05960*, 2018.
- Sanh et al. (2021)
Victor Sanh, Albert Webson, Colin Raffel, Stephen H Bach, Lintang Sutawika, Zaid Alyafeai, Antoine Chaffin, Arnaud Stiegler, Teven Le Scao, Arun Raja, et al.
Multitask prompted training enables zero-shot task generalization.
*arXiv preprint arXiv:2110.08207*, 2021.
- Scetbon et al. (2024)
Meyer Scetbon, Joel Jennings, Agrin Hilmkil, Cheng Zhang, and Chao Ma.
A fixed-point approach for causal generative modeling.
*arXiv preprint arXiv:2404.06969*, 2024.
- Schug et al. (2024)
Simon Schug, Seijin Kobayashi, Yassir Akram, João Sacramento, and Razvan Pascanu.
Attention as a hypernetwork.
*arXiv preprint arXiv:2406.05816*, 2024.
- Shalev-Shwartz and Ben-David (2014)
Shai Shalev-Shwartz and Shai Ben-David.
*Understanding machine learning: From theory to algorithms*.
Cambridge university press, 2014.
- Tong et al. (2023)
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio.
Conditional flow matching: Simulation-free dynamic optimal transport.
*arXiv preprint arXiv:2302.00482*, 2(3), 2023.
- Vaswani et al. (2017)
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.
Attention is all you need.
*Advances in neural information processing systems*, 30, 2017.
- Von Oswald et al. (2019)
Johannes Von Oswald, Christian Henning, Benjamin F Grewe, and João Sacramento.
Continual learning with hypernetworks.
*arXiv preprint arXiv:1906.00695*, 2019.
- Von Oswald et al. (2023)
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov.
Transformers learn in-context by gradient descent.
In *International Conference on Machine Learning*, pages 35151–35174. PMLR, 2023.
- Watts and Strogatz (1998)
Duncan J Watts and Steven H Strogatz.
Collective dynamics of ‘small-world’networks.
*nature*, 393(6684):440–442, 1998.
- Wei et al. (2022a)
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al.
Emergent abilities of large language models.
*arXiv preprint arXiv:2206.07682*, 2022a.
- Wei et al. (2022b)
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al.
Chain-of-thought prompting elicits reasoning in large language models.
*Advances in neural information processing systems*, 35:24824–24837, 2022b.
- Wichrowska et al. (2017)
Olga Wichrowska, Niru Maheswaranathan, Matthew W Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Nando Freitas, and Jascha Sohl-Dickstein.
Learned optimizers that scale and generalize.
In *International conference on machine learning*, pages 3751–3760. PMLR, 2017.
- Zhang et al. (2016)
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.
Understanding deep learning requires rethinking generalization.
*arXiv preprint arXiv:1611.03530*, 2016.

## 7 Related Work

In this section, we formalize and contrast several existing related frameworks encapsulated in our setting. We highlight the technical differences among these approaches through formal definitions and equations, while retaining key references.

### 7.1 Meta-Learning

Meta-learning generally refers to the problem of learning a model (which itself could be a learner, optimizer, or algorithm) that can quickly and efficiently generalize to novel tasks. Formally, let there be a distribution over tasks $p(\mathcal{T})$, with each task $\mathcal{T}$ associated with a dataset $\mathcal{D}_{\mathcal{T}}=\{(\mathbf{x}_{j},\mathbf{y}_{j})\}_{j=1}^{N}$. The objective is to learn meta-parameters ${\bm{\theta}}$ that enable fast adaptation to new tasks $\mathcal{T}^{\prime}$ with dataset $\mathcal{D}_{{\mathcal{T}}^{\prime}}$:

$$ $\displaystyle\min_{\theta}\mathbb{E}_{\mathcal{T}}\mathcal{L}\left({\mathcal{D}}_{\mathcal{T}}^{\text{valid}},U_{k}({\bm{\theta}},\mathcal{D}_{\mathcal{T}}^{\text{train}})\right),$ (10) $$

where:

- •
$U_{k}({\bm{\theta}},\mathcal{D}_{\mathcal{T}}^{\text{train}})$ is an adaptation operator applying $k$ gradient steps or another procedure starting from ${\bm{\theta}}$,
- •
$\mathcal{L}\left({\mathcal{D}}_{\mathcal{T}}^{\text{valid}},\cdot\right)$ is the loss on held-out data $\mathcal{D}_{\mathcal{T}}^{\text{valid}}$.

The popular Model-Agnostic Meta-Learning (MAML) framework (Finn et al., 2017) fits into this paradigm with the update:

$$ $\displaystyle{\bm{\theta}}^{\prime}={\bm{\theta}}-\alpha\nabla_{\bm{\theta}}\mathcal{L}\left({\mathcal{D}}_{\mathcal{T}}^{\text{train}},\cdot\right)$ (11) $$

where the meta-objective optimizes ${\bm{\theta}}$ such that ${\bm{\theta}}^{\prime}$ performs well on $\mathcal{D}_{\mathcal{T}}^{\text{valid}}$. Extensions such as Reptile (Nichol and Schulman, 2018) and Meta-SGD (Antoniou et al., 2018) also follow this principle. These methods ultimately fit within the *parametric modeling* framework, as the model parameters ${\bm{\theta}}$ are explicitly adapted to new tasks (Raghu et al., 2019). Here, the adaptation operator $U_{k}$ can be seen as the analogue of $g_{\bm{\varphi}}$ from our setting.

### 7.2 Amortization

Amortized inference traditionally refers to learning an inference network $q_{\phi}(\mathbf{z}|\mathbf{x})$ to approximate the posterior $p(\mathbf{z}|\mathbf{x})$, as in variational autoencoders (Kingma et al., 2013; Rezende et al., 2014).
Our focus differs by considering amortization at the dataset or task level, rather than per-observation. We formalize this as learning an amortized posterior:

$$ $\displaystyle q_{\phi}(\mathbf{z}\mid\mathcal{D})\approx p(\mathbf{z}\mid\mathcal{D}),$ (12) $$

where $\mathcal{D}$ is a dataset corresponding to a task. This approach implicitly or explicitly learns an optimizer or inference model that can solve novel tasks *zero-shot*. This viewpoint aligns with probabilistic meta-learning frameworks (Garnelo et al., 2018b; Amos et al., 2023), which emphasize amortizing inference across a distribution of tasks. Such methods can be framed as parametric approaches when the target ${\mathbf{z}}$ of interest are all the parameters of the likelihood (Mittal et al., 2025b) or explicit approaches when the likelihood is learned as well and ${\mathbf{z}}$ just represents some latents (Garnelo et al., 2018a, b).

In more general terms, all methods of meta-learning rely on some or the other notion of amortization and depending on the object that is trained to be amortized, we recover parametric, explicit or implicit models.

### 7.3 Learned Optimizers

Learned optimizers learn parameter update functions $h_{\bm{\varphi}}$ conditioned on gradients and optimization states:

$$ $\displaystyle{\bm{\theta}}_{t+1}={\bm{\theta}}_{t}+h_{\bm{\varphi}}(\nabla_{\bm{\theta}}\mathcal{L}({\mathcal{B}}_{\mathcal{T}},{\bm{\theta}}),{\mathbf{s}}_{t}),$ (13) $$

where ${\mathbf{s}}_{t}$ denotes the optimizer’s internal state, e.g., momentum or recurrent memory and ${\mathcal{B}}_{\mathcal{T}}$ represents a mini-batch of task-specific training data. These methods have been studied extensively (Metz et al., 2022b; Knyazev et al., 2024; Metz et al., 2019, 2022a; Li and Malik, 2017; Wichrowska et al., 2017).

While learned optimizers enable scalable amortization of inference, their hypothesis class is limited by dependence on *only first-order gradient information* and parameter updates. They cannot easily represent complex second-order dynamics such as Newton methods, which rely on Hessian information. Furthermore, they require explicit task parameters ${\bm{\theta}}$, restricting them to parametric models and preventing direct applicability to implicit modeling approaches. Here, one can see the recursive application of $h_{\bm{\varphi}}$ as the $g_{\bm{\varphi}}$.

### 7.4 Hypernetworks

Hypernetworks (Ha et al., 2016) are architectures that take a dataset $\mathcal{D}$ as input and generate the weights ${\bm{\theta}}$ of a target neural network to be used for predictions on that dataset:

$$ $\displaystyle{\bm{\theta}}=H_{\bm{\varphi}}(\mathcal{D})$ (14) $$

with the target model output

$$ $\displaystyle{\mathbf{y}}=f(\mathbf{x},{\bm{\theta}})=f(\mathbf{x},H_{\bm{\varphi}}({\mathcal{D}}))$ (15) $$

Such models have been studied extensively (Krueger et al., 2017; Chauhan et al., 2024; Schug et al., 2024; Von Oswald et al., 2019), including recent work that shows standard transformers can be interpreted as hypernetworks (Schug et al., 2024).
Hypernetworks are generally *parametric* since ${\bm{\theta}}$ is explicitly generated. However, if the hypernetwork conditions directly on the query $\mathbf{x}$, it can realize implicit models within our taxonomy. The general case of hypernetworks can, however, be seen as parametric modeling with $H_{\bm{\varphi}}$ being the $g_{\bm{\varphi}}$ component.

### 7.5 In-Context Learning and Prior-Fitted Networks

In-context learning (ICL) solves similar problems by training models (notably transformers) to make predictions conditioned on example input-output pairs provided as context, either by modeling explicit parameters akin to hypernetworks (Mittal et al., 2025b, a) or without explicit parameter updates (Mittal et al., 2024; Elmoznino et al., 2024; Müller et al., 2021; Garg et al., 2022; Von Oswald et al., 2023; Hollmann et al., 2022; Li et al., 2023; Bai et al., 2023).

Formally, given a context $\mathcal{C}=\{(\mathbf{x}_{i},{\mathbf{y}}_{i})\}_{i=1}^{K}$, the model learns to produce predictions $f_{\gamma}(\mathbf{x},\mathcal{D})$
with no parameter updates performed during inference. ICL and PFNs share the same modeling setup when the input is a sequence of examples, though ICL can be more general by conditioning on other forms of task information such as language prompts.

Both are subsumed within our unified framework, categorized as *parametric*, *explicit*, or *implicit* models depending on how the in-context learner is parameterized.

## 8 Overall pipeline

We additionally provide a rough sketch that walks through our pipeline, which consists of two distinct phases.

### 8.1 Meta-Training

For all methods — parametric, explicit, or implicit — we train an amortized model on a diverse set of tasks. All frameworks train a Transformer as the amortized network, with an optional MLP as the predictor $f_{\gamma}$ for the explicit setup. Each task is either derived from a dataset (e.g. , MNIST or ImageNet) or synthetically generated (e.g., linear regression tasks).
Details of task construction are in [Sections˜5](https://arxiv.org/html/2510.11471v1#S5) and [9](https://arxiv.org/html/2510.11471v1#S9), but we provide a concrete example here:

For MNIST, each task involves applying:

- •
A random projection from the original $28\times 28$ image space to a 100-dimensional space.
- •
A random permutation of class labels (e.g., all images labeled ’5’ may be relabeled as ’2’ and so on).

Training Procedure.
All methods are trained using maximum likelihood: the objective is to predict a query’s label given a batch of context examples and the model’s current state. However, training how the state updates is only done greedily – given a state and mini-batch, the model is trained to predict the next best state without taking into account how it impacts future states.

Distinction Between Method Variants.
Owing to the unified framework, we can now succinctly describe the key difference across different amortization methods simply as how this state is defined:

- •
Parametric: The state is the weights of a linear predictor, fixing the mapping from state and query to prediction (denoted $f$).
- •
Explicit: The state is a learned high-dimensional latent vector. It is updated using context examples and gradients, and is then used by a learned prediction function (denoted $f_{\gamma}$).
- •
Implicit: The state is defined per-query, representing the model’s belief about the correct label for that specific query. Note that in this case the state is not a function of the dataset, but a function of (dataset, query).

### 8.2 Evaluation

In-Distribution
After meta-training, we evaluate the model on in-distribution tasks, i.e., tasks drawn from the same distribution as the training tasks. For example, with MNIST:

- •
A new random projection matrix is sampled (from the same distribution used during training).
- •
A subset of training examples are projected to form the context set.
- •
The model then makes predictions on the corresponding projection of test examples.

This setup evaluates whether the amortized model generalizes to unseen tasks that are structurally similar to those seen during training (i.e. new random projections and new test images).

Out-of-Distribution (OoD). To assess cross-dataset generalization, we also evaluate the model on out-of-distribution tasks. For example:

- •
We evaluate the MNIST model on FashionMNIST — context and queries are drawn from the train and test set of FashionMNIST respectively, following the same procedure for random projection as in-distribution.

This tests whether the learned algorithm transfers to tasks with different input distributions. Note that for ImageNet experiments we do not use random projections but instead rely on Dino-v2 embeddings.

## 9 Dataset details

In this section, we provide details of all the datasets and tasks considered for pre-training the different variants of amortized estimators as well as the tasks used for evaluation.

### 9.1 Linear Regression

For the problem of linear regression, we define each task with a corresponding weight vector ${\mathbf{w}}_{\mathcal{T}}$ such that observations in the context and query set follow the mapping $({\mathbf{x}},{\mathbf{y}}={\mathbf{w}}_{\mathcal{T}}^{T}{\mathbf{x}}+\epsilon)$ where ${\mathbf{x}}$ is sampled from a standard normal distribution and $\epsilon$ denotes random noise sampled from a normal distribution with standard deviation of $0.25$. For this problem, we define the family of tasks for pre-training and evaluation corresponding to ${\mathbf{w}}_{\mathcal{T}}$ sampled randomly from the unit normal.

For training we pass a set of $\{({\mathbf{x}}_{i},{\mathbf{y}}_{i})\}_{i}$ as the context and a set of ${\mathbf{x}}$ as the query, where the underlying mapping between them is shared to be the same ${\mathbf{w}}_{\mathcal{T}}$. The task of the in-context learner is to predict either the coefficients ${\bm{\theta}}_{\mathcal{T}}$ that are fed to a prediction network or known mapping (explicit or parametric), or it predicts the prediction corresponding to ${\mathbf{x}}$ directly. Thus, the in-context learner has to learn to model and minimize the underlying validation loss corresponding to the problem.

### 9.2 MNIST / FashionMNIST Classification

Similar to linear regression, we randomly sample a projection matrix $W_{\mathcal{T}}\in\mathbb{R}^{784\times 100}$ from unit normal and a class re-labeling matrix $\pi_{\mathcal{T}}$ such that $\pi_{{\mathcal{T}},i,j}\in\{0,1\}$ and $\sum\pi_{{\mathcal{T}},i,:}=1$ such that each task is defined with the following transformations on observations and labels:

$$ $\displaystyle{\mathbf{x}}\to W_{\mathcal{T}}{\mathbf{x}}\quad{\mathbf{y}}\to\pi_{\mathcal{T}}{\mathbf{y}}$ (16) $$

where ${\mathbf{x}}\in\mathbb{R}^{784}$ and ${\mathbf{y}}$ denotes the one-hot vector corresponding to the class. This operation defined above can be seen as randomly projecting pixel values to a lower dimensional space and re-labeling the classification problem such that the same classes are mapped to the same new class $-$ for e.g. the remapping could map digits $1,4,7$ to class $1$, and so on. That is, this re-labeling can club multiple classes into a single class or rename class labels, but never sends two objects from the same class to different labels.

We follow the same setup for FashionMNIST as well, and then check for OoD evaluation on the other dataset by feeding a set of observations as context and evaluating on queries; note that task specific variables $W_{\mathcal{T}}$, $\pi_{\mathcal{T}}$ are shared between context and query.

### 9.3 ImageNet Classification

We follow the same setup as MNIST / FashionMNIST classification with just a few differences $-$ dimensionality reduction is done using a pre-trained Dino-v2 model and thus this projection operator is shared across tasks and not task-dependent anymore, and each task is defined by sampling a subset of images corresponding to $100$ random classes of ImageNet from which grouping / re-labeling is done exactly as above. That is, for any task the problem is to solve the maximum $100$-way classification problem as opposed to the single $1000$-way problem of ImageNet. Correspondingly, evaluation is done on contexts from ImageNet training set and queries from evaluation set corresponding to multiple $100$-way problems as well as CIFAR-10 and CIFAR-100’s train and test set forming the context and query.

### 9.4 Topological Order Prediction

#### Problem Statement.

We begin by defining *Structural Causal Models (SCMs)*, which formalize the causal generative process over a set of random variables. An SCM defines a distribution over of $d$ endogenous variables $\bm{V}=\{X_{1},\ldots,X_{d}\}\penalty 10000\ \sim\mathbb{P}_{X}$, each determined by a deterministic function of its parents ($F_{i}$) and a corresponding exogenous noise term ($N_{i}$). Specifically, each variable $X_{i}$ is generated as:

$$ $\displaystyle X_{i}=F_{i}(\text{Pa}(X_{i}),N_{i}),\quad\text{with}\quad\text{Pa}(X_{i})\subseteq\bm{V}\setminus\{X_{i}\}$ (17) $$

where the functions $F_{i}$ define the causal mechanisms, $\text{Pa}(X_{i})$ denotes the set of direct causes (parents) of $X_{i}$, and $N_{i}$ is a latent noise variable drawn independently from a distribution $\mathbb{P}_{N}$. Collectively, the SCM is represented as $\mathcal{S}(\mathbb{P}_{N},F,{\mathcal{G}})$, where ${\mathcal{G}}\in\{0,1\}^{d\times d}$ is the adjacency matrix of the causal graph, i.e., ${\mathcal{G}}_{ij}=1$ if $X_{j}\in\text{Pa}(X_{i})$.

Following standard assumptions in the literature, we consider markovian SCMS, where we restrict the causal graph ${\mathcal{G}}$ to be a directed acyclic graph (DAG) and the set of noise variables $\{N_{i},\cdots,N_{d}\}$ to be mutually independent.

Given the DAG assumption, we can obtain a unique topological order $\tau$ associated with the causal graph ${\mathcal{G}}$, and the prediction task is defined as follows.

> Given a dataset of causal variables $D_{X}\in{\mathbb{R}}^{n\times d}$ samples from an unknown SCM $\mathcal{S}(\mathbb{P}_{N},F,{\mathcal{G}})$, the goal is predict the associated topological order $\tau$ from $D_{X}$.

#### Method.

For amortized inference of topological order, we follow the approach from Scetbon et al. (2024). They leverage transformer-based architecture that attend over both the sample dimension ($n$) and the node dimension ($d$) to attend over the context, followed by a linear prediction layer to classify the leaf nodes (no outgoing edge) of the causal graph ${\mathcal{G}}$. Once we obtain the current leaf nodes, we can remove them from the dataset $D_{X}$ and repeat the procedure again to predict leaf nodes. This recursive procedure will terminate in at most $d$ iteration and output the inferred topological order $\hat{\tau}$.

#### Synthetic Data Simulator.

Following (Scetbon et al., 2024), we use the AVICI synthetic data simulator Lorch et al. (2022). This synthetic generator supports diverse structural and functional variations, making it particularly well-suited for training and evaluating models under distribution shifts. We describe below the options avaiable for each compoenent in the SCM.

- •
*Graph Structures.* We can sample causal graphs from a variety of schemes; Erdős–Rényi graphs (Erdos and Renyi, 1959), Scale-free networks (Barabási and Albert, 1999), Watts–Strogatz small-world graphs (Watts and Strogatz, 1998), and Stochastic block models (Holland et al., 1983).
- •
*Noise Distributions.* Exogenous noise variables $\{N_{i}\}$ are sampled from either Gaussian or Laplace distributions, with randomly selected variances.
- •
*Functional Relationships.* Causal relationships are instantiated using either Linear (LIN) models with randomly sampled weights and biases, or with Random Fourier Features (RFF) to generate more complex, non-linear mappings.

To simulate distribution shifts, two families of SCM distributions are defined:

- •
*In-Distribution* ($\mathbb{P}_{\text{in}}$): Graphs are sampled from Erdős–Rényi and scale-free models, noise from Gaussian distributions, and functions are either LIN or RFF using standard parameter ranges.
- •
*Out-of-Distribution* ($\mathbb{P}_{\text{out}}$): Graphs are drawn from Watts–Strogatz or stochastic block models, noise from Laplace distributions, and functions (LIN/RFF) are sampled from disjoint parameter ranges to introduce a shift.

*Parameter Ranges:*

- •
Linear Mechanisms:
–
$\mathbb{P}_{\text{in}}$: weights $\sim U_{\pm}(1,3)$, bias $\sim U(-3,3)$
–
$\mathbb{P}_{\text{out}}$: weights $\sim U_{\pm}(0.5,2)\cup U_{\pm}(2,4)$, same bias range
- •
RFF Mechanisms:
–
$\mathbb{P}_{\text{in}}$: length scale $\sim U(7,10)$, output scale $\sim U(5,8)\cup U(8,12)$, bias $\sim U_{\pm}(-3,3)$
–
$\mathbb{P}_{\text{out}}$: length scale $\sim U(10,20)$, output scale $\sim U(8,12)\cup U(18,22)$, bias $\sim U_{\pm}(-3,3)$

*Train Datasets.* We train the model by randomly sampling SCMs from the $\mathbb{P}_{\text{in}}$ distribution. Each epoch contains datasets with $n=1000$ samples and $d=20$ nodes.

*Test Datasets.* We evaluate performance across four distinct settings that differ in the distribution they induce over SCMs. From each SCM, we sample a two test datasets, one with $n=1000$ samples and $d=20$ nodes; and the with $n=1000$ samples and $d=50$ nodes.

- •
LIN: Linear mechanisms under $\mathbb{P}_{\text{in}}$; total of 9 randomly sampled SCMs with 3 different graph types.
- •
RIN: Nonlinear (RFF) mechanisms under $\mathbb{P}_{\text{in}}$; total of 9 randomly sampled SCMs with 3 different graph types.
- •
LOUT: Linear mechanisms under $\mathbb{P}_{\text{out}}$; total of 6 randomly sampled SCMs with 2 different graph types.
- •
ROUT: Nonlinear (RFF) mechanisms under $\mathbb{P}_{\text{out}}$; total of 6 randomly sampled SCMs with 2 different graph types.

### 9.5 Alphabets

We next turn our attention to generation tasks where the goal is to generate novel samples by inferring the underlying density defined by the context examples and consequently learning to sample from it. We first borrow the task of sampling point clouds resembling alphabets (Atanackovic et al., 2024) where different tasks ${\mathcal{T}}$ correspond to different alphabets with different scaling and rotation operations applied. Out of all the capitalized alphabets in the English language, we reserve {D,H,N,T,X,Y} solely for evaluation, i.e. they are not provided for training on during pre-training. We then consider evaluation on both in-distribution and out-of-distribution alphabets.

### 9.6 Gaussian Mixture Model

Similar to the above, we consider the problem of sampling similar to a mixture of Gaussians provided as in-context examples. Here, each task is defined by an underlying number of clusters $n_{{\mathcal{T}}}$ as well as corresponding means $\mu_{{\mathcal{T}},i}$ which are sampled from a normal distribution with 5 standard deviation. Samples from this GMM are obtained by fixing the standard deviation corresponding to each cluster as 0.3 and then drawing samples from this mixture distribution. We train the in-context estimator on a constant stream of new tasks synthetically generated and then evaluate them on new tasks sampled randomly. Our tasks involve both a simple version of GMM which is in $2$-dimensional observed space as well as a more complex GMM which is $5$-dimensional. The number of clusters are uniformly sampled from a maximum of 100 clusters.

## 10 Metrics

For metrics, we consider the $l_{2}$ loss for the regression problem while for the classification problems, we consider error as a metric, which is $100-$Accuracy, which we use for the topological prediction task as well. For the generative modeling tasks, we consider the $2$-Wasserstein and $1$-Wasserstein distance which are based on optimal transport and can be described as

$$ $\displaystyle{\mathcal{W}}_{p}=(\inf_{\pi}\frac{1}{N}\sum_{i}\|{\mathbf{x}}_{i}-\hat{{\mathbf{x}}}_{\pi_{i}}\|_{p}^{p})^{1/p}$ (18) $$

where $\pi$ describes a permutation matrix, of which we use $p=1$ or $p=2$.

## 11 Experimental Details

### 11.1 Model Parameterization

We parameterize the amortized model, $g_{\bm{\varphi}}$ in the case of parametric and explicit model, and $f_{\gamma}$ in the case of implicit model, as a transformer with $512$ hidden dimensions, $2048$ hidden dimensions in the feed-forward neural network and $8$ heads and $8$ layers. We use gelu activation function and perform normalization first in PyTorch’s version of transformers. In addition, we consider a specific learnable linear encoder for gradients as well as observations, where observations are $({\mathbf{x}},{\mathbf{y}})$ pairs and gradients correspond to $\nabla_{\bm{\theta}}{\mathcal{L}}({\mathbf{y}},f_{\gamma}({\mathbf{x}},{\bm{\theta}}))|_{{\bm{\theta}}_{\mathcal{T}}}$ for parametric and explicit models while implicit models cannot use gradients.

In addition, for implicit models we append $\hat{{\mathbf{y}}}$ to query observations and re-use the observation encoder to embed queries in the same space. We additionally use a learnable query embedding that is added to the query tokens for implicit model to be able to differentiate context from query. For the ablation conducted on Pre-MLP state that is carried over, we leverage a separate learnable query embedding matrix that does not share weights with the observation encoder. Finally, to model predictions or vector fields in the case of generative modeling, we leverage a similar linear decoder that maps from $512$ dimensions to dimensions of the prediction, parameters, latents or vector field depending on the task and modeling setup.

It is important to note that our model should be able to make predictions or model parameters conditioned on arbitrary number of context examples, and hence should be trained in a manner that allows for varying context length $-$ this is similar to decoder only transformers that learn to make predictions for every context length in parallel. Given the difference in our setting from autoregressive language modeling, we provide two ways of implementing this amortization through the use of either a causally masked transformer or a non causal transformer, the details of which are provided in the next subsections.

Finally, for modeling iterative amortized inference, we use a simple setup where the output from one batch $-$ ${\bm{\theta}}$ or $\hat{{\mathbf{y}}}$ depending on the amortization framework used $-$ is fed back into the transformer with another batch after being detached from the computation graph. Here, ${\bm{\theta}}$ is just fed as another token with or without its gradient information as an additional token, while for implicit models $\hat{{\mathbf{y}}}$ is appended to its corresponding query, hence the state represents the predictions corresponding to the query that get updated after every iteration $-$ note that the state is tied to a query here. Given a number of steps, e.g. $k$, we iterate over this process $k$ times and compute gradients each time, accumulating them before taking a gradient step. Given that we detach the state before feeding it again, it prevents backpropagation through time and thus defines a greedy procedure.

### 11.2 Training Objective

The training procedure for all the models can be seen as optimization over the following loss for parametric and explicit models, where the batch sizes can contain variable number of observations $N$ which is randomly sampled from $1-100$.

$$ $\displaystyle\arg\min_{\gamma,{\bm{\varphi}}}\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}_{{\mathcal{T}}}\mathbb{E}_{{\mathcal{B}}_{\mathcal{T}}^{\text{valid}}\subseteq{\mathcal{D}}_{\mathcal{T}}^{\text{valid}}}\mathbb{E}_{{\mathcal{B}}_{\mathcal{T}}^{(t),\text{train}}\subseteq{\mathcal{D}}_{\mathcal{T}}^{\text{train}}}$ $\displaystyle\frac{1}{N}\sum_{({\mathbf{x}},{\mathbf{y}})\in{\mathcal{B}}_{\mathcal{T}}^{\text{valid}}}{\mathcal{L}}({\mathbf{y}},f_{\gamma}({\mathbf{x}},g_{\bm{\varphi}}({\mathcal{B}}_{\mathcal{T}}^{(t),\text{train}},{\bm{\theta}}_{\mathcal{T}}^{(t)})))$ (19) $\displaystyle{\bm{\theta}}_{\mathcal{T}}^{(t)}$ $\displaystyle=g_{\bm{\varphi}}({\mathcal{B}}_{\mathcal{T}}^{(t-1),\text{train}},\text{sg}({\bm{\theta}}^{(t-1)}_{\mathcal{T}}))$ (20) $$

In contrast, implicit models consider the following learning paradigm

$$ $\displaystyle\arg\min_{\gamma,{\bm{\varphi}}}\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}_{{\mathcal{T}}}\mathbb{E}_{{\mathcal{B}}_{\mathcal{T}}^{\text{valid}}\subseteq{\mathcal{D}}_{\mathcal{T}}^{\text{valid}}}\mathbb{E}_{{\mathcal{B}}_{\mathcal{T}}^{(t),\text{train}}\subseteq{\mathcal{D}}_{\mathcal{T}}^{\text{train}}}$ $\displaystyle\frac{1}{N}\sum_{({\mathbf{x}},{\mathbf{y}})\in{\mathcal{B}}_{\mathcal{T}}^{\text{valid}}}{\mathcal{L}}({\mathbf{y}},f_{\gamma}({\mathbf{x}},{\mathcal{B}}_{\mathcal{T}}^{(t),\text{train}},\hat{{\mathbf{y}}}^{(t)})))$ (21) $\displaystyle\hat{{\mathbf{y}}}^{(t)}$ $\displaystyle=f_{\gamma}({\mathcal{B}}_{\mathcal{T}}^{(t-1),\text{train}},[{\mathbf{x}},\text{sg}(\hat{{\mathbf{y}}}^{(t-1)})])$ (22) $$

Figure: Figure 5: Masking procedure for the causally masked parametric and explicit model, where the context evolves according to a causal mask $-$ the matrix of black and white squares describes the masking procedure where white blocks denote masks in the matrix $-$ where each token in parallel consequently predicts parameters of interest conditioned on past batch of data and previous state. This mimics having variable sized dataset and processing for variable dataset sizes in parallel. The output corresponding to the last token is the ${\bm{\theta}}_{\mathcal{T}}^{(t)}$ that gets fed back recurrently.
Refer to caption: x11.png

Figure: Figure 6: Masking procedure for the causally masked implicit model, where the context evolves according to a causal mask and the queries can attend somewhere in between, which conditions them on that particular position and the positions behind it. Since we use multiple queries in parallel which can look at different number of past contexts through this manipulation of the masking, we obtain parallel processing of multiple dataset sizes.
Refer to caption: x12.png

### 11.3 Causal Masked Model

Given that the implicit model differs considerably from the parametric and explicit model in terms of conditional independence assumptions defined, we need to handle the two cases somewhat differently to enable efficient parallelized implementations. We discuss the details below.

Parametric and Explicit. For the parametric and explicit model, we consider a simple causal masking over observations such that the model predicts ${\bm{\theta}}_{\mathcal{T}}$ corresponding to each token conditioned on the previous token, and the loss is aggregated by using all the intermediate ${\bm{\theta}}_{\mathcal{T}}$’s obtained to make predictions on a validation set and averaging them. This procedure is highlighted in [Fig.˜5](https://arxiv.org/html/2510.11471v1#S11.F5).

Implicit. In contrast to the parametric approach, the implicit framework cannot deal with this in such a simple manner since the query is ${\mathbf{x}}$ itself which cannot remain in the same position as well as attend over variable sized context. To resolve this issue, we process in parallel multiple points from the validation set and manipulate the masking matrix to randomly let them attend to some token $j$ and its predecessors. We use a causal transformer so this ensures that if the query looks at index $j$ and before, then that query models prediction using a dataset of size $j$ instead of context length. We use multiple queries in parallel, randomly choosing the corresponding $j$ so that we can, in-parallel, model predictions for variable sized datasets in a single forward pass (though they have to correspond to different queries). A pictorial representation of this is showcased in [Fig.˜6](https://arxiv.org/html/2510.11471v1#S11.F6).

Figure: Figure 7: We show the case of non causal masked transformer for parametric and explicit setup where at each iteration, a differently sized context is provided as input to predict the next state. The size of the context is randomly sampled. At evaluation we use the full context.
Refer to caption: x13.png

### 11.4 Non-Causal Model

For the non causal masked version, at each training iteration we just randomly sample a length $n$ and then consider a batch of only $n$ observations for a forward pass. The benefits of this approach is that it can allow for conditioning on the $n$ observations in-context without having to rely on a causal decomposition of attention but the downside is that the number of observations have to be kept fixed for a training iteration leading to unbiased gradients but they may have high variance since each training iteration only relies on a single sample ($n$) of the number of observations which then differ across training iterations. Thus, it is not possible to leverage parallel computation for a variety of number of observations within a single forward pass.

A benefit of this approach is that for parametric and explicit models, it can allow more than one token to denote the latents which is computationally infeasible for the masked case. For explicit models in particular, we consider $8$ tokens for latents which are of $128$ dimensions each; note that having this for masked case is possible but it will require complex masking procedures as well as efficient sparse attention kernels otherwise the context length would blow up relatively quickly.

We refer to [Figs.˜7](https://arxiv.org/html/2510.11471v1#S11.F7) and [8](https://arxiv.org/html/2510.11471v1#S11.F8) for details regarding the non causal transformer model for the parametric, explicit and implicit cases.

Figure: Figure 8: We highlight the case of non-causal transformer model for the implicit model where at each iteration a subset of observations are provided as context; their number randomly sampled. At evaluation we use the full context.
Refer to caption: x14.png

## 12 Additional Ablations

### 12.1 Causal vs Non Causal Transformer

We perform a thorough comparison between using a causal and non-causal transformer in all three amortization regimes $-$ parametric, explicit, and implicit. We generally see mixed results which we hypothesize is because while the non-causal method is more expressive algorithm, it has a higher variance of gradients during training as one cannot process in parallel multiple dataset sizes and thus leverage larger effective batch size which makes them inefficient to train and the optimization process largely more unstable.

Parametric.

Figure: Figure 9: Parametric with Gradient Information. We look at the comparison between causal and non-causal transformer for a parametric model that leverages gradients but not data.
Refer to caption: x15.png

Figure: Figure 10: Parametric with Observations Information. We look at the comparison between causal and non-causal transformer for a parametric model that leverages data but not gradients.
Refer to caption: x16.png

Figure: Figure 11: Parametric with Observations and Gradient Information. We look at the comparison between causal and non-causal transformer for a parametric model that leverages both gradients and data.
Refer to caption: x17.png

For the parametric model, we consider the causal method that predicts weights of the linear predictor in parallel conditioned on ${\mathbf{x}}_{1:n}$ for all $n$ in parallel by leveraging a causal mask, as described in the section above whereas the non-causal framework leverages a more expressive transformer since the attention is not restricted to causal attention but at every training iteration, a single $n$ is sampled at random and conditioning is done on only a subset of $n$ observations (this $n$ is changed at every iteration). We can see that this is an unbiased estimate of the gradient but with only a single sample of $n$, and thus suffers from larger variance but allows for more expressive architecture. We highlight the performance difference of the two approaches for various number of iterations of amortization as well as modalities of information $-$ gradient or observations or both $-$ in [Figs.˜9](https://arxiv.org/html/2510.11471v1#S12.F9), [10](https://arxiv.org/html/2510.11471v1#S12.F10) and [11](https://arxiv.org/html/2510.11471v1#S12.F11).

Explicit.

Figure: Figure 12: Explicit with Gradient Information. We look at the comparison between causal and non-causal transformer for an explicit model that leverages gradients but not data.
Refer to caption: x18.png

Figure: Figure 13: Explicit with Observation Information. We look at the comparison between causal and non-causal transformer for an explicit model that leverages data but not gradients.
Refer to caption: x19.png

Figure: Figure 14: Explicit with Observation and Gradient Information. We look at the comparison between causal and non-causal transformer for an explicit model that leverages both gradients and data.
Refer to caption: x20.png

Similar to the parametric modeling setup, we look at the explicit amortization in the case of both a causal and a non-causal transformer with different number of iterations as well as different modalities of information $-$ gradient or observations or both $-$ in [Figs.˜12](https://arxiv.org/html/2510.11471v1#S12.F12), [13](https://arxiv.org/html/2510.11471v1#S12.F13) and [14](https://arxiv.org/html/2510.11471v1#S12.F14).

Implicit. We refer to [Fig.˜15](https://arxiv.org/html/2510.11471v1#S12.F15) for experimental details contrasting the causal and non-causal transformer model.

As highlighted at the start of the section, our results on the comparisons are mixed as one model is more expressive while the other leads to optimization challenges due to potential variance in gradients since it is a single sample monte carlo estimate.

Figure: Figure 15: Implicit with Observation Information. We look at the comparison between causal and non-causal transformer for an implicit model, which can only leverage data.
Refer to caption: x21.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 4
presentation: 2
questions: 1: It would be helpful to include a concise summary of the methods and an algorithmic description, especially for the iterative amortized inference component. Because the framework spans multiple cases, a step-by-step outline (inputs, updates, outputs) would make the overall system easier to follow.

2: A more thorough comparison against baselines from multiple categories (e.g., meta-learning, context learning, learned optimizers) would better demonstrate the framework’s effectiveness. Reporting results across diverse, representative baselines would clarify where the unified approach helps and where specialized methods still have an edge.
soundness: 3
strengths: 1: This paper investigates the similarities among different amortized methods, which are typically treated as distinct research areas. The conclusions drawn from this unified analysis are meaningful, as they help reveal the common underlying principles shared across these approaches.

2: The authors categorize various amortized methods based on whether the functions f and g are trainable or fixed. This categorization offers an interesting and novel perspective for understanding the relationships between different amortized learning paradigms.

3: The authors provide sufficient theoretical context and analysis to support the effectiveness of the proposed framework, thereby strengthening its conceptual soundness and empirical credibility.
summary: In this paper, the authors present a unified framework that encompasses various amortized learning paradigms, including meta-learning, context learning and learned optimization. Within this framework, a learned function $ f_{\gamma}$ operates in conjunction with another function $ g_{\phi}$. Specifically, $f_{\gamma}$ performs inference based on the knowledge that $g_{\phi}$ has acquired from a dataset $ D_{t}$. This formulation integrates the key characteristics of different amortized methods, allowing each individual approach to be viewed as a special case}within the broader framework. Furthermore, the authors categorize tasks into three classes, parametric, implicit, and explicit, which help distinguish among different amortized models according to whether $f_{\gamma}$ and $g_{\phi}$  are trainable or fixed. In addition, they propose a scalable sequential inference mechanism, where the output from each step serves as the input to the next, effectively modeling the process as a sequence-based system.
weaknesses: 1: The presentation of the paper is somewhat confusing. It lacks a clear clarification of the unified framework or an algorithmic summary. The purely descriptive explanations make it difficult to understand how the proposed unified system actually works in practice.

2: The explanation of the iterative amortized inference process is unclear. The distinction between the proposed method and existing approaches such as meta-learning is vague, which makes it challenging to identify the novel contributions of this work.

3: Since the framework aims to unify different amortized methods, the authors should provide a more thorough comparison with existing approaches across multiple baselines. This would help ensure that the unification does not lead to a significant performance drop relative to specialized methods.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: - The paper mentions improved scalability through iterative amortization. Could the authors clarify whether this scalability refers primarily to model size, dataset size, task complexity, or another dimension? A more detailed evaluation along the chosen dimension would strengthen the claim.
- How would recent works on meta-learning with in-context learning (e.g., [1,2]) fit into the proposed taxonomy?
- The experimental tables report error rates rather than average accuracy. Is there a specific reason for this choice? Reporting accuracy would make it easier to compare with previous studies.
- How is the number of steps for IAI selected? Do additional steps continue to yield improvements, or is there a point of diminishing returns?
- Given the iterative refinement process, IAI appears well suited for continual or lifelong learning setups. Have the authors considered this application, or do they foresee challenges in extending IAI to such scenarios?
soundness: 4
strengths: The unified perspective proposed in the paper is valuable, as it clearly articulates how different amortized learning methods relate to one another within a common framework. The taxonomy provides a useful conceptual structure that can guide future research in this area. Moreover, the proposed IAI approach effectively addresses the scalability limitations of previous methods. The empirical results are convincing, demonstrating consistent improvements over existing baselines in both in-distribution and out-of-distribution settings.
summary: This paper proposes Iterative Amortized Inference (IAI), a unified framework that generalizes amortized learning approaches (meta-learning, ICL, prompt tuning and learned optimizers) under a common mathematical framework. Furthermore, the authors provide a categorization of amortized models into parametric, implicit, and explicit regimes. Finally, they also propose a mechanism inspired by stochastic optimization, which processes mini-batches of task data incrementally.
weaknesses: - The paper could benefit from more general reference to the meta-learning literature. Including references to survey papers in the main text could help the reader to better understanding the meta-learning problem, which is only detailed in the appendix.
- While Section 5 discusses memory usage and computational cost, the analysis lacks details on the runtime required to train a k-step iterative model.
- It would be valuable to discuss theoretical guarantees (e.g., convergence bounds) for IAI.
- Minor issue: it would be helpful to illustrate $g_\psi$ in Figure 1.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 4
questions: 1. For parametric and explicit setup, the authors treat $\theta$ as a token to be input to transformer, along with other tokens which represent context. Considering the possible very large discrepancy in dimensions (e.g.,$ W_T \in R^{784\*100} $) while context units is about 8*8 in MNIST), how to feed $\theta$ as an input token to transformer? And how to unify different modalities of input tokens (parameter, data, gradient)?
The reviewer think the traditional way to modulate $\theta$ by some functions using the output of $h_\psi$ [5], rather than feeding $\theta$ into transformer, might be more reasonable and practical.

2. The principle idea of introducing iteration into inference, and the implementation, seems to be similar to auto-regressive training that has been commonly applied in ICL with transformer. What is the difference? 

3. The usage and meaning of  "amortization" seems different from literature [6]. I am curious where this term originates from, and why this paper use it in this way.

[5] Fast and flexible multi-task classification using conditional neural adaptive processes. NIPS2019

[6] Memory efficient meta-learning with large images. NIPS2021
soundness: 3
strengths: 1. The conceptual framework and taxonomy is nice and reasonable. The f(x, g(D_T)) formulation provides a clear and powerful lens for understanding the underlying mechanics of ICL, learned optimizers, and meta-learners, highlighting their shared principles.
The proposed taxonomy of parametric, implicit, and explicit amortization is intuitive and effectively organizes the space of amortized learning models.

2. Novel and well-motivated method. Iterative Amortized Inference is a novel solution to the critical problem of scaling in-context adaptation to large datasets. The idea of processing data in mini-batches at inference time is a natural and scalable extension of existing paradigms, merging the strengths of forward-pass and optimization-based approaches.
summary: This paper presents a unified framework for amortized learning, aiming to connect and clarify the relationship between methods like meta-learning, in-context learning (ICL), and learned optimizers. The authors propose that these diverse approaches can be understood through a common functional decomposition involving a task-adaptation function and a prediction function. Building on this framework, they introduce a taxonomy that categorizes models into three regimes: parametric, implicit, and explicit amortization, based on how they encode shared inductive biases and adapt to new tasks.
The paper also identifies a key weakness in existing methods: their limited ability to scale to large datasets during inference due to constraints like context length. To overcome this, the authors introduce Iterative Amortized Inference, a novel technique that refines a solution iteratively over mini-batches of task data.
weaknesses: 1. The unification of ICL and learned optimizers, and the taxonomy, are not novel. ICL has been viewed as a form of meta-learners in literature [1,2,4]. 
And the f(x, g(D_T)) framework also has exists in literature for a not short time [1,3].
Especially, [1] has unified ICL and meta-learners explicitly and theoretically, under the same f(x, g(D_T)) framework. As learned optimizer is a subset of meta-learners, this paper exactly falls into the existing paradigm. However,  [1] has not even been referred, as a very closely related work.

2. Limitation of proposed method:
the baseline comparison and application scenario discussion are insufficient. Note that most experiments follows meta-learning setting, where numerous meta-learning methods are feasible, but this paper has only compared with 1-step amortization with transformer (e.g., standard ICL). Even restricting in this scope, it seems that training set size per task would play a critical role in the effect of iterating over one-step, which should be ablation studied and discussed. Existing results are far from to be able to see when and what advantage IAI would bring comparing with feasible baselines in practice.

3. Some technical design of IAI is questioned. Please refer to Questions 1.

4. The final implementation of IAI has hardly distinct with existing auto-regressive training, which is a common practice in ICL with transformer. Please refer to Question 2.

[1] Why In-Context Learning Models are Good Few-Shot Learners?. ICLR2025

[2] General-purpose in-context learning by meta-learning transformers. arXiv:2212.04458

[3] Conditional neural processes. ICML 2018. and Neural Processes. ICML 2018 workshop

[4] Learning to Learn with Contrastive Meta-Objective. arXiv:2410.05975

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 2
questions: The errors in table 2 are quite large. Is there any intuition why?
soundness: 3
strengths: This paper provides a framework that unifies three disparate ideas, meta-learning, learned optimizers, and in-context learning. The connection of in-context learning to the others is especially insightful. I believe that this kind of work that combines different topics and provides a bird's-eye view is valuable for the community, as it has the potential to bring to light new avenues of research.

The framework itself is mathematically rigorous and clearly defined, and it encompasses many previous algorithms. Experiments are done on a variety of tasks, including classification, regression, and generation. Results provide practical lessons, e.g. parametric models outperform explicit models.
summary: This paper unifies meta-learning, learned optimizers, in-context learning, etc. into a common framework of amortized learning, i.e. learning how to reuse shared information across tasks to improve performance. In this framework, methods vary along two axes, how meta-testing data is converted into task-specific parameters, and how those parameters are combined with queries to obtain predictions. Those two processes are either fixed or learnable, leading to three categories of methods. Based on the categorization, the paper then proposes two general methods that iteratively refines task-specific parameters using mini-batches of meta-testing data (which subsumes some existing algorithms). Experiments are done on regression, classification, and generation datasets showing the power of the framework.
weaknesses: The experiments are incomplete in some respects. First, only one network architecture is considered - it would strengthen the work to see how the method scales. Second, the parametric and explicit models only consider regression and classification tasks, but they can also be used for generation.

Moreover, the paper can be unclear at times. The main idea is abstract, so more explanation would be helpful, for example by:
- Adding an explanation to the caption of figure 1
- Explaining what likelihood is being referred to in line 206
- In line 322, explaining what is the causal structure being masked
- At the end of section 4, comparing and contrasting the proposed methods to existing algorithms


Decision: Reject