Tree Cross Attention

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.

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.

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

More context on the baseline IO Perceiver: The authors need a background section for IO perceiver so the work is self-contained. With the current version, IO perceiver, though being a famous and well-cited paper, is not clearly stated.

Thank you for bringing this to our attention. We have added a background section for Perceiver IO.

More general evaluation: [...] 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 a significant amount of work to be put into actually scaling it to these real settings.

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.

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. [...] it will be important to justify the design with proper inference wall-clock time measurements. [...] Questions: Does the algorithm run faster than full attention in terms of 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 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.

Thank you for your prompt response!

my concerns about the practical applicability of the method remain given its extreme high cost to possibly to replace cross attention, therefore I will maintain my current score.

We would like to reiterate that Tree Cross Attention (TCA) trades off the memory and runtime. Although the runtime increases with TCA, the memory usage (and % tokens) has decreased significantly as shown in Figure 4 and the table results.

The trade-off between memory and runtime can be controlled by the design of the tree. Tree Cross Attention sequentially searches down a tree for the relevant tokens. As such, the factor which primarily affects the overall runtime is the height of tree ($H$) which is dependent on the branching factor $b_f \geq 2$. The relationship between the height of the tree, the branching factor, and the number of tokens is as follows: $H = \lceil\log_{b_f}(N)\rceil$. In our results, we showed the performance with using a binary tree ($b_f=2$). Notably, this corresponds to a tree design with a longer runtime but uses few tokens. By selecting a higher branching factor, the runtime can be decreased at the expense of requiring more tokens as we show in the table below. This feature is not available in standard cross attention.

From these results, it is clear that TCA enables more control over the memory and computation time trade-off. Thus making it more tailored to many practical use cases (especially the low memory settings)

We would like to thank the reviewer for their very supportive feedback. We are pleased to see your enthusiasm for our work.

Would TCA work out of the box for masked language modeling, e.g. BERT?

TCA supports masked inputs and can be applied similarly to self-attention by setting the query vectors in TCA the same as the context vectors. As such, TCA can be applied to masked language modelling.

Did you try using Gumbel-softmax for training the tree expansion policy?

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.

What are the barriers to applying TCA to self-attention? Would aggregation become the most expensive operation?

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.

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.

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.

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 modelling.

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

Theoretical complexity is fine, but the paper should also report wall-clock time for ReTreever and compare it with the full Transformer+Cross Attention and Perceiver IO models.

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. Due to space limitations, we have added these results to the appendix.

I am guessing the tree approach is not parallelizable on accelerated devices like GPUs

The tree approach has a sequential aspect to it related to the height. However, significant 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.

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.

Using ReTreever-full does not make sense and it only confuses the understanding of the reader in my opinion. Either remove it, or add more details like why there is a performance gap between the full cross-attention and ReTreever-full given both are using 100% of the tokens.

Thank you for your suggestion! We have removed ReTreever-Full from the main paper as it is not the focus.

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.

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.