only O(d log T ) memory with the same time complexity O(T 2d). The core idea behind
memory eﬃcient attention is similar to “lazy softmax” [50] where the denominator
of the softmax for the dot product of queries and keys can be calculated in the later
stage. KDEFormer [51] suggests reducing the denominator of the softmax function to
a variant of the kernel density estimation (KDE) problem, and an eﬃcient KDE solver
can be further utilized to accelerate the multiplication of the attention matrix with the
value matrix. The trick is based on reducing the number of columns of the attention
matrix (referred to as A := exp(QK ⊤/√d)) using importance sampling. KDEFormer
delivers a O(T md) time complexity and O(T m) memory complexity, where m < T
is a small number. MEGA [52] introduces a moving average equipped gated attention
Beyond Efficiency