# An extension of linear self-attention for in-context learning

arXiv: https://arxiv.org/abs/2503.23814 · HTML: https://mlsanity.com/paper/2503.23814/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 14 of its 22 references that are in the map.
- First posted on arXiv 31 Mar 2025.

## Builds on (most cited first, 14 of 14)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Error bounds for approximations with deep ReLU networks](https://mlsanity.com/paper/1610.01145/) (2016)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [A Survey on Mixture of Experts in Large Language Models](https://mlsanity.com/paper/2407.06204/) (2024)
- [Optimal approximation of piecewise smooth functions using deep ReLU neural networks](https://mlsanity.com/paper/1709.05289/) (2017)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Why Deep Neural Networks for Function Approximation?](https://mlsanity.com/paper/1610.04161/) (2016)
- [Automatic Domain Adaptation by Transformers in In-Context Learning](https://mlsanity.com/paper/2405.16819/) (2024)

## Abstract (arXiv)

In-context learning is a remarkable property of transformers and has been the focus of recent research. An attention mechanism is a key component in transformers, in which an attention matrix encodes relationships between words in a sentence and is used as weights for words in a sentence. This mechanism is effective for capturing language representations. However, it is questionable whether naive self-attention is suitable for in-context learning in general tasks, since the computation implemented by self-attention is somewhat restrictive in terms of matrix multiplication. In fact, we may need appropriate input form designs when considering heuristic implementations of computational algorithms. In this paper, in case of linear self-attention, we extend it by introducing a bias matrix in addition to a weight matrix for an input. Despite the simple extension, the extended linear self-attention can output any constant matrix, input matrix and multiplications of two or three matrices in the input. Note that the second property implies that it can be a skip connection. Therefore, flexible matrix manipulations can be implemented by connecting the extended linear self-attention components. As an example of implementation using the extended linear self-attention, we show a heuristic construction of a batch-type gradient descent of ridge regression under a reasonable input form.