Paper: Fast Multipole Attention: A Divide-and-Conquer Attention Mechanism for Long Sequences (ICLR-2024)

Reviewer comment (Experiments):
Discussion on other efficient attentions that also uses multi-resolutions is limited. For example, MRA attention [1] also provides a multi-resolution approximation of self-attention. Also, since the proposed method is an approximation to the full self-attention matrix, it would be better to compare the approximation quality of the proposed method with other approximation baselines.

Author response:
Yes, thanks for pointing us to MRA, which is an important contribution we were not aware about and is highly relevant. We now discuss similarities and differences with MRA and present numerical comparisons, see the response to the second referee.

In short, while MRA and FMA are based on similar high-level principles, the main advance of our FMA relative to MRA attention is that our FMA approach **learns the basis functions** that compute group quantities, as opposed to the fixed wavelet family used in MRA attention.

Our new numerical results show that learning the basis functions in hierarchical methods leads to substantially better accuracy for language tasks compared to MRA attention, for similar memory and time cost.

Furthermore, the comment on comparing approximation quality of the attention matrix while fixing projection weights $W_Q$, $W_K$ and $W_V$ (as obtained for the pretrained full attention, for example) is interesting. However, the bottomline in our tests shows that when we train $W_Q$, $W_K$ and $W_V$ with FMA, the resulting models are more accurate than doing this with MRA or H-matrix, so *learning* the basis functions that compute group quantities in FMA leads to a large accuracy improvement compared to the results of training MRA with fixed wavelet basis. We believe this final bottomline is more representative of the desired outcome (an accurate language model with low memory cost) than direct comparison with the attention matrix as trained for a full attention model (which is too expensive to compute in practice for long enough sequence lengths).

Reviewer comment (Evaluation):
Error analysis in 3.6 only analyzes the error between $QK^\top$ and $Q\tilde{K}^\top$, which is too rudimentary. How does the approximation impact the final output of self-attention?

Author response:
Thank you for the suggestion. We have improved our theoretical analysis and provided the two following results on the error estimations of the attention matrix and its exponential.

**Theorem 1 (Error Estimation of the Attention Matrix Approximation).**

Let $\widetilde{C}$ be our two-level FMA attention matrix with $p=1$. Suppose for every $j,\ell \in \{0,\ldots,m-1\},t\in \{0,1,2,3\},$ and $i\in\{0,\ldots, n-1\}$, we have 

$||Q_i^T||_2\leq\beta_1$ 

and 

$||K_{tm+j}^T - K_{tm+\ell}^T||_2\leq \beta_2,$

and the kernel weights are normalized in the sense that  $\sum\limits_{j=0}^{m-1}\alpha_j=1$.  Then $||\widetilde{C} -C||_F$ is bounded above by

$m\beta_1\beta_2\sqrt{6\sum\limits_{j=0}^{m-1}\alpha_j^2 }.$

Note that the results can be easily extended to higher level $L>1$ and $p>1$. 

In addition, inspired by the estimation derived in  Zeng et al. (2022), we have a similar estimation of $||\exp(\widetilde{C}) -\exp(C)||_F$ where our set of components for the approximation is obtained from the hierarchical structure (see Fig. 1b).

**Theorem 2.**

Let $\widetilde{C}$ be the approximated attention matrix obtained from our $(L+1)$-level FMA with fixed kernel weights and $\delta_k$ be the $m_k$-th largest coefficient. Suppose 

$||Q_i^T||_2\leq\beta_1$

and

$||K_{j_1}^T - K_{j_2}^T||_2\leq \beta_2,$

for $(i,j_1)$ and $(i,j_2)$ belong to the same level $I_\ell$, for $\ell = 1,\ldots L$ and $i\in\{0,\ldots, n-1\}$. Then 

$||\exp(\widetilde{C}) -\exp(C)||_F^2$

is bounded above by

$\sum\limits_{\ell=1}^LC_{2r}\delta_k^2 \left(3nm_k -6m_k^2 \right).$

Here $r = 2\beta_1\beta_2$ and $C_{2r} = 1+ \exp(2r)-\exp(r)$.

It is worth noting that if $\delta_k$ decays fast as $k$ increases (which we can see in fast multipole method, and can expect for 1D sequences), the error in the above estimation becomes small, illustrating the power of the fast multipole method for this kind of data.

Reviewer comment (Experiments):
Experiments are only performed on 512, 1K, and 2K sequence lengths, which are not very long. Full self-attention also performs efficiently on these sequence lengths (based on my knowledge and the efficiency comparison on Figure 3). The efficiency benefits of efficient attentions are very limited on these sequence lengths. The authors should evaluate the methods on longer sequences.

