Paper: Tree Cross Attention (ICLR-2024)

Reviewer comment (Evaluation):
In the evaluation, the authors focus on %tokens metric. How does that translate to wall clock speed up? Or does the tree structure introduce operations that are hard to take advantage of by modern hardware?

Author response:
We would like to thank the reviewer for the comments and helpful feedback.

> In the evaluation, the authors focus on % tokens metric. How does that translate to wall clock speed up? 

The % tokens do not directly translate to wall clock time. In this work, we focused on % tokens as the metric since it is (1) neither hardware nor implementation specific and (2) % tokens translates towards memory which is a major limiting factor scalability, particularly for (i) attention-based models which are memory intensive and (ii) low-memory devices which are popular (e.g., IoT devices). In contrast, wall clock time (runtime) is highly dependent on the efficiency of the implementation and the hardware. Nonetheless, in the table below, we compare the wall-clock time of the modules used during inference by ReTreever, Transformer+Cross Attention, and Perceiver IO, i.e., Tree Cross Attention, Cross Attention, and Perceiver IO's CA respectively.

We evaluate the modules on the Copy Task for both a GPU and CPU. Although Tree Cross Attention is slower, all methods only require milliseconds to perform inference. Both Tree Cross Attention and Cross Attention learn to solve the task perfectly. However, Tree Cross Attention uses only $3.5$% of the tokens. As such, Tree Cross Attention trades off between the number of tokens and computational time. Perceiver IO, however, fails to solve the task for the same number of tokens. 

|         Model        | % Tokens |     Accuracy    | CPU time (in Milliseconds) | GPU Time (in Milliseconds) |
|:--------------------:|:--------:|:---------------:|:--------------------------:|:--------------------------:|
|    Cross Attention   |  100.0%  | **100.0 ± 0.0** |            12.05           |            1.61            |
| Tree Cross Attention | **3.5%** | **100.0 ± 0.0** |            19.31           |            9.09            |
|   Perceiver IO's CA  | **3.5%** |    13.4 ± 0.2   |          **10.98**         |          **1.51**          |



> Or does the tree structure introduce operations that are hard to take advantage of by modern hardware?


The tree approach has a sequential aspect to it related to the height. However, the major aspects of it can still be parallelized on GPUs. For example, during the tree construction, left and right subtrees can be constructed at the same time since they do not share children nodes. Aggregation can also be performed at the same time for nodes at the same depth. Our implementation leverages parallelization in some aspects of the tree approach, but we have not utilized full-parallelization which can lead to further code optimization and runtime efficiency. 

We will be sharing our implementation alongside the camera-ready. We hope it will help improve the efficiency of future implementations of ReTreever and Tree Cross Attention.

Reviewer comment (Evaluation):
My main question regarding the evaluation is in the weakness section.

Author response:
We would like to thank the reviewer for the comments and helpful feedback.

> In the evaluation, the authors focus on % tokens metric. How does that translate to wall clock speed up? 

The % tokens do not directly translate to wall clock time. In this work, we focused on % tokens as the metric since it is (1) neither hardware nor implementation specific and (2) % tokens translates towards memory which is a major limiting factor scalability, particularly for (i) attention-based models which are memory intensive and (ii) low-memory devices which are popular (e.g., IoT devices). In contrast, wall clock time (runtime) is highly dependent on the efficiency of the implementation and the hardware. Nonetheless, in the table below, we compare the wall-clock time of the modules used during inference by ReTreever, Transformer+Cross Attention, and Perceiver IO, i.e., Tree Cross Attention, Cross Attention, and Perceiver IO's CA respectively.

We evaluate the modules on the Copy Task for both a GPU and CPU. Although Tree Cross Attention is slower, all methods only require milliseconds to perform inference. Both Tree Cross Attention and Cross Attention learn to solve the task perfectly. However, Tree Cross Attention uses only $3.5$% of the tokens. As such, Tree Cross Attention trades off between the number of tokens and computational time. Perceiver IO, however, fails to solve the task for the same number of tokens. 

