Linear Transformers with Learnable Kernel Functions are Better In-Context Models (2024)

Abstract: Advancing the frontier of subquadratic architectures for Language Models (LMs) is crucial in the rapidly evolving field of natural language processing. Current innovations, including State Space Models, were initially celebrated for surpassing Transformer performance on language modeling tasks. However, these models have revealed deficiencies in essential In-Context Learning capabilities – a domain where the Transformer traditionally shines. The Based model emerged as a hybrid solution, blending a Linear Transformer with a kernel inspired by the Taylor expansion of exponential functions, augmented by convolutional networks. Mirroring the Transformer’s in-context adeptness, it became a strong contender in the field. In our work, we present a singular, elegant alteration to the Based kernel that amplifies its In-Context Learning abilities evaluated with the Multi-Query Associative Recall task and overall language modeling process, as demonstrated on the Pile dataset.

## introduction large language models (llms) are revolutionizing the field of natural language processing and establishing new benchmarks across various tasks @xcite . nevertheless, despite their triumphs, most of these models are built on transformer frameworks that employ attention mechanisms. these mechanisms scale poorly with long text sequences, leading to impractical computational complexity for extending contextual processing @xcite . to address this constraint, several alternatives to transformers were proposed. @xcite suggested replacing the exponential function in the attention mechanism with the kernel function to change the order of computations and thus move away from quadratic complexity of the * * work done while at tinkoff. sequence length. however, when compared to vanilla transformers, this approach leads to a drop in performance. furthermore, the kernel function selection is a topic still in need of consideration. an alternative way to define a linear model is to utilize state space models (ssms) @xcite @xcite , which are capable of producing quality that is comparable to transformers when measured with perplexity on language modeling. notably, both linear transformers @xcite and ssms can be described as recurrent neural networks (rnns) @xcite , which have their limitations when it comes to managing lengthy dependencies within texts since memory capacity can be overrun as the volume of information increases. additionally, while the hidden state of rnns is larger for linear transformers than for ssms, the latter showed higher text modeling quality. the introduction of the based model @xcite attempted to address the abovementioned challenges by utilizing a hybrid architecture @xcite based on a linear transformer with a novel kernel function derived from a taylor expansion of an exponential function. @xcite demonstrated that the based model was less prone to performance issues when working with longer content than other models when assessed on the multi-query associative recall (mqar) task. nonetheless, even the based model experiences a drop in performance when faced with extensive contexts relative to the conventional transformer architecture. a profound comprehension of the processes occurring within the based architectures is essential for their advancement. upon examining how attention scores are distributed, we argue that the kernel function previously adopted in based cannot be considered optimal, resulting in limitations when dealing with lengthy context and small model when testing our enhanced architecture on the mqar task, we found that rebased surpasses the original based model across a variety of contexts and model sizes. additionally, after training with the pile dataset @xcite , we observed that rebased performs better than its predecessor at in-context learning and excels at modeling associative dependencies measured through improved perplexity metrics. ## recent work the vanilla transformer architecture @xcite , although widely used in nlp @xcite @xcite , suffers from growing computational and memory demands (o (𝑑 * 𝑁 2 ) as sequence lengths (𝑁) and head size (𝑑) increase). while this is not much of a problem when it comes to shorter sequences, it becomes a significant bottleneck when working with longer ones. several alternative architectures were proposed to address this issue. @xcite suggested substituting the attention mechanism's exponential function, which is meant to measure the similarity between queries and keys, with a product of kernel functions that can be separately evaluated for queries and keys. this kernel-based approach reshapes the computation within the attention mechanism, cutting the time and memory complexity to o (𝑑 2 * 𝑁). additionally, during inference, it supports sampling sequences with linear complexity regarding length, similar to rnns @xcite . in a different approach, state space models (ssms) borrow from control theory to offer a simplified structure akin to rnns, but without activation functions across time @xcite @xcite . the mamba model, also known as s6 @xcite , stands out in this category, displaying enhanced learning of shortterm dependencies in texts compared to existing pre-trained llms @xcite . despite these advancements, there is no standard way to fully evaluate these innovative architectures to assess their performance limits. one standard evaluation method is to pre-train a language model and assess its perplexity with a given dataset, but this may not truly reflect the model's ability to manage long context dependencies. another option is to use the long range arena (lra) benchmark, which involves classification tasks with long input sequences. though some new models have outperformed transformers in the lra, it is believed that the benchmark is capable of introducing bias in the comparison @xcite . one promising evaluation approach is to test an architecture's in-context learning abilities. @xcite introduced the concept of associative recall (ar), a task where the model learns to copy a token from a sequence after a certain point. however, while in fu et al. (2023a) the associative recall task was implemented with a goal to retrieve only one token, @xcite noted that this task could be considered overly simplistic. this led to the creation of the multi-query associative recall (mqar) task, which requires retrieving multiple tokens from the context. findings on mqar indicate that while newer models may compete with transformer in terms of perplexity, they can still struggle with long contexts at small model sizes because of their limited in-context learning capabilities. meanwhile, transformers remain robust against such factors. lastly, @xcite introduced linear transformer with a new kernel function (namely based), showcasing enhanced performance on the mqar task when compared to mamba. despite this improvement, compared to traditional transformers, the problem of decline in performance when handling long sequences with smaller models still remains. addressing this challenge is the primary goal of our paper. ## background 3.1 linear transformers to fully grasp the based architecture, it is vital to first discuss the original transformer model. the attention mechanism, which is central to the transformer's functionality, evaluates the output 𝑦 𝑖 for each position 𝑖 as follows to address this scalability problem, katharopoulos et al. (2020) suggested replacing the direct computation of similarity between 𝑄 and 𝐾 with a transformation through a non-linear kernel function 𝜙(•). this allows for the following approximation: sim(𝑄 𝑖 , 𝐾 𝑗 ) ≈ 𝜙 𝑇 (𝑄 𝑖 )𝜙(𝐾 𝑗 ). by implementing this kernel, the linear transformer computes 𝑦 𝑖 as by calculating 𝜙(𝐾 𝑗 )𝑉 𝑇 𝑗 ∈ r 𝑑×𝑑 upfront, the complexity of the attention mechanism transitions to linear with the sequence length, addressing the inefficiencies of the original model. ## revisiting based in our study, we explore the fundamental requirements for kernel functions. we examine the exponential function and its approximate representation, as depicted in figure 2 . we observe a limitation in the approximation since its minimal value is fixed at 0.5. this is problematic for handling long sequences, as it is difficult to assign a near-zero attention score to specific token pairs. ideally, we want to be able to diminish the attention scores we propose to learn the scale and shift of the parabola jointly with the model and make it possible to zero out the similarity value. see section 4 for more details and section 5.1 for experimental setup description. to zero, which would require significantly larger values elsewhere in the normalization process with the based model. to rectify this issue, a straightforward approach would be to adjust the lowest point of the kernel function to zero. however, this solution prompts us to ask why the minimum value of the kernel function should occur precisely at 𝑞 𝑇 𝑘 = -1. as used in the original transformer, the traditional exponential similarity function increases monotonically, but the quadratic kernel has an optimal value to which it decreases and then ascends from. therefore, to decrease attention in the transformer, one would aim to minimize 𝑞 𝑇 𝑘. in contrast, the ideal 𝑞 𝑇 𝑘 should be exactly -1 for the based method. otherwise, the attention score would increase. this condition may induce less-than-ideal training outcomes and degrade the model's accuracy. these challenges lead us to conjecture that if the quadratic kernel is used to calculate the similarity between 𝑞 and 𝑘, we must consider the range of potential 𝑞 𝑇 𝑘 values and create adjustable parameters for the parabolic function to align with these values during training. simplifying for clarity, let us look at a one-dimensional scenario. we can express the trainable parameters of the kernel function in relation to the affine transformation of 𝑞 and 𝑘 as such here, 𝜙(•) represents a quadratic function. the model can learn any quadratic function with a determined minimum value by adjusting its parameters. we can, therefore, simplify the kernel function to interestingly, our transformation resembles the application of layer normalization @xcite , minus the normalization itself. we hypothesize whether normalizing 𝑞 and 𝑘 before the kernel function could improve the model's performance. our suspicion is confirmed when normalization enhances results, as demonstrated in a later ablation study. consequently, our refined rebased model incorporates layer normalization. in the following sections, we provide an in-depth analysis and conduct comprehensive experiments to validate the effectiveness of these modifications. ## conclusion and future work in this paper, we present rebased, a novel architecture for sub-quadratic attention computation. for our model, we analyzed the base architecture and proposed to develop it even further by using polynomial kernels with learnable parameters and adding normalization before the kernel evaluation. while incorporating layer normalization into model training was attempted previously @xcite , our method integrates this normalization directly into the kernel function. with this simple architectural change, we achieved results that outperformed based on mqar and language modeling with the pile dataset tasks. we analyzed the internal representations of re-based, based, and vanilla attention modules, and concluded that rebased resembles attention more than based. notably, while based uses a taylor expansion of an exponential function, a rebased kernel function is different from the exponent but shows better performance. our research suggests that using a second-order polynomial might be insufficient for the best performance, and indicates that more sophisticated learnable kernels could be utilized to improve the performance of trained models. normalization could further improve various kernel functions. this highlights a need for researchers to revisit kernel-based methods with the goal of enhancing their adaptability and efficiency. our findings reveal a disparity in handling the mqar task between attention-based models and others such as based, specifically as sequence lengths increase. attention models excel on longer sequences, significantly outperforming their nonattention counterparts. these results highlight the necessity of further research into strategies that could bridge this gap in order to reach the performance of attention-based methods. perhaps the superior aspects of attention mechanisms could be matched or exceeded by other models, especially on tasks that require associative recall, such as machine translation @xcite . future research could give insight into this, leading to improved models for processing long sequences.