Title: Provable In-Context Learning of Nonlinear Regression with Transformers
Primary area: learning theory
Keywords: In-Context Learning, Transformers, Nonlinear Regression, Learning Dynamics

Abstract: The transformer architecture has revolutionized machine learning by processing input sequences into outputs. A defining feature is in-context learning (ICL)—the ability to perform unseen tasks from prompts without updating model parameters. Early theoretical work focused on linear tasks, and recent studies have begun exploring nonlinear functions. Yet a rigorous analysis of the training dynamics—how transformers learn such complex tasks—remains elusive. This paper presents the first formal analysis of ICL training dynamics for a broad class of nonlinear regression functions. We analyze the stage-wise dynamics of attention during training: attention scores between a query token and its target features rise rapidly at first, then gradually converge to one, while attention to irrelevant features decays more slowly and can oscillate. Our analysis explicitly characterizes how general non-degenerate $L$-Lipschitz task functions shape attention weights, identifying the Lipschitz constant $L$ as the key factor governing the convergence dynamics. Leveraging these insights, for two distinct regimes depending on whether $L$ is below or above a threshold, we derive different time bounds to guarantee near-zero prediction error.
Despite convergence time depending on the task, we prove query tokens ultimately focus on highly relevant prompt tokens, demonstrating transformers’ robust ICL capability for unseen functions.