|         Model        | % Tokens |     Accuracy    | CPU time (in Milliseconds) | GPU Time (in Milliseconds) |
|:--------------------:|:--------:|:---------------:|:--------------------------:|:--------------------------:|
|    Cross Attention   |  100.0%  | **100.0 ± 0.0** |            12.05           |            1.61            |
| Tree Cross Attention | **3.5%** | **100.0 ± 0.0** |            19.31           |            9.09            |
|   Perceiver IO's CA  | **3.5%** |    13.4 ± 0.2   |          **10.98**         |          **1.51**          |



> Or does the tree structure introduce operations that are hard to take advantage of by modern hardware?


The tree approach has a sequential aspect to it related to the height. However, the major aspects of it can still be parallelized on GPUs. For example, during the tree construction, left and right subtrees can be constructed at the same time since they do not share children nodes. Aggregation can also be performed at the same time for nodes at the same depth. Our implementation leverages parallelization in some aspects of the tree approach, but we have not utilized full-parallelization which can lead to further code optimization and runtime efficiency. 

We will be sharing our implementation alongside the camera-ready. We hope it will help improve the efficiency of future implementations of ReTreever and Tree Cross Attention.

Reviewer comment (Experiments):
**Speedup by attending to fewer context tokens**: One claimed benefit of the method is that it attends to fewer tokens to context when performing the task, which I assume would result in an inference speedup. But the work does not explicitly measure if TCA runs faster than CA. On the other hand, TCA does in most of the cases leads to a slight performance degradation, and it will be important to justify the design with proper inference wall-clock time measurements.

Author response:
In this work, we focused on % tokens as the metric since it is (1) neither hardware nor implementation specific and (2) % tokens translates towards memory which is a major limiting factor scalability, particularly for (i) attention-based models which are memory intensive and (ii) low-memory devices which are popular (e.g., IoT devices). In contrast, wall clock time (runtime) is highly dependent on the efficiency of the implementation and the hardware. Nonetheless, in the table below, we compare the wall-clock time of the attention mechanisms.


We evaluate the modules on the Copy Task for both a GPU and CPU. Although Tree Cross Attention is slower, all methods only require milliseconds to perform inference. Both Tree Cross Attention and Cross Attention learn to solve the task perfectly. However, Tree Cross Attention uses only $3.5$% of the tokens. As such, Tree Cross Attention trades off between the number of tokens and computational time. Due to space limitations, we have added these results to the appendix. 

We will be sharing our implementation alongside the camera-ready. We hope it will help improve the efficiency of future implementations of ReTreever and Tree Cross Attention.

|         Model        | % Tokens |     Accuracy    | CPU time (in Milliseconds) | GPU Time (in Milliseconds) |
|:--------------------:|:--------:|:---------------:|:--------------------------:|:--------------------------:|
|    Cross Attention   |  100.0%  | 100.0 ± 0.0 |            12.05           |            1.61            |
| Tree Cross Attention | 3.5% | 100.0 ± 0.0 |            19.31           |            9.09            |

Reviewer comment (Experiments):
**More general evaluation**: The authors claim to have proposed this general architecture, but the evaluating tasks are specific and not as general as expected. The tasks evaluated in the baseline work — IO perceiver (e.g., pretraining MLM, optical flow and multimodal encoding) seem harder and more general than the ones performed in this paper, (i.e., copy task, gp regression on image completion and time series on human activities). It would be nice to see experimental setups with more significant implications like pre-training, yet there seems to be significant amount of work to be put into actually scale it to these real settings.

