An extension of linear self-attention for in-context learning
This addresses a specific bottleneck in transformer architectures for in-context learning, offering a flexible method for matrix manipulations, but it is incremental as it builds on existing linear self-attention.
The paper tackles the limitation of naive self-attention for in-context learning by extending linear self-attention with a bias matrix, enabling it to output constant matrices, input matrices, and multiplications of up to three matrices, and demonstrates this with a heuristic construction for batch gradient descent in ridge regression.
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.