Value-aware Approximate Attention (2021)

Abstract: Following the success of dot-product attention in Transformers, numerous approximations have been recently proposed to address its quadratic complexity with respect to the input length. However, all approximations thus far have ignored the contribution of the *value vectors* to the quality of approximation. In this work, we argue that research efforts should be directed towards approximating the true output of the attention sub-layer, which includes the value vectors. We propose a value-aware objective, and show theoretically and empirically that an optimal approximation of a value-aware objective substantially outperforms an optimal approximation that ignores values, in the context of language modeling. Moreover, we show that the choice of kernel function for computing attention similarity can substantially affect the quality of sparse approximations, where kernel functions that are less skewed are more affected by the value vectors.

## introduction the transformer architecture @xcite has been widely successful in a wide range of natural language processing tasks, including machine translation @xcite , language modeling @xcite , question-answering @xcite , and many more. transformers pre-trained on large amounts of text with a language modeling (lm) objective have become the standard in nlp, exhibiting surprising amounts of linguistic and world knowledge @xcite @xcite @xcite . the contextualizing component of the transformer is the attention layer where all positions in an input sequence of length l aggregate information from the entire sequence in parallel. at its core, given l query, key and value vec-tors, the dot-product attention function outputs foot_0 softmax(qk )v where the softmax function is applied row-wise on the matrix @xmath0, consisting of similarity scores of the query-key pairs. unfortunately, computing ω(l • l) similarity scores is prohibitive for long sequences. to alleviate this, past work proposed to compute an approximation of softmax(qk ). one major line of research focused on sparse attention variants, where only a few similarity scores are computed per position, and the rest are ignored. methods differ by which query-key pairs are selected @xcite @xcite @xcite @xcite . a second line of research explored dense variants @xcite @xcite (cf. @xcite for a survey). e.g., instead of computing the attention scores exactly for only a few querykey pairs, @xcite ) compute an approximation of scores for all pairs. in this work, we point to a lacuna in current research on efficient transformers. while recent work focused on approximating the attention scores softmax(qk ), the true target of approximation should be the output of the attention sub-layer, namely @xmath1, which also includes the value vectors, v . we show that ignoring value vectors leads to unwarranted consequences both theoretically and empirically. to demonstrate the importance of value-aware approximation, we analyze optimal sparse attention, that is, the case where, in hindsight, the model computes dot product similarity only with the most similar key vectors, while still ignoring the value vectors. we show that in the popular masked language modeling (mlm) setup, optimal sparse attention dramatically under-performs compared to an optimal approximation of the true output of the attention sub-layer, h, leading to an error increase of 8-20 points. next, by theoretically focusing on the case where queries compute similarity to the single most similar key vector, we show that approximating softmax(qk ) is equivalent to approximating h when the value vectors v satisfy strong orthogonality and norm constraints. conversely, when they do not, ignoring v can lead to an unbounded approximation error. second, we discuss the kernel-based view of attention, where efficiency is gained by replacing the exponential kernel (corresponding to softmax) with other kernel functions @xcite . we theoretically show that while in the exponential kernel case (corresponding to softmax), the effect of the norm of the value vectors is potentially small, switching to other kernels can dramatically increase the importance of the value vectors. we empirically test this by comparing optimal sparse attention given different kernel functions, and see that indeed approximation quality decreases when replacing the exponential kernel, to conclude, we theoretically and empirically show that approximating the attention score matrix alone is insufficient, and propose that the research community should instead approximate the true output of the sub-attention layer, which importantly includes value vectors. our code and trained models are available at https://github.com/ ag1988/value_aware_attn . ## background we review the kernel-based view of attention @xcite , which will be instructive in §3. ## optimal sparse attention prior methods for approximating attention have ignored the contribution of the values vectors v . as the true output of the attention sub-layer also depends on v , a natural question is whether it is possible to design better approximation methods by incorporating v , and if so, how much improvement is even possible? to answer this, we focus on sparse attention, and analyze the difference between an oracle sparse approximation that considers the value vectors, and an oracle approximation that does not. that is, we look at the difference between the two approximations from the perspective of expressivity, ignoring any memory and computational constraints. we denote an optimal value-aware approximation that uses r key vectors per query by optimal-vaware-r, and an optimal approximation that ignores value vectors by optimal-v-oblivious-r. we define optimal-v-oblivious-r as the output of eq. 2 in which s is selected to be the r indices with the highest attention scores α i 's. this is a natural baseline since this is what current sparse methods are trying to emulate. we now explicitly derive and analyze the value-aware objective. ## experiments we empirically verify our observations in the context of training causal and masked language models, which are known to strongly correlate with performance on downstream applications @xcite . ## conclusions in this work, we provide theoretical and empirical evidence against current practice of focusing on approximating the attention matrix in transformers, while ignoring the value vectors. we propose a value-aware objective and argue that the efforts to develop more efficient transformers should consider this objective function as a research target.