Spectraformer: A Unified Random Feature Framework for Transformer

2024-05-24

arXiv.org

AI Summary

Existing random feature methods for linearizing Transformer attention lack systematic kernel function design and evaluation. Method: This paper proposes Spectraformer—the first comparable, scalable, unified framework that explicitly models kernel function structure in linearized attention. By jointly designing nonlinear component functions (e.g., ReLU, exp, sin) and structured or random weight matrices, it systematically uncovers the complementarity and task dependency of different kernels for long-range text modeling. Contribution/Results: Evaluated on all three text-based Long-Range Arena (LRA) benchmarks, Spectraformer demonstrates that kernel selection critically impacts accuracy—achieving significant improvements while preserving linear time and space complexity. The framework is fully open-sourced.

Technology Category

• Machine Learning: Kernel Methods
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

Linearization of attention using various kernel approximation and kernel learning techniques has shown promise. Past methods use a subset of combinations of component functions and weight matrices within the random features paradigm. We identify the need for a systematic comparison of different combinations of weight matrices and component functions for attention learning in Transformer. In this work, we introduce Spectraformer, a unified framework for approximating and learning the kernel function in linearized attention of the Transformer. We experiment with broad classes of component functions and weight matrices for three textual tasks in the LRA benchmark. Our empirical findings indicate that different kernels are good at different tasks and that kernel choice is fundamental to performant models. Our code is available at: https://github.com/dukenguyenxyz/spectraformer .

Problem

• Systematic comparison of weight matrices and component functions for Transformer attention
• Unified framework for approximating and learning kernel functions in attention mechanisms
• Achieving state-of-the-art performance with random feature-based efficient Transformers

Innovation

• Unified framework for Transformer kernel approximation
• Combines weight matrices and component functions systematically
• Achieves state-of-the-art performance on Long Range Arena