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

Abstract: Advancing the frontier of subquadratic architectures for LLMs (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.

• The paper enhances in-context learning by integrating learnable kernel functions into linear transformer architectures, reducing computational cost on long sequences.
• It introduces a quadratic kernel with trainable parameters and normalization, achieving improved perplexity metrics on associative recall tasks.
• Empirical evaluations on MQAR and language modeling demonstrate ReBased’s superior performance over sub-quadratic models and its predecessor.

ReBased: Elevating In-Context Learning in Linear Transformers through Learnable Kernel Functions

Introduction

In the ever-expanding field of NLP, the proficiency of LLMs in capturing long-context dependencies is paramount. Despite the widespread acclaim of Transformer-based models for their effectiveness, the quadratic computational cost associated with their attention mechanisms poses a significant bottleneck, especially for long sequences. This has prompted the exploration of alternative architectures, notably Linear Transformers and State Space Models (SSMs), to reduce computational demands while striving to maintain or surpass the performance of traditional Transformers. The Based model, a notable advancement in this direction, employs a hybrid architecture integrating Linear Transformers with a kernel function inspired by the Taylor expansion of the exponential function, showing potential in tackling the in-context learning limitations of its predecessors. This paper introduces ReBased, a modification to the Based model that further exploits the advantages of learnable kernel functions and normalization to enhance its in-context learning capabilities and overall performance.

Recent Developments

The quest for efficiency in processing long text sequences has led to the exploration of various alternatives to the Transformer architecture. Linear Transformers propose a kernel-based approach to address the quadratic computational complexity, while SSMs offer a simplified structure that avoids activation functions across time. Despite these innovations, challenges persist in matching the original Transformer's performance, especially in tasks requiring retrieval of information from long sequences.

ReBased: A Novel Approach

ReBased emerges as a novel variant of the Linear Transformer, targeting the limitations inherent in the Based model. By revisiting the kernel function used in Based and introducing learnable parameters alongside normalization, ReBased aims to provide a more flexible and effective method for in-context learning. These enhancements facilitate the model's ability to diminish attention scores to zero for irrelevant token pairs, thus potentially improving performance on tasks involving long sequences.

Methodology

ReBased's architecture roots in modifying the Based model's kernel function, leveraging a quadratic function with trainable parameters to achieve a more adaptable attention mechanism. The inclusion of normalization steps before applying this kernel function is a critical innovation, drawing parallels with the benefits observed from Layer Normalization in improving model training dynamics.

Experimental Insights

The empirical evaluation involved two primary tasks: the Multi-Query Associative Recall (MQAR) task and language modeling on the Pile dataset. Across different contexts and model sizes, ReBased demonstrated advantageous performance over the Based model and other sub-quadratic architectures. Notably, ReBased showed a better capability to manage associative dependencies, indicated by improved perplexity metrics in both associative recall and non-associative contexts when compared to its predecessors.

Theoretical Implications and Future Prospects

The development of ReBased underscores the potential of leveraging learnable kernel functions in Linear Transformer architectures to enhance in-context learning capabilities. This approach marks a significant step toward addressing the scalability challenges associated with processing long sequences, paving the way for more efficient and powerful NLP models. Future investigations might explore the integration of complex kernel functions and further refinements in normalization techniques to unlock new levels of performance in language modeling and beyond.

Concluding Remarks

In summary, ReBased represents a promising advancement in the development of efficient architectures for NLP tasks, particularly those requiring the processing of extensive contexts. By innovatively refining the kernel function and incorporating normalization, ReBased enhances the in-context learning strengths of Linear Transformers. This paper not only highlights the potential of ReBased in bridging the performance gap with the traditional Transformer model but also sets the stage for future explorations in optimizing model efficiency and effectiveness for large-scale NLP applications.