An extension of linear self-attention for in-context learning
Katsuyuki Hagiwara
TL;DR
Addressing limitations of naive self-attention for in-context learning, the paper introduces extended linear self-attention (ELSA) that augments the input transformation with a bias term to enable flexible matrix computations. The core mechanism is $ELSA({\bf H}) = ({\bf H}{\bf W}_3+{\bf B}_3)({\bf H}{\bf W}_1+{\bf B}_1)^{\top}({\bf H}{\bf W}_2+{\bf B}_2)$, which can realize outputs such as a constant matrix, the input itself, or products of two or three matrices, enabling skip connections. The authors demonstrate a heuristic implementation of a batch-type gradient descent for ridge regression using ELSA under two input forms: a designed input form and a naturally enumerated input form, all within the ELSA framework. This work highlights the potential of extended linear primitives for in-context computation in transformer-like models and suggests future work on training dynamics, nonlinear extensions, and broader algorithmic demonstrations.
Abstract
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.