Efficient Attention via Control Variates
This work addresses the efficiency-accuracy trade-off in attention mechanisms for machine learning practitioners, offering a more flexible and effective solution.
The paper tackles the approximation gap between random-feature-based attention (RFA) and softmax attention by using control variates, resulting in a novel attention mechanism that outperforms state-of-the-art efficient methods on vision and language tasks.
Random-feature-based attention (RFA) is an efficient approximation of softmax attention with linear runtime and space complexity. However, the approximation gap between RFA and conventional softmax attention is not well studied. Built upon previous progress of RFA, we characterize this gap through the lens of control variates and show that RFA can be decomposed into a sum of multiple control variate estimators for each element in the sequence. This new framework reveals that exact softmax attention can be recovered from RFA by manipulating each control variate. Besides, it allows us to develop a more flexible form of control variates, resulting in a novel attention mechanism that significantly reduces the approximation gap while maintaining linear complexity. Extensive experiments demonstrate that our model outperforms state-of-the-art efficient attention mechanisms on both vision and language tasks.