Author response:
We agree with the reviewer that linearly scaling transformers are more advantageous for long sequence lengts, and have added results for sequence length 4096, see response to the first reviewer above. This confirms that for long sequence lengths the full transformer indeed needs very large amounts of memory.

Reviewer comment (Experiments):
No performance comparison to full self-attention on Table 1 and Table 2. Fused full self-attention (such as Flash Attention [2]) is integrated to PyTorch 2. The authors should give a performance and efficiency comparison to the full self-attention.

Author response:
We have included comparison with full attention; see the response to the first reviewer. This confirms the advantages of efficient transformers.

Reviewer comment (Experiments):
On efficiency comparison, the batch size is 1, so compute load is small. In this case, kernel launch latency and memory latency might account for a large portion of the overall latency. Custom CUDA kernels usually have some advantages on these latency due to fused operations. It would be more convincing to provide latency for large batch size.

Author response:
Regarding this question, we use batch size 64 and include the numerical results in the updated Figure 2.

Reviewer comment (Experiments):
To the best of my knowledge, Linformer, Nystromformer, H-Transformer do not support causal attention, how are auto-regressive language modeling experiments performed on these baselines?

Author response:
Although the authors in Nystromformer and H-Transformer do not report numerical results on autoregressive language modeling tasks, causal masking of these models can be achieved by applying a suitable causal mask. For example, for H-attention, causal masking is achieved by applying a causal mask to the blocks of the approximated attention matrix at every level, as in https://github.com/lucidrains/h-transformer-1d.
 
In the original submission we mistakenly cited the Linformer paper, but we actually used [1].
 
[1] Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention (https://arxiv.org/abs/2006.16236)

Reviewer comment (Reproducibility):
Does this method support KV cache during inference, and will it introduce more overhead? When profiling inference time for Figure 3, is KV cache used?

Author response:
We haven't implemented KV cache in our current code. Since in our experiments there is no overlap in K and V in each iteration, there is, in fact, no need to apply KV cache.

Reviewer comment (Experiments):
The major issue with the paper has to do with the experimental evaluation. Although real world language modeling tasks are used, a major baseline is lacking, namely the vanilla transformer. This absence is even more important given that the setup is unusually small with a 6 layer transformer only. For instance, Reformer reports results with 12 layers on enwik8 which are significantly better at 1.19 bits per dimension. How would vanilla transformer and FMA compare at those sizes?

Author response:
We agree with the reviewer that more extensive experimental evaluation is warranted.

We have added the requested results in extended tables in the paper:
* full comparison with the vanilla transformer: **see extended Tables 1 and 2 (additions in blue), and updated Figure 2**

**It can be seen that full attention, as expected, has the best
accuracy but its memory consumption grows quadratically and becomes impractical for large sequence lengths.** Efficient transformers aim to obtain linear scaling in memory use while retaining good accuracy (so the models can use longer sequences).

* comparison with the memory compressed transformer: **see extended Table 1 (additions in blue)**

**Compressed attention reduces the memory by a factor of 3 but loses precision and still scales quadratically in memory; for large sequence lengths, it has worse accuracy and larger memory use than FMA.**

Reviewer comment (Experiments):
Another quite significant omission is the lack of any ablations. The method has several hyper parameters that change from experiment to experiment. There is absolutely no ablation study to provide insight on the effect of group size m and the number of levels on performance as well as the FLOPS and memory required.

Author response:
* an ablation study on the effect of group size m (and number of levels) and rank p: **see new Table 3 in supplement**

It can be seen, as expected, that bpc improves as group size m grows at the cost of larger memory use.
Similarly, bpc improves, as expected, as rank p grows at the cost of larger memory use.

FMA is quite flexible in allowing the user to tune m and p.

The choices for m and p in the experiments were suitable for the problem types and sizes we considered, compared to competing methods.

Reviewer comment (Experiments):
How does memory compressed transformer perform in the same experiments? Dividing the computational and memory cost by 3 would be quite enough for the tested sequence lengths.

Author response:
We have added the requested results in extended tables in the paper:
* comparison with the memory compressed transformer: **see extended Table 1 (additions in blue)**

**Compressed attention reduces the memory by a factor of 3 but loses precision and still scales quadratically in memory; for large sequence lengths, it has worse accuracy and larger memory use than FMA.**

Reviewer comment (Reproducibility):
Minor, but why would the causal linear attention require more memory than the non-causal version? This seems like a bug in the implementation.

Author response:
**Question on memory use of linear attention:**

Regarding the question about the memory of the causal linear attention and the non-causal one, we use the codes from https://github.com/lucidrains/linear-attention-transformer
and report the memory. There could be a bug in the implementation. We will investigate this issue in the future.