MetaLA: Unified Optimal Linear Approximation to Softmax Attention Map
Various linear complexity models, such as Linear Transformer (LinFormer), State Space Model (SSM), and Linear RNN (LinRNN), have been proposed to replace the conventional softmax attention in Transformer structures. However, the optimal design of these linear models is still an open question. In this work, we attempt to answer this question by finding the best linear approximation to softmax attention from a theoretical perspective. We start by unifying existing linear complexity models as the linear attention form and then identify three conditions for the optimal linear attention design: 1) Dynamic memory ability; 2) Static approximation ability; 3) Least parameter approximation. We find that none of the current linear models meet all three conditions, resulting in suboptimal performance. Instead, we propose Meta Linear Attention (MetaLA) as a solution that satisfies these conditions. Our experiments on Multi-Query Associative Recall (MQAR) task, language modeling, image classification, and Long-Range Arena (LRA) benchmark demonstrate that MetaLA is more effective than the existing linear models.
Code (1)
Tasks
image-classificationImage ClassificationLanguage ModelingLanguage ModellingMethods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
Linear $Q$-Learning Does Not Diverge: Convergence Rates to a Bounded Set
$Q$-learning is one of the most fundamental reinforcement learning algorithms. Previously, it is widely believed that $Q$-learning with linear function approximation (i.e., linear $Q$-learning) suffers from possible dive…
Q-LearningRethinking the Global Convergence of Softmax Policy Gradient with Linear Function Approximation
Policy gradient (PG) methods have played an essential role in the empirical successes of reinforcement learning. In order to handle large state-action spaces, PG methods are typically used with function approximation. In…
Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency
Linear attention has attracted interest as a computationally efficient approximation to softmax attention, especially for long sequences. Recent studies have explored distilling softmax attention in pre-trained Transform…
Speaking of Language: Reflections on Metalanguage Research in NLP
This work aims to shine a spotlight on the topic of metalanguage. We first define metalanguage, link it to NLP and LLMs, and then discuss our two labs' metalanguage-centered efforts. Finally, we discuss four dimensions o…
Inference on Optimal Dynamic Policies via Softmax Approximation
Estimating optimal dynamic policies from offline data is a fundamental problem in dynamic decision making. In the context of causal inference, the problem is known as estimating the optimal dynamic treatment regime. Even…
Causal InferenceDecision Makingvalid