Beyond Linear Attention: Softmax Transformers Implement In-Context Reinforcement Learning

• URL: http://arxiv.org/abs/2605.07333v1
• Date: Fri, 08 May 2026 06:37:51 GMT
• Title: Beyond Linear Attention: Softmax Transformers Implement In-Context Reinforcement Learning
• Abstract summary: In-context reinforcement learning studies agents that, after pretraining, adapt to new tasks by conditioning on additional context without parameter updates.<n>This paper provides the first theoretical understanding of ICRL without making the unrealistic linear attention simplification.<n>We show that, with certain parameters, the layerwise forward pass of a Transformer with such softmax attention is equivalent to iterative updates of a weighted softmax temporal difference (TD) learning algorithm.
• Score: 26.624140780367103
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: In-context reinforcement learning (ICRL) studies agents that, after pretraining, adapt to new tasks by conditioning on additional context without parameter updates. Existing theoretical analyses of ICRL largely rely on linear attention, which replaces the softmax function in the standard attention with an identity mapping. This paper provides the first theoretical understanding of ICRL without making the unrealistic linear attention simplification. In particular, we consider the standard softmax attention used in practice. We show that, with certain parameters, the layerwise forward pass of a Transformer with such softmax attention is equivalent to iterative updates of a weighted softmax temporal difference (TD) learning algorithm. Here, weighted softmax TD is a new RL algorithm that performs policy evaluation in kernel space and adopts both linear TD and tabular TD as special cases. We also prove that under a certain contraction condition, the policy evaluation error decays as the number of layers grows, with the identified parameters above. Finally, we prove that those parameters are a global minimizer of a pretraining loss, explaining their emergence in our numerical experiments.

Related papers

• The Entropic Bound for Transformers: Why Static Rank Fails and Attention-Native Rank Recovers [0.0]
  We study the Entropic Bound, a spectral notion of task-intrinsic capacity for Transformers.<n>We first prove that, in a linear attention surrogate, the intrinsic rank $r*$ of the token-mixing operator is a tight lower bound.<n>We then ask whether this transfers to real attention.<n>Motivated by this, we introduce an attention-native intrinsic rank.
  arXiv  Detail & Related papers  (2026-07-25T05:27:52Z)
• Fast Rates for Inverse Reinforcement Learning [12.651802060560636]
  We show that maximum likelihood estimation (MLE) and Min-Max-IRL are equivalent at the population level.<n>We derive novel results on the derivatives of the soft-optimal value function with respect to reward parameters.
  arXiv  Detail & Related papers  (2026-05-14T09:07:31Z)
• Dynamic Rank Reinforcement Learning for Adaptive Low-Rank Multi-Head Self Attention in Large Language Models [0.0]
  We propose Dynamic Rank Reinforcement Learning (DR-RL), a novel framework that adaptively optimize the low-rank factorization of Multi-Head Self-Attention (MHSA) in Large Language Models (LLMs)<n>DR-RL maintains downstream accuracy statistically equivalent to full-rank attention while significantly reducing Floating Point Operations (FLOPs)<n>This work bridges the gap between adaptive efficiency and theoretical rigor in MHSA, offering a principled, mathematically grounded alternative to rank reduction techniques in resource-constrained deep learning.
  arXiv  Detail & Related papers  (2025-12-17T21:09:19Z)
• Softmax as Linear Attention in the Large-Prompt Regime: a Measure-based Perspective [16.076157672455867]
  We develop a measure-based framework for studying single-layer softmax attention under both finite and infinite prompts.<n>For i.i.d. Gaussian inputs, we lean on the fact that the softmax operator converges in the infinite-prompt limit to a linear operator acting on the underlying input-token measure.
  arXiv  Detail & Related papers  (2025-12-12T18:54:52Z)
• Statistical Advantage of Softmax Attention: Insights from Single-Location Regression [0.0]
  We study the dominance of softmax over alternatives in large language models.<n>We show that softmax achieves the Bayes risk, whereas linear attention fundamentally falls short.<n>We discuss the connection with optimization by gradient-based algorithms.
  arXiv  Detail & Related papers  (2025-09-26T06:21:30Z)
