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

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Mittal, Sarthak, Mahajan, Divyat, Lajoie, Guillaume, Pezeshki, Mohammad
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866917008719740928
author Mittal, Sarthak
Mahajan, Divyat
Lajoie, Guillaume
Pezeshki, Mohammad
author_facet Mittal, Sarthak
Mahajan, Divyat
Lajoie, Guillaume
Pezeshki, Mohammad
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2510_11471
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Iterative Amortized Inference: Unifying In-Context Learning and Learned Optimizers
Mittal, Sarthak
Mahajan, Divyat
Lajoie, Guillaume
Pezeshki, Mohammad
Machine Learning
Artificial Intelligence
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.
title Iterative Amortized Inference: Unifying In-Context Learning and Learned Optimizers
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2510.11471