Author response:
We initially considered evaluating on Perceiver IO's tasks. However, Perceiver IO's experiments used a very large amount of computing resources for their experiments which was significantly outside our capabilities. For example, their experiments used "a batch size of 1024 and 64 TPUs"  (Appendix A.1 in Perceiver IO's paper), "a larger batch size of 8192 and train on 256 TPUs" (Appendix A.2 in Perceiver IO's paper), and "1024 batch size on 64 TPUs"  (Appendix A.2 in Perceiver IO's paper). Specifically for their pre-training experiments in Perceiver IO, the authors used 256 TPUs. In contrast, each of our experiments was run on 1 GPU at a time (Nvidia P100 GPU (16 GB)). 

To ensure our experiments were general and with significant implications, we considered (1) different forms of data such as images and time series and (2) popular and well-benchmarked tasks within their respective literatures:  Neural Processes (GP Regression and Image Completion) and Time Series (Human Activity). Time Series models are popular in various important fields such as healthcare, climate science, ecology, and astronomy. Similarly, Neural Processes have been applied to a wide array of impactful settings (see the recent survey (Jha et al., 2023) for a full list) such as climate modelling and detecting biomarkers in fMRI scans. 

The datasets we considered are the standard benchmark datasets for evaluating the performance of Neural Processes (GP regression and image completion) and Time Series (Human Activity) models. Our results are competitive with that of state-of-the-art while requiring only a logarithmic number of tokens. Due to space limitations, these comprisons were included in the Appendix. Specifically, the tables in the Appendix include results comparing the performance of ReTreever with a total of $18$ baselines: $9$ baselines for Neural Processes and $9$ baselines for Time Series.


---

Saurav Jha, Dong Gong, Xuesong Wang, Richard E. Turner, Lina Yao. "The neural process family: Survey, applications and perspectives." arXiv preprint arXiv:2209.00517, 2023.

Reviewer comment (Evaluation):
Does the algorithm run faster than full attention in terms of wall-clock time?

Author response:
In this work, we focused on % tokens as the metric since it is (1) neither hardware nor implementation specific and (2) % tokens translates towards memory which is a major limiting factor scalability, particularly for (i) attention-based models which are memory intensive and (ii) low-memory devices which are popular (e.g., IoT devices). In contrast, wall clock time (runtime) is highly dependent on the efficiency of the implementation and the hardware. Nonetheless, in the table below, we compare the wall-clock time of the attention mechanisms.


We evaluate the modules on the Copy Task for both a GPU and CPU. Although Tree Cross Attention is slower, all methods only require milliseconds to perform inference. Both Tree Cross Attention and Cross Attention learn to solve the task perfectly. However, Tree Cross Attention uses only $3.5$% of the tokens. As such, Tree Cross Attention trades off between the number of tokens and computational time. Due to space limitations, we have added these results to the appendix. 

We will be sharing our implementation alongside the camera-ready. We hope it will help improve the efficiency of future implementations of ReTreever and Tree Cross Attention.

|         Model        | % Tokens |     Accuracy    | CPU time (in Milliseconds) | GPU Time (in Milliseconds) |
|:--------------------:|:--------:|:---------------:|:--------------------------:|:--------------------------:|
|    Cross Attention   |  100.0%  | 100.0 ± 0.0 |            12.05           |            1.61            |
| Tree Cross Attention | 3.5% | 100.0 ± 0.0 |            19.31           |            9.09            |

Reviewer comment (Experiments):
Did you try using Gumbel-softmax for training the tree expansion policy?

Author response:
We have not tried using Gumbel-softmax for training the expansion policy but that would be an interesting direction to explore as an alternative to our REINFORCE-based policy training method.

Reviewer comment (Evaluation):
What are the barriers to applying TCA to self-attention? Would aggregation become the most expensive operation?

Author response:
TCA can be applied to self-attention by setting the query vectors the same as the context vectors. Empirically, we found that aggregation to be a rather efficient operation. For example, on a batch of sequences of length $256$, it took 0.012864 ± 0.002792 seconds.

Reviewer comment (Presentation):
In the last paragraph of 3.1, I was a little confused about why k-d trees were needed as I was only thinking about 1D sequences. Having a picture of an image tree and some more prose about different domains would be really nice for motivating and showing the generality of the method.

Author response:
Thank you for letting us know! We have added in Section 3.1 more details regarding the construction of k-d trees in different domains (sequences and images) to improve the clarity.

Reviewer comment (Writing):
While including cross attention in the model name pragmatically implies the method is not intended directly for self-attention, it would be nice to add a footnote that the focus is not on autoregressive modeling.

Author response:
Thank you for your suggestion. We have updated the paper to include a footnote that clarifies the focus is not on autoregressive modelling.

Reviewer comment (Experiments):
It would be good if a similar row (as given in Table 2) can be added to Table 1 for Perceiver IO with increased latent tokens that matches the performance of TCA on the copy task.

Author response:
We would like to thank the reviewer for their very constructive feedback and for their support.


> It would be good if a similar row (as given in Table 2) can be added to Table 1 for Perceiver IO with increased latent tokens that matches the performance of TCA on the copy task.

In the following table, we include results evaluating Perceiver IO with increased number of latent tokens on varying lengths of the copy task ($N \in \{128, 256, 512, 1024\}$). Specifically, we tested using $100\%$ tokens to evaluate its performance. Although Perceiver IO was able to solve the copy task for $N=128$. Perceiver's performance degraded sharply as the difficulty of the task increased even with $100\%$ tokens. Furthermore, due to Perceiver IO's encoder which aims to map the context tokens to a set of latent tokens, we ran out of memory (on a Nvidia P100 GPU (16 GB)) trying to train the model on a sequence length of 1024.

|                      | $N=128$  | $N=128$     | $N=256$  | $N=256$     | $N=512$  | $N=512$     | $N=1024$ | $N=1024$   |
|----------------------|----------|-------------|----------|-------------|----------|-------------|----------|------------|
| Method               | % Tokens | Accuracy    | % Tokens | Accuracy    | % Tokens | Accuracy    | % Tokens | Accuracy   |
| Cross Attention      | 100.0%   | 100.0 ± 0.0 | 100.0%   | 100.0 ± 0.0 | 100.0%   | 100.0 ± 0.0 | 100.0%   | 99.9 ± 0.2 |
| Random               | ---      | 8.3 ± 0.0   | ---      | 8.3 ± 0.0   | ---      | 8.3 ± 0.0   | ---      | 8.3 ± 0.0  |
| Tree Cross Attention | 10.9%    | 100.0 ± 0.0 | 6.3%     | 100.0 ± 0.0 | 3.5%     | 100.0 ± 0.0 | 2.0%     | 99.6 ± 0.6 |
| Perceiver IO         | 10.9%    | 17.8 ± 0.0  | 6.3%     | 15.2 ± 0.0  | 3.5%     | 13.4 ± 0.2  | 2.0%     | 11.6 ± 0.4 |
| Perceiver IO         | 100.0%   | 100.0 ± 0.0 | 100.0%   | 63.6 ± 45.4 | 100.0%   | 14.1 ± 2.3  | 100.0%   | OOM        |


To analyze the reason behind the degradation in performance of Perceiver IO, we evaluated Perceiver IO on the copy task with a varying number of latent tokens (see table below), i.e., $L \in \{16,32,64,128\}$ and varying sequence lengths $N \in \{128,256,512\}$.  Generally, we found that performance improved as the number of latents increased.  However, the results were relatively unstable. For some runs, Perceiver IO was able to solve the task completely. Other times, Perceiver IO got stuck in lower accuracy local optimas. For longer sequences $N=512$, increasing the number of latents did not make a difference as the model simply got stuck in poor local optimas. 

We hypothesize that the poor performance is due to the incompatibility between the nature of the Copy Task and Perceiver IO's model. The copy task (1) has high intrinsic dimensionality that scales with the number of tokens and (2) does not require computing higher-order information between the tokens. In contrast, Perceiver IO (1) is designed for tasks with lower intrinsic dimensionality by compressing information via its latent tokens and iterative attention encoder, and (2) computes higher-order information via a transformer in the latent space. Naturally, these factors make learning the task difficult for Perceiver IO. Due to space limitations, we have added these tables and analyses to the appendix. 

| Perceiver IO's Num. Latents ($L$) | Copy Task ($N=128$) Accuracy | Copy Task ($N=256$) Accuracy | Copy Task($N=512$) |
|:---------------------------:|:-------------------:|:-------------------:|:------------------:|
|            $L=16$           |    39.83 ± 40.21    |     15.16 ± 0.03    |    13.41 ± 0.15    |
|            $L=32$           |    58.92 ± 47.44    |    34.75 ± 39.14    |    19.41 ± 12.07   |
|            $L=64$           |    79.45 ± 41.11    |    27.10 ± 21.23    |    14.00 ± 0.72    |
|           $L=128$           |    100.00 ± 0.00    |    63.50 ± 39.37    |    13.23 ± 0.42    |

Reviewer comment (Evaluation):
Building on the previous point, wall-clock times for the tree construction and bottom-up aggregation should be reported too.

Author response:
> Building on the previous point, wall-clock times for the tree construction and bottom-up aggregation should be reported too.

For $256$ tokens, the wall-clock time for constructing a tree is $0.076 ± 0.022$ milliseconds. The wall-clock time for the aggregation is $12.864 ± 2.792$ milliseconds. We have added these numbers to the appendix.