🤖 AI Summary
This work addresses the limited expressiveness of linear self-attention in in-context learning. We propose an extended linear self-attention mechanism augmented with a learnable bias matrix, overcoming the fundamental limitation of standard linear attention—which can only approximate matrix multiplication—while preserving O(n) time complexity. For the first time, our mechanism exactly represents key linear operators, including precise matrix multiplication and identity mapping (i.e., skip connections). Through formal matrix-theoretic analysis and heuristic input-structure design, we prove its capacity to represent constant matrices, input matrices, and arbitrary two- or three-matrix products. Furthermore, we construct an interpretable, batched ridge regression gradient-descent module grounded in this mechanism; empirical evaluation on few-shot regression tasks demonstrates algorithm-level generalization. This establishes a foundation for native, interpretable algorithmic reasoning within Transformer architectures.
📝 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.