• EKPC: Elastic Knowledge Preservation and Compensation for Class-Incremental Learning [53.88000987041739]
  Class-Incremental Learning (CIL) aims to enable AI models to continuously learn from sequentially arriving data of different classes over time.<n>We propose the Elastic Knowledge Preservation and Compensation (EKPC) method, integrating Importance-aware importance Regularization (IPR) and Trainable Semantic Drift Compensation (TSDC) for CIL.
  arXiv  Detail & Related papers  (2025-06-14T05:19:58Z)
• Sculpting Subspaces: Constrained Full Fine-Tuning in LLMs for Continual Learning [19.27175827358111]
  Continual learning in large language models (LLMs) is prone to catastrophic forgetting, where adapting to new tasks significantly degrades performance on previously learned ones.<n>We propose a novel continual full fine-tuning approach leveraging adaptive singular value decomposition (SVD)<n>We evaluate our approach extensively on standard continual learning benchmarks using both encoder-decoder (T5-Large) and decoder-only (LLaMA-2 7B) models.
  arXiv  Detail & Related papers  (2025-04-09T17:59:42Z)
• In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention [52.159541540613915]
  We study how multi-head softmax attention models are trained to perform in-context learning on linear data.<n>Our results reveal that in-context learning ability emerges from the trained transformer as an aggregated effect of its architecture and the underlying data distribution.
  arXiv  Detail & Related papers  (2025-03-17T02:00:49Z)
• Extreme Q-Learning: MaxEnt RL without Entropy [88.97516083146371]
  Modern Deep Reinforcement Learning (RL) algorithms require estimates of the maximal Q-value, which are difficult to compute in continuous domains. We introduce a new update rule for online and offline RL which directly models the maximal value using Extreme Value Theory (EVT) Using EVT, we derive our Extreme Q-Learning framework and consequently online and, for the first time, offline MaxEnt Q-learning algorithms.
  arXiv  Detail & Related papers  (2023-01-05T23:14:38Z)
• Enhancing Classifier Conservativeness and Robustness by Polynomiality [23.099278014212146]
  We show howconditionality can remedy the situation. A directly related, simple, yet important technical novelty we subsequently present is softRmax. We show that two aspects of softRmax, conservativeness and inherent robustness, lead to adversarial regularization.
  arXiv  Detail & Related papers  (2022-03-23T19:36:19Z)
• Sparse Attention with Linear Units [60.399814410157425]
  We introduce a novel, simple method for achieving sparsity in attention: we replace the softmax activation with a ReLU. Our model, which we call Rectified Linear Attention (ReLA), is easy to implement and more efficient than previously proposed sparse attention mechanisms. Our analysis shows that ReLA delivers high sparsity rate and head diversity, and the induced cross attention achieves better accuracy with respect to source-target word alignment.
  arXiv  Detail & Related papers  (2021-04-14T17:52:38Z)
• Direction Matters: On the Implicit Bias of Stochastic Gradient Descent with Moderate Learning Rate [105.62979485062756]
  This paper attempts to characterize the particular regularization effect of SGD in the moderate learning rate regime. We show that SGD converges along the large eigenvalue directions of the data matrix, while GD goes after the small eigenvalue directions.
  arXiv  Detail & Related papers  (2020-11-04T21:07:52Z)
• AdaS: Adaptive Scheduling of Stochastic Gradients [50.80697760166045]
  We introduce the notions of textit"knowledge gain" and textit"mapping condition" and propose a new algorithm called Adaptive Scheduling (AdaS) Experimentation reveals that, using the derived metrics, AdaS exhibits: (a) faster convergence and superior generalization over existing adaptive learning methods; and (b) lack of dependence on a validation set to determine when to stop training.
  arXiv  Detail & Related papers  (2020-06-11T16:36:31Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.