Tree Cross Attention
Authors: Leo Feng
Frederick Tung
Hossein Hajimirsadeghi
Yoshua Bengio
Mohamed Osama Ahmed
Venue: ICLR 2024 poster (2024)
Keywords: Attention
Retrieval
Tree
Primary area: general machine learning (i.e., none of the above)
TL;DR: We propose (1) Tree Cross Attention (TCA), a module based on Cross Attention that only requires a logarithmic $O(\log(N))$ number of tokens, and (2) ReTreever, a flexible token-efficient architecture that outperforms Perceiver IO.

Abstract: Cross Attention is a popular method for retrieving information from a set of context tokens for making predictions. At inference time, for each prediction, Cross Attention scans the full set of $\mathcal{O}(N)$ tokens. In practice, however, often only a small subset of tokens are required for good performance. 
Methods such as Perceiver IO are cheap at inference as they distill the information to a smaller-sized set of latent tokens $L < N$ on which cross attention is then applied, resulting in only $\mathcal{O}(L)$ complexity. 
However, in practice, as the number of input tokens and the amount of information to distill increases, the number of latent tokens needed also increases significantly. 
In this work, we propose Tree Cross Attention (TCA) - a module based on Cross Attention that only retrieves information from a logarithmic $\mathcal{O}(\log(N))$ number of tokens for performing inference. 
TCA organizes the data in a tree structure and performs a tree search at inference time to retrieve the relevant tokens for prediction. 
Leveraging TCA, we introduce ReTreever, a flexible architecture for token-efficient inference. 
We show empirically that Tree Cross Attention (TCA) performs comparable to Cross Attention across various classification and uncertainty regression tasks while being significantly more token-efficient. 
Furthermore, we compare ReTreever against Perceiver IO, showing significant gains while using the same number of tokens for inference.

Published as a conference paper at ICLR 2024
TREE CROSS ATTENTION
Leo Feng
Frederick Tung
Mila Université de Montréal & Borealis Al 1eo. fengemila. quebec
Borealis AI
Frederick. wtw.rrcck..unnbbbbaaiiisa... com
Hossein Hajimirsadeghi
Yoshua Bengio
Borealis AI hossein, ha jimi rsadeghi @boreal isai. com
Mila Université de Montréal
yoshua bengi o@mi la quebec
Mohamed Osama Ahmed
Borealis AI
mohamed. o www...mmemebbbbbalaaaooooom
ABSTRACT
Cross Attention a popular method for retrieving information from set of con- text tokens for making predictions. At inference time, for each prediction, Cross Attention scans the full set of O(N) tokens. In practice, however, often only a small subset of tokens are required for good performance. Methods such Per- ceiver 10 are cheap at inference as they distill the information to a smaller-sized set of latent tokens T. N on which cross attention is then applied, resulting in only (L) complexity. However, in practice, as the number of input tokens and the amount of information distill increases, the number of latent tokens needed also increases significantly. In this work, we propose Tree Cross Atten- tion (TCA) - module based on Cross Attention that only retrieves information from a logarithmic Ollog(N) number of tokens for performing inference. TCA organizes the data tree structure and performs tree search at inference time to retrieve the relevant tokens for prediction. Leveraging TCA, we introduce Re- Treever, a flexible architecture for token-efficient inference. We show empirically that Tree Cross Attention (TCA) performs comparable to Cross Attention across various classification and uncertainty regression tasks while being significantly more token-efficient. Furthermore, compare ReTreeverr against Perceiver IO, showing significant gains while using the same number tokens for inference.
INTRODUCTION
With the rapid growth in applications of machine learning, important objective is make infer- ence efficient both in terms of compute and memory. NVIDIA (Leopold) 2019) and Amazon (Barr) 2019) estimate that 80-90% of the ML workload is from performing inference. Furthermore. with the rapid growth in cow-memory/copput domains (e.g. 1oT devices) and the popularity of attention mechanisms in recent years, there is strong incentive design more efficient attention mechanisms for performing inference.
Cross Attention (CA) is a popular method at inference time for retrieving relevant information from a set of context tokens. CA scales linearly with the number of context tokens ONN). However, in reality many of the context tokens are not needed, making CA unnecessarily expensive in practice. General-purpose architectures such as Perceiver IO (Jaegle et al., (021) perform inference cheaply by first distilling the contextual information down into a smaller fixed-sized set of latent tokens L <<> N. When performing inference, information is instead retrieved from the fixed-size set of latent tokens (L). Methods that achieve efficient inference via distillation are problematic since (1) problems with high intrinsic dimensionality naturally require large number of latents, and (2) the number of latents (capacity of the inference model) is hyperparameter that requires specifying before training. However, in many practical problems, the required model's capacity may not be

---

Published as a conference paper at ICLR 2024
known beforehand. For example, in settings where the amount of data increases overtime (e.g., Bayesian Optimization. Contextual Bandits, Active Learning, etc.... the number of latents needed in the beginning and after many data acquisition steps can be vastly different.
In this work, we propose (1) Tree Cross Attention (TCA), a replacement for Cross Attention that performs retrieval, scaling logarithmically ([log(N) with the number of tokens. TCA organizes the tokens into a tree structure. From this, it then performs retrieval via a tree search, starting from the root. As result, TCA selectively chooses the information retrieve from the tree depending on a query feature vector. TCA leverages Reinforcement Learning (RL) to learn good representations for the internal nodes of the tree. Building on TCA, we also propose (2) ReTreever, a flexible architecture that achieves coken-efficien inference.
In our experiments, we show (1) TCA achieves results competitive to that of Cross Attention while only requiring a logarithmic number of tokens, (2) Re'Treever outperforms Perceiver I0 on var- ious classification and uncertainty estimation tasks while using the same number of tokens, (3) ReTreever's optimization objective can leverage non-differentiable objectives such as classification accuracy, and (4) TCA's memory usage scales logarithmically with the number tokens unlike Cross Attention which scales linearly.
2 BACKGROUND
2.1 ATTENTION
Attention retrieves information from a context set Xg as follows:
QKT
Attention(Q. K, K) softmax(
Vd V
where Q is the query matrix, K is the key matrix, is the value matrix, and d is the dimension of the key and query vectors. In this work, we focus on Cross Attention where the objective is to retrieve information from the set of context tokens X. given query feature vectors, i.e., K and V are embeddings of the context tokens Xe while Q is the embeddings of a batch of query feature vectors. In contrast, in Self Attention, K, V. and Q are embeddings of the same set of token X and the objective is to compute higheroorder information for downstream calculations.
2.1.1 PERCEIVER IO
Perceiver 10 (Jaegle et al. 2021) is general attention-based neural network architecture applica- ble to various tasks. Perceiver 10 is composed of a stacked iterative attention encoder NRD -> RLxD) and a Cross Attention module where N is the number of context tokens and L is hyperpa- rameter. Perceiver IO's encoder aims to compress the information of the context tokens to smaller constant (L) number of latent tokens (Latents) whose initialization is learned. More specifically, each of the encoder's stacked iterative attention blocks work as follows:
Latents <- SelfAttention (CoossAtteniiinnaaaaattaaoa De))
When performing inference, Perceiver I0 simply retrieves information from the set of latents to make predictions via cross attention. As a result, all the necessary information needed for per- forming inference must be compressed into its small fixed number latents. However, this is not practical for problems with high intrinsic dimensionality.
2.2 REINFORCEMENT LEARNING
Markov Decision Processes (MDPs) are defined as a tuple M = AS Po. H), where and denote the state and action spaces respectively. T' a) the transition dy- namics, H(re-1|9t, az, 11+1) the reward function, the discount factor, and the horizon. In the standard Reinforcement Learning setting, the objective is to optimize policy T[als) that maximizes the expected discounted return Fr,,T --1 7t. R(re41 18t; ae, 8141]]-
the initial state distribution, y E (0, 1]
2

---

Published as a conference paper at ICLR 2024
Used at Inference
Encodings
Input
(N D)
Tree Constructor
Array
Encoder
T
Retrieval
Salocacd Encodinge (ogthi Di
KV
Query
Cross Altention
Array
Output
(M D)
Tree Cross Attention
Figure 1: Architecture Diagram of ReTreeveer. Input Array comprises set of N context tokens which are fed through an encoder to compute a set of context encodings. Query Array denotes a batch of M query feature vectors. Tree Cross Attention organizes the encodings and constructs tree T. At inference time, given a query feature vector, logarithmic-sized subset of nodes (encodings) retrieved from the T. The query feature vector retrieves information from the subset of encodings via Cross Attention and makes a prediction.
Ag-Agg/h,
h - Agg(h,hh
Rg- Ag-AAg(h,,t,
R At A.)
fü Ag h,)
h - Agg((t, An))
hg - A,,
Rig
Hts
Ato
Ru
hia
R1
Figure 2: Diagram the aggregation procedure performed during the Tree Construction phase. The aggregation procedure is performed bottom-up beginning from the parents of the leaves and ending at the root of the tree. The complexity of this procedure is (NN) but this only needs to be performed once for set of context tokens. Compared to the cost of performing multiple predictions, the one- time cost of the aggregation process minor.
3 TREE CROSS ATTENTION
We propose Tree Cross Attention (TCA), a token-efficient variant of Cross Attention. Tree Cross Attention is composed of three phases: (1) Tree Construction, (2) Retrieval, and (3) Cross Attention. In the Tree Construction phase (R - T), TCA organizes the context tokens (R-V x into a tree structure (T) such that the context tokens the leaves in the tree. The internal (i.e., non-leaf) nodes the tree summarise the information its subtree. Notably, this phase only needs to be performed once for set of context tokens.
The Retrieval and Cross Attention performed multiple times at inference time. During the Re- trieval phase (T R - (N))))()))) the model retrieves a logarithmic-sized selected subset of nodes (3) from the tree using query feature vector R D (or batch of query feature vectors RMDD). Afterwards, Cross Attention (I (logg())))
R RD) is performed by retrieving the information from the subset of nodes with the query feature vector. The overall complexity of infer- ence is Ollgg(N) (N)) per query feature vector. We detail these phases fully in the below subsections.
After describing the phases of TCA, we introduce ReTreever, a general architecture for token- efficient inference. Figure T illustrates the architecture of Tree Cross Attention and Re Reeeee...
3

---

Published as a conference paper at ICLR 2024
ho
f1
h2
lt3
h4
hts
h6
h
igg
hg
A10
11
h12
13
f1
Figure 3: Example result of Retrieval phase. The policy creates path from the tree's root to its leaves, selecting a subset of nodes: the terminal leaves in the path and the highest-level unexplored ancestors of the other leaves. The green arrows represent the path (actions) chosen by the policy m. The red arrows represent the actions rejected by the policy. The green nodes denote the subset of nodes selected, i.e., S { ha, ha, }.0). The grey nodes denote nodes that were explored some point but not selected. The red nodes denote the nodes that were not explored or selected.
3.1 TREE CONSTRUCTION
The tokens are organized tree T such that the leaves of the tree consist of all the tokens. The internal nodes (i.e., non-leaf nodes) summarise the information of the nodes in their subtree. The information stored in a node has two specific desideratas, summarising the information in its subtree needed for performing (1) predictions and (2) retrieval (i.e., finding the specific nodes in its subtree relevant for the query feature vector).
The method of organizing the data in the tree is flexible and can be done either with prior knowledge of the structure of the data, with simple heuristics, or potentially learned. For example, heuristics used to organize the data in traditional tree algorithms be used to organize the data, e.g., ball- tree (Omohundro) 1989) or k-d tree (Cormen et al., 2006).
After organizing the data in the tree, aggregator function Agg P(R)) -> Rd (where P denotes the power set) is used to aggregate the information of the tokens. Starting from the parent nodes of the leaves of the tree, the following aggregation procedure is performed bottom-up until the root of the tree: It, = Agg (h..l u E C. where C. denotes the set of children nodes of a node v and h, denotes the vector representing the node. Figure Saaureeiiluse illluttrae the aggregation process.
In our experiments, we consider a balanced binary tree. To organize the data, we use a k-d tree. During their construction, k-d trees select an axis and split the data evenly according to the median, grouping similar data together. For sequence data such as time series, the data is split according to the x-axis (e.g., time). For images, this means that the pixels are split according the or y-axis, i.., vertically or horizontally. To ensure the tree is perfectly balanced, padding tokens are added such that the values are masked during computation. Notably, these padding nodes are less than the number of tokens. As such, it does not affect the big-O complexity.
3.2 RETRIEVAL
To learn informative internal node embeddings for retrieval, we leverage policy learned via Reinforcement Learning for retrieving a subset of nodes from the tree. Algorithm [ describes the process of retrieving the set of selected nodes S from the tree T. Figure 3 illlustrates an example of the result of this phase. Figures to II] in the Appendix illustrate a step-by-step example of the retrieval described in Algorithm
v denotes the node of the tree currently being explored. At the start of the search, it is initialized as < Root(T). The policy's 's state is the set of children nodes of the current node being explored C, and the query feature vector q, i.e., s = (C,, q). Starting from the root, the policy samples one of the children nodes, i.e., o' To(a q) where a C., to further explore for more detailed
4

---

Published as a conference paper at ICLR 2024
information retrieval. The nodes that are not being further explored added to the set of selected nodes 8 + & J (Co (o)) Afterwards, is updated (v t- v') with the new node and the process is repeated until is a leaf node (i.e., the search is complete).
Since the height of a balanced tree (log(N)), such, the number of nodes that are explored and retrieved is logarithmic: ||||| = (log(N)).
Algorithm Retrieval
Input:
Tree Structure (T), policy (To), and query feature vector
Output: Set of selected nodes (S)
3.3 CROSS ATTENTION
V - Ro)t(T)
S t 0
The retrieved set of nodes (S) is used input for Cross Attention alongside the query feature vec- tor q. Since |||| (log([N)), this results in Cross Attention with overall logarithmic complexity (log((N) complexity per query feature vector. In contrast, applying Cross Attention to the full set of tokens has linear complexity O(N). Notably, the set of nodes S has full receptive field of the entire set of tokens, i.e., either a token is part of the se- lected set of nodes or is descendant (in the tree) of one of the selected nodes.
while v is not leaf do v' ( CalC...))
where C,, D Cp denotes
S (C, (e))
children nodes of V x v' end while S < SUr
3.4 RETREEVER: RETRIEVAL VIA TREE CROSS ATTENTION
In this section, we propose ReTreever (Figure D. a general-purpose model that achieves token- efficient inference leveraging Tree Cross Attention. The architecture is similar in style Per- ceiver IO's 's (Figure 5 in the Appendix).
In the case of Perceiver IO, the model composed of (l) an iterative attention encoder NRNxD LLxD) which compresses the information into a smaller fixed-sized L set of latent tokens and (2) Cross Attention module used during inference for performing information retrieval from the set of latent tokens. As result, Perceiver IO's inference complexity is (LL).
In contrast, ReTreever composed of (1) an encoder (IR NxD Nx)) and (2) Tree Cross Atten- tion (TCA) module used during inference for performing information retrieval from tree structure. Unlike Perceiver IO which compresses information via specialized encoder to achieve efficient in- ference, ReTreever's 's inference is token-efficient irrespective of the encoder, scaling logarithmically (log(N) with the number of tokens. As such, the choice of encoder is flexible and can be, for example, Transformer Encoder or efficient versions such as Linformer, ChordMixer, etc.
3.5 TRAINING OBJECTIVE
The objective ReTreever optimises consists of three components with hyperparameters ARL and ACA. denoting the weight of the terms:
LRefreever LICA + ARLCAL + ACACCa
LTcA aims tackle the first desiderata, i.e., learning node representations in the tree that summarise the relevant information in its subtree for making good predictions.
STCA = Loss -osssoossoooo (CrosssAttenniionnaa
LR. tackles the second desiderata, i.e., learning internal node representations for retrieval (the RL policy т) within the tree structure
log(N (N)
LRI.
R log no (arls)) + a.H(To(-))s))
'The described objective is the standard REINFORCE loss (without discounting 7 1) using entropy bonus for sparse reward environment.
5

---

Published as a conference paper at ICLR 2024
The horizon is the height of the tree log( H denotes entropy, and R denotes the reward/objective we want ReTreever to maximise Typically this reward corresponds to the negative TCA loss R = -CCCA, e.g.. Negative MSE, Log-Likelihood, and Negative Cross Entropy for regression, uncertainty estimation, and classification tasks, respectively. However, crucially, the reward does not need to be differentiable. As such, R can also be an objective we are typically not able to optimize directly via gradient descent, c.g., accuracy for classification tasks (R = for correct classification and R for incorrect classification).
Furthermore, at inference time, the objective of the policy and the Cross Attention module are in- tertwined in that the policy aims to select relevant nodes for the Cross Attention. As such, to improve training. the weights are shared between the policy and the Cross Attention module. More specifi- cally, the policy is parameterized as an attention module where the policy's ' action probabilities are the attention weights of the context (i.e., child nodes C.)) given the query (i.e., q).
Lastly, to improve the early stages of training and encourage TCA to learn good node representa- tions, LCA is included:
LGs = LosssCrossAtttnttion/ Leavess))))) v)
EXPERIMENTS
Our objective in the experiments is: (1) Compare Tree Cross Attention (TCA) with Cross Attention in terms of the number of tokens used and the performance. We answer this in two ways: (i) by comparing Tree Cross Attention and Cross Attention directly on a memory retrieval task and (ii) by comparing our proposed model ReTrreever which uses TCA with models which use Cross Attention. (2) Compare the performances of ReTreever and Perceiver while using the same number of tokens. In terms of analyses, we aim to: (3) Highlight how ReTreever and TCA can optimize for non-differentiable objectives using CRI, improving performance. (4) Show the importance of the different loss terms in the training objective. (5) Empirically show the rate at which Cross Attention and TCA memory grow with respect to the number of tokens.
To tackle those objectives, we consider benchmarks for classification (Copy Task, Human Activity) and uncertainty estimation settings (GP Regression and Image Completion). Our focus is on (I) comparing Tree Cross Attention with Cross Attention, and (2) comparing ReTrreever with general- purpose models (i) Perceiver I0 for the same number of latents and (ii) Transformer (Encoder) + Cross Attention. Full details regarding the hyperparameters are included in the appendix
For completeness, in the appendix, we include the results for many other baselines (including re- cent state-of-the-art methods) for the problem settings considered in this paper. The results show the baseline Transformer + Cross Attention achieves results competitive with prior state-of-the-aa Specifically, we compared against the following baselines for GP Regression and Image Comple- tion: LBANPs (Feng et al., 2023a). TNPs (Nguyen & Grover 2022). NPs (Garnelo et al. 201 8b).
BNPs (Lee et al.] CNPs (Garnelo et al., 2018)), CANPs (imm et al. 2019), ANPs (Kim et al.,
2019). . and BANPs (Lee et al., (020). We compared against the following baselines for Human Ac-
tivity: SeFT (Horn et al., 2020), RNN-Decay (Che et al. 2018). IP-Nets (Shukla & Marlin 2019), (Chen et al. 2018), L-ODE-ODE (Rubanova et al.] (019),, and TTAD--Enc (Shukla & Marlin| 2021). Additionally, we include results in the Appendix for Perceiver 10 with a vari- ous number of latent tokens, showing Perceiver I0 requires several times more tokens to achieve comparable performance to Re'rreeve
L-ODE-RNN (Chen
4.1
COPY TASK
We first verify the ability of Cross Attention and Tree Cross Attention (TCA) to perform retrieval.
The models are provided with sequence of length N = 2*, beginning with a [BOS] (Beginning of Sequence) token, followed by a randomly generated palindrome comprising of 2 with [EOS] (End of Sequence) token. The objective of the task is to predict the second half (2t-1) of the sequence given the first half 22k-1 tokens) of the sequence as context. To make a correct
digits, ending
Thhe only reward R. is given at the end of the episode. As such, the undiscounted return for each timestep is R.
'The code is available at https//gthaaaccoooaaaaaaaalleeet

---

Published as a conference paper at ICLR 2024
N
256 Accuracy
N
512 Accuracy
N
1024 Accuracy
Method
% Tokens 100.0%
% Tokens
% Tokens 100.0%
Cross Attention
100.0 + 0.0
100.0%
100.0 + 0.0
99.9 + 0.2 8.3
Random
8.3 13::000 0.0
8.3
Perceiver IO
6.3%
15.2 0.0
3.5% 3.5%
0.2 100.0
2.0% 2.0%
11.6 0.4
TCA
6.3%
100.0 0.0
99.6 0.6
Table I: Copy Task Results with accuracy (higher is better) and % tokens (lower better) metrics.
Method
% Tokens
RBF
Matern 5/2
Transformer + Cross Attention
100%
1.35 + 0.02
0.91 + 0.02
Perceiver IO
14.9%
1.06 + 0.05
0.58 + 0.05
Perceiver IO (L 32)
68.0%
1.25 + 0.04
0.78 + 0.05
ReTrreever
14.9%
1.25 + 0.02
0.81 + 0.02
Table 2: 1-D Meta-Regression Experiments with log- ligelikliihodd (higher is better) and % tokens (lower is better) metrics.
prediction for index 2* (where 0 < i 2*-1),, the model must retrieve information from its input/context sequence at index + 1. The model is evaluated on its accuracy on 3200 randomly generated test sequences. As a point of comparison, we include Random as a baseline, which makes random predictions sampled from the set of digits (0 9), [EOS], and [BOS] tokens.
Results. We found that both Cross Attention and Tree Cross Attention were able solve this task perfectly (Table D. In comparison, 'CA requires ~ 50x fewer tokens than Cross Attention. Furthermore, we found that Perceiver IO's performance for the same number of tokens was dismal (15.2 + 0.0% accuracy at N 256), further dropping in performance increased. This result is expected as Perceiver 10 aims to compress all the relevant information for any predictions into small fixed-sized set of latents. However, all the tokens are relevant depending on the queried index. As such, methods which distill information to a lower dimensional space are insufficient. In contrast, TCA performs tree-based memory retrieval, selectively retrieving small subset of tokens for predictions. As such, TCA is able to solve the task perfectly while using the same number of tokens Perceiver.
the length of the sequence
4.2
UNCERTAINTY ESTIMATION: GP REGRESSION AND IMAGE COMPLETION
We evaluate ReTreever popular uncertainty estimation settings used in (Conditional) Neural Pro- cesses literature and which have been benchmarked extensively (Table 13 in Appendix) (Garnelo)
et al. 2018a
Kim et al. 2019 Lee et al. 2020 Nguyen & Grover 2022 Feng et al. 2023a b).
4.2.1
GP REGRESSION
The goal of the GP Regression task is to model an unknown function given N points. During training. the functions are sampled from GP prior with RBF kernel fi GP k) where m(x) 0 and k(x,
2 ) The hyperparameters of the kernel are randomly sampled according to | ~ [40066 1.0), of 1(0.1, 1.0), N ~ U3, 47). and M U(3, 50 N). After training, the models are evaluated according to the log-likelihood of functions sampled from GPs with RBF and Matern 5/2 kernels.
of exp(
we see that ReTreever (1.25 + 0.02) outperforms Perceiver IO (1.06 0.05) by large margin while using the same number of tokens for inference. To how many latents Perceiver 10 would need to achieve performance comparable to ReTreever, we varied the number of latents (see Appendix Table [3 for full results). We found that Perceiver 10 needed 4.6x the number of tokens to achieve comparable performance.
Results. In Table 2
4.2.2 IMAGE COMPLETION
The goal the Image Completion task to make predictions for the pixels an image given a random subset of pixels of an image. The CelebA dataset comprises coloured images of celebrity

---

Published as a conference paper at ICLR 2024
CelebA
EMNIST
Method
% Tokens
LogLikelihood
% Tokens
LogLikelihood
Transformer Cross Attention
100%
3.88 + 0.01
100%
1.41 = 0.00
Perceiver IO ReTrreever
4.6%
3.20 + 0.01
4.6%
1.25 + 0.01
4.6%
3.52 + 0.01
4.6%
1.30 + 0.01
Table Image Completion Experiments. The methods are evaluated according to the log- ligliiklihood (higher is better) and % tokens (lower is better) metrics.
Model
Tokens
Accuracy
Transformer Cross Attention
100%
89.1 + 1.3
Perceiver IO
14%
87.6 + 0.3
ReTreever
14%
88.9 + 0.4
Table 4: Human Activity Experiments with accuracy (higher is better) and % tokens (lower is better) metries.
faces. The images are downsized to size 32 x 32. The EMNIST dataset comprises black and white images of handwritten letters with a resolution of 32 x 32. We consider 10 classes. Following prior works, the values of the images are rescaled to be between [-1, 1J, the values to be between [-0.5, 0.5], and the number of pixels is sampled according to N U[3 197) and M U[3 200 N).
Results. In Table 3 we see that ReTrreever outperforms Perceiver IO significantly both CelebA and EMNIST while using the same number of tokens. Compared with Transformer + Cross Atten- tion, Re' Trevver uses ~ 21 x fewer tokens, making it significantly more token-efficient.
4.3 TIME SERIES: HUMAN ACTIVITY
The human activity dataset consists of 3D positions of the waist, chest, and ankles (12 features) collected from five individuals performing various activities such as walking, lying, and standing. We follow prior works preprocessing steps, constructing a dataset of 6, 554 sequences with 50 time points. The objective of this task is classify each time point into one of eleven types of activities.
Results. Table 4 show similar conclusions as our prior experiments. (1) ReTreever performs com- parable to Transformer Cross Attention, requiring far fewer tokens (86% less), and (2) ReTreever outperforms Perceiver IO significantly while using the same number of tokens.
4.4 ANALYSES
Optimising non-differentiable objectives using LRI. In classification, typically, accuracy the metric we truly care about. However, accuracy is a non-differentiable objective, so instead in deep learning, we often use cross entropy so we can optimize the model's 's parameters via gradient descent. However, in the case of RL, the reward does not need to be differentiable. As such, we compare the performance of ReTreever when optimizing for the accuracy and the negative cross entropy. Table shows that accuracy the RL reward improves upon using negative cross entropy. This is expected since (1) the metric we evaluate our model on is accuracy and not cross entropy, and (2) accuracy as entropy
reward simpler compared negative optimize. More specifically, a reward using accuracy as the metric is either: R = 0 for an incorrect prediction or for a correct prediction. However, reward based on cross entropy can have vastly different values for incorrect predictions.
reward, making it easier for policy to
N 256 Accuracy
N
512
N 1024
Method
% Tokens
% Tokens
Accuracy
Tokens
Accuracy
TCA (Acc)
6.3%
100.0 = 0.0
3.5% 3.5%
100.0 4: 0.0
2.0% 2.0%
99.6 + 0.6
TCA (Neg. CE)
6.3%
99.8 - 0.3
95.5 - 3.8
80.3 = 14.8
Table 5: Comparison of Accuracy and Negative Cross Entropy the reward.
8

---

Published as a conference paper at ICLR 2024
Effect of Tokens en Memory Usage
Training Curve (Vary Aca)
Training Curve (Vary And
100
100
CrossAttention
Memory Usage (GBs)
TCA (Ours)
90-
90-
Accuracy
Accuracy
-4
33
Aca = 0.0 Aua 1.0
- 1.0 - 0.0
2
Aci 1.0 Arc 1.0
Aoe 1oa-1.0A4 -1.0
Aca m 10.0 Ar, 1.0
Ace m 1.0 = 10.0
50
500
1000
1500
2000
10 20
30
50
10
20
30
40
Number of Tokenss
Number of Steps (x 1e3)
Number of Steps (x 1e3)
Figure 4: Analyses Plots. (left) Memory usage plot comparing the rate in which memory usage grows inference time relative to the number of tokens. (middle) Training curve with varying weights (Aca) for the Cross Attention loss term. (right) Training curve with varying weights (ARz) for the RL retrieval loss term.
Memory Usage. Tree Cross Attention's memory usage (Figure 4 (left)) is significantly more effi- cient than Cross Attention, growing logarithmically in the number of tokens compared to linearly.
Importance of the different loss terms Aaa and XRI. Figure 4 (right) shows that the weight of the RL loss term is crucial to achieving good performance. ARL = 1.0 performed the best. If the weight of the term is too small ARL {0.0, 0.1} or too large ARL 10.0, the model is unable to solve the task. Figure (middle) shows that the weight of the Cross Attention loss term improves the stability of the training, particularly in the early stages, Setting Aca = 1.0 performed the best. Setting too small values for Aca E {0.0, 0.1} causes some erraticness in the early stages of training. Setting too large of a value ACA 10.0 slows down training.
RELATED WORK
There have been few works which have proposed to leverage tree-based architecture (Tai al., 2015 Nguyen et al. 2020 Madaan et al.] 2023 Wang et al.] 2019) for attention, the closest of which our work is Treeformer (Madaan et al. 2023). Unlike prior works (including Treeformer) which focused on replacing self-attention in transformers with tree-based attention, in our work, we focus on replacing cross attention with tree-based cross attention mechanism for efficient memory retrieval for inferences. Furthermore, Treeformer (1) uses a decision tree to find the leaf that closest resembles the query feature vector, (2) TF-A has a partial receptive field comprising only the tokens in the selected leaf, and (3) TF-A has worst-case linear complexity in the number of tokens per query. In contrast, TCA (1) learns a policy to select a subset of nodes, (2) retrieves a subset of nodes with full receptive field, and (3) has a guaranteed logarithmic complexity.
As trees are a form of graphs, Tree Cross Attention bears a resemblance to Graph Neural Networks (GNNs). The objectives, however, of GNNs and TCA are different. In GNNs, the objective is to perform edge, node, or graph predictions. However, the goal of TCA is to search the tree for a subset of nodes that is relevant for a query. Furthermore, unlike GNNs which typically consider the tokens to correspond one-to-one with nodes of the graph, TCA only considers the leaves of the tree to be tokens. We refer the reader surveys on GNNs (Wu al. 2020
Thomas al. 2023).
CONCLUSION
In this work, we proposed Tree Cross Attention (TCA), variant of Cross Attention that only re- quires a logarithmic Olog(N)) number of tokens when performing inferences. By leveraging RL, TCA can optimize non-differentiable objectives such as accuracy. Building on TCA, we introduced ReTreever, a flexible architecture for tion and uncertainty prediction tasks, showing (1) TCA achieves performance comparable to Cross Attention while being significantly more token efficient and (2) ReTreever outperforms Perceiver I0 while using the same number tokens for inference.
inference. We evaluate across various classifica-

---

Published as a conference paper at ICLR 2024
REFERENCES
Jeff Barr. Amazon ec2 update- inf'l instances with aws inferentia chips for high performance cost-effecctive inferencing, Dec 2019. URL httpps///ww...aaaoo.oo//ooo/ogg/
Zhengping Che, Sanjay Purushotham, Kyunghyun Cho, David Sontag, and Yan Liu. Recurrent neural networks for multivariate time series with missing values. Scientific reports, 8(1):6085, 2018.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein. Introduction to algorithmss Third Edition. 2006.
Leo Feng, Hossein Hajimirsadeghi, Yoshua Bengio, and Mohamed Osama Ahmed. Latent bot- tlenecked attentive neural processes. In International Conference on Learning Representations, 2023a. URL
Leo Feng, Frederick Tung, Hossein Hajimirsadeghi, Yoshua Bengio, and Mohamed Osama Ahmed. Constant memory attentive neural processes. arXiv preprint arXiv:2305. 14567, 2023b.
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Ali Eslami. Conditional neural processes. In International Conference Machine Learning, pp. 1704-175 1713. PMLR, 2018a.
Marta Gareel,, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXiv:1807.0162. 2018b.
Max Horn, Michael Moor, Christian Bock, Bastian Rieck, and Karsten Borgwardt. Set functions for time series. In International Conference on Machine Learning, pp. 4353 PMLR, 2020.
Andrew Jaegle, Sebastian Borgeaud, Lean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al. Perceiver io: A general architecture for structured inputs & outputs. In International Conference on Learning Representations, 2021.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee whye Teh. Attentive neural processes. 2019.
Juho Lee, Yoonho Lee, Jungtaek Kim, Eunho Yang, Sung Ju Hwang, and Yee Whye Teh. Boot- strapping neural processes. Advances in neural information processing systems, 33:6606-6615. 2020.
George Leopold. Aws to offer nvidia's 14 gpus for ai inferencing. Mar 2019.
URL
https: http:///www..a.a...../////////// ai www.innneeeeeeeeeeeee pla.eeemm/
Lovish Madaan, Srinadh Bhojanapalli, Himanshu Jain, and Prateck Jain. Treeformer: Dense gradi- ent trees for efficient attention computation. In International Conference on Learning Represen- tations, 2023.
Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. In International Conference on Machine Learning, pp. 16569-16594, PMLR, 2022.
Xuan-Phi Nguyen, Shafiq Joty, Steven Hoi, and Richard Socher. Tree-structured attention with hierarchical accumulation. In International Conference on Learning Representations, 2020. URI. Mtttps:/openzevvew.
hwwp:/nnee/ffouummreeee forum? - HJxK5pEYVI
Stephen M Omohundro. Five balitree construction algorithms. Institute Berkeley, 1989.
International Computer Science
10

---

Published as a conference paper at ICLR 2024
Yulia Rubanova, Ricky TQ Chen, and David K. Duvenaud. Latent ordinary differential equations for saaeuee time series. Advances in neural information processing systems, 32, 2019.
Satya Narayan Shukla and Benjamin Marlin. ceapollation-redicttin networks for irregularly sam- pled time series. In International Conference on Learning Representations, 2019.
Satya Narayan Shukla and Benjamin Marlin, Multi-time attention networks for irregularly sampled time series. In International Conference on Learning Representations, 2021.
Kai Sheng Tai, Richard Socher, and Christopher D Manning. Improved semantic representations from tree-structured long short-term memory networks. arXiv preprint arXiv:1503.00077 2015.
Josephine Thomas, Alice Moallemy-Ouree Silvia Beddar-Wiesing, and Clara Holzhüter. Graph neural networks designed for different graph types: A survey. Transactions Machine Learn- ing Research, 2023. ISSN 2835-8856. URL https: attpss//ppenieeiisnnttt.oo for.mnee//ooummiio
h4BYt 14BYY279uy?
Yau-Shian Wang, Hung-Y Yi Lee, and Yun-Nung Chen, Tree transformer: Integrating tree structures into self-attention. arXiv preprint arXiv: 1etp:9999999999 2019.
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. A comprehensive survey on graph neural networks. JEEE transactions on neural networks and learning systems, 32(1):4-22) 24, 2020.

---

Published as a conference paper at ICLR 2024
A APPENDIX: ILLUSTRATIONS
A.1 BASELINE ARCHITECTURES
Figures 5 and illustrate the difference in architectures between the main baselines considered in the paper. Transformer + Cross Attention is composed of Transformer encoder (I D + R and Cross Attention module (R (RMxx xIR NRRx -> I M) Perceiver IO is composed of an iterative
attention encoder (R Nx - R and Cross Attention module (IR MR&Mx x x R - RMxD).
In contrast, ReTreever (Figure is composed of flexible encoder (R NR Cross Attention module.
R R and a Tree
Empirically, we showed that Transformer Cross Attention achieves good performance. However, its Cross Attention is inefficient due to retrieving from the full set of tokens. Notably, Perceiver IO achieves its inference-ting efficiency by performing a compression via iterative attention encoder. To achieve token-efficient inferences, the number of latents needs to be significantly less than the number of tokens originally, i.e., << N. However, this not practical in hard problems since setting low values for L results in significant information loss due to the latents being bottleneck. In contrast, ReTreever able to perform token-efficient inference while achieving better performance than Perceiver 1O for the same number of tokens. ReTreever does this by using Tree Cross Attention to retrieve the necessary tokens, only needing a logarithmic number of tokens log(NV << N, making it efficient regardless of the encoder used. Crucially, this also means that ReTreever is more flexible with its encoder than Perceiver IO, allowing for customizing the kind of encoder depending on the setting.
Used Inference
Input Array (N x D)
Encoding (L D)
Iterative Attention Encoder
KV
Latent Array (L x D)
Output Array (M D)
Cross Attention
Output
Figure 5: Architecture Diagram of Perceiver IO. The model composed of an iterative attention encoder (TR NxD - RRLD) and Cross Attention module MRD RRLDD MxD)
Used at Inference
Encoding (N x D)
Input Array (N D)
Transformer Encoder
KV
Query Array (M D)
Cross Attention
Output
Figure 6: Architecture Diagram of Transformer + Cross Attention. The model is composed of Transformer encoder NxD Nx and Cross Attention module MD x + M M)
A.2 EXAMPLE RETRIEVAL
In Figures 7 to 11
illustrate example of the retrieval process described in Algorithm
12

---

Published as a conference paper at ICLR 2024
h2
h1
h2
h3
h4
ht5
h6
hr
hg
hg
1110
h1
h12
h13
h1
Figure 7: Example Retrieval Step 0 Illustration. The diamond icon on node indicates the node that will be explored. The retrieval process begins at the root, i.e., v f- 0 where the index of root node. Yellow nodes denote the nodes that have yet to be explored and may be explored in the future. Grey nodes denote the nodes that are or have been explored but have not been added to the selected set S. Red nodes denote the nodes that have not been explored and will not be explored in the future. Green nodes denote the subset of nodes selected, i.e., S.
ho
Pby
ht
hg
h4
ht5
h6
h
igg
hg
h10
h
h
h
fi
Figure 8: Example Retrieval Step Illustration. Green arrows represent the path (actions) chosen by the policy æ for a query feature vector. Red arrows represent the actions rejected by the policy. The right child of e = 0 was rejected, it is added to the selected set S. As consequence. descendant nodes of v's right child will never be explored, so they are coloured red in the diagram. Tentatively, S {h2). Because the policy selected the left child of 0, thus have 1
13

---

Published as a conference paper at ICLR 2024
hoo
h1
h2
h2y
h55
h6
hg
hg
h10
h1
h12
1
11
Figure 9: Example Retrieval Step Illustration. The left child
is rejected this time, so Consequently, the descendant nodes
added to the selected set S (green). Tentatively, $ = (h2, ha). of e's left child will never be explored, so they are coloured red. Because the policy selected the right child of e = 1, we thus have <- 4.
ho
h1
lt2
h3
h4
h2s
h6
h
hg
fly
N10
h1
h12
h13
ft
Figure 10: Example Retrieval Step Illustration. The right child of is rejected this time, so is added to the selected set S. Tentatively, = {h2, ha, J11). The descendant nodes of v's right child are be coloured red, but there are none. Because the policy selected the left child of 4, we thus have e t-
ho
ft
h2
hg
h4
h23
h6
h9
hg
hg
h1o
ft
h12
h
ft
Example Retrieval Step 4 Illustration. its end. At the end of the tree search, we simply add the node to the selected set, resulting in S {/2 ha, hg, h1o)
Figure 11:
9 is leaf, the tree search has reached
14

---

Published as a conference paper at ICLR 2024
Method
Tokens
CelebA
EMNIST
Transformer Cross Attention
100%
3.88 0.01
1.41 1.20
Perceiver I0
4.6%
3.20
P[Treever-Randoon
4.6%
3.02
1.26
ReTreever-KKD KD-2o
4.6%
4.50
1.30
ReTreever-KK KD-21
4.6%
3.52 0.01
1.26 + 0.02
ReThrevver-KD-RP
4.6%
3.51 + 0.01
1.29 + 0.02
Table 6: Tree Design Analyses Experiments. We compare the performance of ReTreever with var- ious heuristics for structuring the tree (1) organized randomly TeTeeerr-Randonn (2) sorting via the first index of the image data the image data Sitreeverr-D-zz and (4) sorting the value after a Random Projection (ReTreever- KD-RP).
and (3) sorting via the second index 21 of
GP Regression RBF
Image Completion
Method
% Tokens
Matern 5/2
% Tokens
CelebA
ReTreever
14.9%
1.25 + 0.02
0.81 + 0.02
4.6% 100%
3.52 + 0.01
Re RuTevveer-Fll 100%
1.33 0.02
0.90 + 0.01
3.77 0.01
Table 7: Comparison of ReTreever with ReTreever-Fur which naively seleets all the leaves of the tree instead of leveraging the learned policy to select subset of the nodes.
B APPENDIX: ADDITIONAL EXPERIMENTS AND ANALYSES
B.I ADDITIONAL ANALYSES
B.1.1 IMPORTANCE OF TREE DESIGN
Unlike sequential data where sorting according to its index is natural way organize the data, it is less clear how to organize other data modalities (e.g., pixels) as leaves in balanced tree. As such, in this experiment (Table compare the performance of ReTreever with various heuristics for structuring the tree (1) organized randomly ReeTreever-Ranoml (2) sorting via the first index To of the image data SerFreever-KD-ryy and (3) sorting via the second index I1 of the image data (ReTreever-KKD -DDev) 11,, and (4) sorting the value after a Random Projection [Rerreever-Kl -DD-PP ReTreever- KD-RP generates random matrix w E IR [Rdim(x)x1 x1 that is fixed during training and orders the context tokens in the tree according to w" instead.
In Table we see that ReTreever outperforms Perceiver IO significantly when splitting according to either of the axes. In contrast, organizing the tree randomly causes ReTreever to perform worse than Perceiver I0. There is, however. minor difference between KD-zo and KD-21- We hypothesize that this in part due image data being multi-diimensoonll and not all dimensions are equal.
B.1.2 RETREEVER-FULL
ReTreever leverages RL to learn good node representations for informative retrieval. However, ReTreever limited to only using the learned policy for retrieval inference time. For example, instead of retrieving subset of the nodes according to policy, we can (as proof of concept) choose to retrieve all the leaves of the tree "ReTreever-Full Full" (see Table 7).
Results. On CelebA and GP Regression, we see that ReTrrever-Full achieves performance close state-of-the-art. For example, on CelebA, ReTreever-Full achieves 3.77 0.01, and Transformer Cross Attention achieves 3.88 + 0.01. On GP, ReTreever-Full achieves 1.33 + 0.02, and Trans- former + Cross Attention achieves 1.35 + 0.02. In contrast, ReTreever achieves significantly lower results. The performance gap between Re Reereev and ReTreever-Full Full suggests that the performance of ReTreever can vary depending on the number of nodes (tokens) selected from the tree. As such, an interesting future direction is the design of alternative methods for selecting subsets of nodes.
'Random Projection is popular in dimensionality reduction literature as it preserves partial information regarding the distances between the points.
15

---

Published as a conference paper at ICLR 2024
N 256
N
512
N
1024 Accuracy
Method
% Tokens
Accuracy
% Tokens
Accuracy
% Tokens 100.0%
Cross Attention
100.0%
100.0 0.0
100.0%
100.0 10 0.0
99.9 + 0.2 8.0 11.6
Random
3.3::.00
8.0
Perceiver IO
6.3%
1.0
3.5%
13.4
2.0%
Perceiver IO
100.0%
63.6 45.4
100.0%
14.1 2.3
100.0%
OOM 99.6 + 0.6
TCA
6.3%
100.0
3.5%
100.0 - 0.0
2.0%
Table 8: Copy Task Results with accuracy (higher is better) and % tokens (lower is better) metrics.
Perceiver IO
Copy Task
Num. Latents (L)
== 128
NV 256
N 512
L 16
39.83 = 40.21
15.16 0.03
13.41 #.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
63.50 = 39.37
13.23 0.42
Table 9: Analysis of Perceiver IO's performance relative to the number of latents (L) and the se- quence length (N) on the Copy Task
B.1.3 PERCHIVER IO's COPY TASK PERFORMANCE
In Table 8 we include results evaluating Perceiver 10 with increased number of latent tokens on varying lengths of the copy task. Specifically, we tested using 100% tokens to evaluate its perfor- mance. Perceiver IO's performance degraded sharply the difficulty of the task increased. Further- more, due Perceiver IO's encoder which aims to map the context tokens to a set of latent tokens, ran out of memory (on a Nvidia P100 GPU (16 GB)) trying to train the model sequence length of 1024.
To analyze the reason behind the degradation in performance of Perceiver IO, we evaluated Perceiver IO (Table 9 with a varying number of latent tokens L E {16, 32, 64, 128} and varying sequence found that performance improved
the number of lengths N (128, 256, 512). Generally, relatively unstable. For some runs, Perceiver IO was latents increased. However, the results were able solve the task completely. Other times, Perceiver 10 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 1O'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) 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.
B.1.4 RUNTIME ANALYSES
Runtime of Tree Construction and Aggregation. For 256 tokens, the wall-clock time for con- structing a tree with our implementation was 0.076 + 0.022 milliseconds. The wall-clock time for the aggregation is 12.864 + 2.792 milliseconds.
Runtime of Tree Cross Attention-based models vs. vanilla Cross Attention-based models. In table 10 we measured the wall-clock time of the modules used during inference time, i.c., Re- Treever's Tree Cross Attention module (R Mg xD x + MRMx Transformer+Crow Attention's 's Cross Attention module (R MDD B + DMDD), and Perceiver IO's Cross Attention module (I. Mx x R Lx MxD).
We evaluated the modules the Copy Task for both GPU and CPU. Although Tree Cross At- tention is slower, all methods only require milliseconds to perform inference. Both Tree Cross Attention and Cross Attention learn solve the task perfectly. However, Tree Cross Attention uses
16

---

Published as a conference paper at ICLR 2024
Model
% Tooenns 100.0%
Accuracy
CPU time
GPU Time
Cross Attention
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
Table 10: Comparison of runtime (in milliseconds) for the inference modules of Re ReTree,er, Trans-
former + Cross Attention, and Perceiver I0, i.c., Tree Cross Attention, Cross Attention, and Per- ceiver IO's CA respectively. Runtime is reported the average of 10 runs.
Model
Tree Height (H)
% Tokens
CPU Time (in ms)
GPU Time (in ms)
Cross Attention
N/A
100.0 %
12.05
1.61
TCA (b, 256)
1
100.0 %
16.21
2.05
TCA (b, =32)
2
15.2 %
14.05
3.99
TCA (bp =16)
2
12.1 %
13.25
4.13
TCA (b, =8)
3
7.0
15.53
5.73
TCA (b, =4)
3
3.9 %
16.30
6.45
TCA (b,- = 2)
8
3.5
21.92
9.83
Table Memory-Runtine trade-off with respect to the tree height (H).
only 3.5% of the tokens. As such, Tree Cross Attention tradesoff between the number of tokens and computational time. Perceiver IO, however, fails to solve the task for the same number of tokens.
The effect of tree design searches down a tree for the relevant tokens. As such, the factor which primarily affects the overall
Memory-Runting trade-off. Tree Cross Attention sequentially
runtime the height tree (H) which is dependent on the branching factor by 2. The relation- ship between the height of the tree, the branching factor, and the number of tokens is as follows: H [lgg,, (N)]. In our results, we showed the performance by using binary tree (b, = since this corresponds to a tree design that uses few tokens. By selecting higher branching factor, the runtime can be decreased at the expense of requiring more tokens (see Table IT]. This feature is not available in standard cross attention. such, TCA enables more control over the memory and computation time trade-off
B.2
ADDITIONAL EXPERIMENTS
B.2.1 COPY TASK
N 128
Method
% Tokens
Accuracy
Cross Attention
100.0%
1.0.0
Random
8.0
Perceiver 10
10.9%
17.8
Perceiver 10
100.0%
1.0
TCA (Acc)
10.9%
100.0
TCA (CE) 10.9%
100.0 0.0
Table 12: Copy Task Results with accuracy (higher is better) and % tokens (lower is better) metries.
we include additional results for the Copy Task where N 128. Similar to previous results, we see that Cross Attention and TCA are able to solve this task perfectly. In contrast, Perceiver IO is not able solve this task for the same number of tokens. We also see both TCA using accuracy and TCA using negative cross entropy as the reward were able to solve the task perfectly.
Results. In Table 12
17

---

Published as a conference paper at ICLR 2024
Method
Tokens
RBF
Matern 5/2
CNP Garnelo et al.
2018a
0.26
0.04 0.02
CANP Kim al.
2019)
0.62 # 0.00
NP Garnelo et al.
2018b
0.27
0.07 #.01
ANP (Kim
2019
0.81 0.00
0.63 0.00
BNP (Lee et al.
2020
+
0.18
BANP (Lee et 2020)
0.01
0.66
TNP-D (Nguyen & Grover
[2022
1.39 0.00
+
LBANP (Feng et 2023a)
1.27
0.05
CMANP
(Feng et al.
20236)
1.22
0.00
Perceiver 10 (L 4)
8.5%
0.03
0.56
Perceiver IO (L 8)
17.0%
0.03
0.65
Perceiver (L 16)
34.0%
0.05
0.75
Perceiver IO (L 32)
68.0%
1.25
0.78
Perceiver IO (L 64)
136.0%
1.30
0.02
Perceiver IO (L 128)
272.0%
1.04
Transformer Cross Attention
100%
1.35 1.35 0.02
0.91+ 0.02
Perceiver IO Re'Treever
14.9%
1.06 + 0.05
0.58 + 0.05
14.9%
1.25 + 0.02
0.81 + 0.02
Table 13: 1-D RegaaReeression Experiments with log- liellikelihood metric (higher is better).
sinke) sinta xt
((x) offfecd wintpi xl/
18
11
-15 -25
13
19
220
-14
40
Figure 12: Visualizations of GP Regression
18

---

Published as a conference paper at ICLR 2024
Method
% Tokens
CelebA
CNP
2.05
CANP
0.06
NP
0.02
ANP
0.00
BNP
2.76
BANP
0.00
TNP-D
3.09
LBANP (8)
3.05
LBANP (128)
0.02
CMANP
0.05
Perceiver 10 (L = 8)
4.1%
0.03
Perceiver IO (L 16)
8.1%
3.02
Perceiver IO (I. = =32)
16.2%
3.00
Perceiver IO (L =64)
32.5%
3.61 0.03
Perceiver IO (L 122)
65.0%
3.02
Transformer + Cross Attention
100%
3.88 0.01
Perceiver IO
4.6%
3.20 - 0.01
ReTreeever
4.6%
3.52 0.01
Table CelebA Image Completion Experiments. The methods are evaluated according to the log- likelihood (higher is better).
B.2.2 GP REGRESSION
Results. In Table 13 we show results for all baselines on GP Regression. We see that Transformer Cross Attention performs comparable to the sfate-of-the-art, outperforming all Neural Process baselines except for TNP-D by significant margin. We evaluated Perceiver IO with varying number of latents L. We see that Perceiver 10 (L = 32) uses ~ 4.6x the number of tokens to achieve performance comparable to ReTreeverr. ReTreever by itself already outperforms all NP baselines except for LBANP and TNP-D. Visualizations for different kinds of test functions are included in Figure 12
B.2.3 IMAGE COMPLETION (CELEBA)
show results for all baselines on CelebA. We that Transformer Cross Attention achieves results competitive with the state-of-the-art. We evaluated Perceiver IO with varying numbers of latents L. We that Perceiver I0 (L 32) uses 3.5x the number of tokens to achieve performance comparable to the ReTreever. Visualizations are available in Figure B3
Results. Table 14
B.2.4 IMAGE COMPLETION (EMNIST)
Results. In Table 15 we see that Transformer + Cross Attention achieves results competitive with state-of-the-are outperforming all baselines except TNP-D.
B.2.5 TIME SERIES: HUMAN ACTIVITY
Results. Table 16 show results on the Human Activity task comparing with several time series baselines. We see that Transformer Cross Attention outperforms all baselines except for miTAND-Ene. The confidence intervals are, however, very close overlapping. We would like to note, however. that mTAND-Ene bidirectional RNNs, and reference points. In contrast, Transformer + Cross Attention a general- purpose model made by just leveraging simple attention modules. We hypothesize the performance could be further improved by leveraging some features mTAND-Enc. but this outside the scope of our work.
carefully designed for time series, combining attention,
Notably, ReTreever also outperforms all baselines except for mTAND-Enc. Comparing, Trans- former Cross Attention and ReTreever, that ReTreever achieves comparable performance
19

---

Published as a conference paper at ICLR 2024
N = 25
N = 50
N 75
N = 100
N 125
N 150
Figure 13: Visualizations of CelebA32
Method
% Tokens
EMNIST
CNP
0.00
CANP
0.01
NP
0.09
ANP
0.98
BNP
BANP
1.00
TNP-D
1.06
LBANP (8)
1.01
LBANP (128)
1.00
CMANP
1.00
Transformer + Cross Attention
100%
1.41
Perceiver I0
4.6%
1.25 0.00
ReTreever
4.6%
1.30 0.01
Table 15: EMNIST Image Completion Experiments. log-likelihood (higher is better).
The methods are evaluated according to the
20

---

Published as a conference paper at ICLR 2024
Model
% Tokens
Accuracy
RNN-Impute CChe al.
2018
0.4
RNN- RNN-Decay Che et al
2018
8.50
RNN GRU-D Che et al.
2018
0.5
IP-Nets (Shukla & Marlin
2019)
8.7
SeFT (Hom et al. 20200
8.2
ODE-RNN (Rubanova et al.
2019)
0.8
L-ODE-R RNN (Chen al. 2018)
83.8
L-ODE-DDE (Rubanova et al. 2019)
87.0 + 2.8
mTAND-Ene (Shukla Marlm
2021
90.7 0.2
Transformer Cross Attention
100%
89.1 1.3
Perceiver IO
14%
87.6 + 0.4
ReTreever
14%
88.9 0.4
Table 16: Human Activity Experiments with Accuracy metric.
while using far fewer tokens. Furthermore, we see that ReTreever outperforms Perceiver 10 signifi- cantly while using the same number of tokens.
C APPENDIX: IMPLEMENTATION AND HYPERPARAMETERS DETAILS
C.1 IMPLEMENTATION
For the experiments on uncertainty estimation, we use the official repositories for TNPs htwp:/ittwu..ooouuuunnnn////PPPP
and LBANPs https///gguubbccoo/oreeliiiAllli bottlenecked-ang anp). The GP regression and Image Classification (CelebA and EMNIST) datasets are available the same links. For the Human Activity dataset, we use the official repository for mTAN (Multi-Time Attention Networks) The Human Activity dataset is available in the same link. Our Perceiver 10 baseline is based the popular Perceiver (IO) repository htttss//gihhbb.ccm///cdirrressss
http:///ttuubbcoo///mmllllllb///
py). We report the baseline results listed Nguyen & Grover (2022) and (Shukla & Marlin) (2021). When comparing methods such as Perceiver, Re'Treever, Transformer + Cross Attention, LBANP, and TNP-D, we use the same number blocks (i.e., CMABs, iterative attention, and transformer) in the encoder.
C.2 HYPERPARAMETERS
Following previous works Neural Processes (LBANPs and TNPs), ReTreever in the encoder for all experiments. Our aggregator function is a Self Attention (Transformer) module whose output is averaged. Following standard RL practices, test time, TCA and Re- Treever's policy selects the most confident actions: u argmax, (a(a))) We tuned Ant (RL loss term weight) and λсл (CA loss term weight) between {0.0, 0.1, 1.0, 10.0) the Copy Task (see analysis in main paper) with N 256. We also tuned (entropy bonus weight) between (0.0, 0.01, 0.1, 1.0]. We found that simply setting ARL λAA 1.0 and a 0.01 performed well in practice on the Copy Task where N = 256. As such, for the purpose of consistency, we 0.01 was for all tasks. We used ADAM optimizer with
layers
set ARL AC 1.0 and standard learning rate of 5e 4. All experiments were run with 3 seeds. Cross Attention and Tree Cross Attention tended to get stuck a local optimum on the Copy Task, so to prevent this, dropout was set to 0.1 for the Copy Task for all methods. For image data, we selected a single axis to split the data according to the k-d tree. In the case of EMNIST, the data was split according to the first axis. In the case of CelebA, the data split according to the second axis.
C.3 COMPUTE
Our experiments were run using a mix of Nvidia GTX 1080 Ti (12 GB) or Nvidia Tesla P100 (16 GB) GPUs. GP regression experiments took 2.5 hours. EMNIST experiments took 1.75 hours. CelebA experiments took ~ 15 hours. Human Activity experiments took ~ 0.75 hours.
21


Reviews:
[{"review_id": "eq3NhPbWih", "reviewer": "Reviewer_RnYk", "summary": "This paper proposes TCA, a tree-based cross attention module to reduce the complexity of cross attention from $O(N)$ to $O(\\log(N))$, where $N$ is the number of tokens used for cross attention. Given $N$ tokens, TCA first constructs a balanced binary tree representation using standard methods like K-D tree, where the leaf nodes are the token embeddings, and the internal node representations are aggregated using the two children of the internal node. TCA uses reinforcement learning to learn good internal node representations. This construction is only performed once for a set of context tokens. Now, for a given query vector, a tree search is performed to select a subset of nodes ($O(\\log(N))$ size) of the tree for cross attention, resulting in $O(\\log(N))$ complexity for retrieval. Using TCA, the paper further proposes ReTreever, a general-purpose retrieval model that achieves token-efficient inference. The paper compares the ReTreever models with other token-efficient retrieval models like Perceiver IO and show impressive gains over the baseline - little to no drop in performance while leveraging only a small subset of tokens for cross attention.", "questions": "I have asked most of my questions in the weakness section. If the authors can address my questions and add the relevant latency benchmarks too, I am willing to increase my score.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- The paper is well written and builds the theory coherently.\n- The proposed cross-attention architecture, TCA, along with the general purpose retrieval model, ReTreever is novel.\n- Because ReTreever uses reinforcement learning to learn the internal node representations, the reward used for optimization can be non-differentiable like accuracy, which improves performance over a reward based on cross entropy because the reward model is simpler in case of accuracy.\n- The reasoning behind each of the loss terms in $\\mathcal{L_{ReTreever}}$ is well-explained and it also uses leaf-level cross attention loss to make the training faster.\n- The empirical results on various tasks like copying, uncertainty estimation are impressive using ReTreever, and the paper also has good ablation studies to test the various components of the proposed approach.", "weaknesses": "- 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.\n- 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. I am guessing the tree approach is not parallelizable on accelerated devices like GPUs, but it would be good to see if there's considerable decrease in latency on CPUs.\n- Building on the previous point, wall-clock times for the tree construction and bottom-up aggregation should be reported too.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ne4PbK1WHW", "reviewer": "Reviewer_779Z", "summary": "This paper presents a method for reducing the test-time computational cost of attention. Namely, TCA walks down a tree to attend to a single set of sibling leaves, while only attending to a compressed version of the other leaves. This allows the complexity of attention to be logarithmic in the total number of leaves $N$ (the sequence length). This comes at the cost of training a policy to traverse the tree, which must be trained via REINFORCE, an aggregator that compresses and composes leaf representations, as well as defining the tree itself.\n\nExperiments on a copy task, GP regression, image completion, human activity classification show that the method is efficient and performant. Additional analysis highlights the method's memory efficiency compared to full attention.", "questions": "## Questions\n1. Would TCA work out of the box for masked language modeling, e.g. BERT?\n2. Did you try using Gumbel-softmax for training the tree expansion policy?\n3. What are the barriers to applying TCA to self-attention? Would aggregation become the most expensive operation?\n\n## Suggestions\n1. 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.\n2. 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.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I enjoyed reading the paper. The paper is well-written and easy to follow. The idea is simple and clever.", "weaknesses": "Overall, I believe the paper is pretty complete. I am mostly curious about how to make this method work for self-attention and (masked) autoregressive modeling.\n\nLarger-scale experiments would be appreciated, as the current experiments are quite small-scale. Presumedly the challenges of training the tree expansion policy would increase with harder datasets.\n\nOne suggestion for a larger-scale experiment would be training a translation or summarization model and replacing the encoder attention with tree cross attention.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WkyYgWXIPD", "reviewer": "Reviewer_Wmtm", "summary": "This work proposes a tree cross attention (TCA) architecture to effectively encode context into a tree structure when performing attention, such that the amount of attended tokens will be less than the actual number of context tokens. The architecture is co-trained with the objective to effectively retrieve context through graph-based searching, and utilize context for performing the task. Experiments on a specific set of tasks (i.e., copy task, gp regression on image completion and time series on human activities) show that TCA is able to achieve a similar performance as full cross attention while attending to much fewer tokens. \n\nThe architecture proposed is generally applicable to a lot of settings broadly, yet the evaluation is too specific and does not adequately proves the generality of the method.", "questions": "- Does the algorithm run faster than full attention in terms of wall-clock time?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The idea to construct context as a tree is interesting and could have broad implications in constructing context for language models use cases including agent trajectories, in-context learning examples and retrieved documents and more.", "weaknesses": "- **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. \n- **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. \n- **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kYq70q0i78", "reviewer": "Reviewer_rpWY", "summary": "1. The authors introduce a new module called Tree Cross Attention (TCA) that reduces the number of tokens required for efficient inference while maintaining comparable performance to Cross Attention. \n\nTCA organizes data in a tree structure and performs a tree search at inference time to retrieve relevant tokens for prediction. Specifically, TCA only retrieves information from a logarithmic O(log(N)) number of tokens for performing inference, while Cross Attention scans the full set of O(N) tokens. \n\n2. The authors also present ReTreever, a flexible architecture for token-efficient inference that incorporates TCA. \n\n3. Empirically, the paper demonstrates the effectiveness of TCA and ReTreever on various classification and uncertainty regression tasks.", "questions": "1. My main question regarding the evaluation is in the weakness section.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and straightforward to understand.\n\n2. The empirical results that the authors show are impressive. Specifically, the authors show that Perceiver IO's performance drops significantly as the length of the sequence increases, while ReTreever is able to maintain high accuracy across a range of sequence lengths.\n\n3. The proposed method, ReTreever, is able to perform token-efficient inference while achieving better performance than Perceiver IO for the same number of tokens. ReTreever does this by using Tree Cross Attention to retrieve the necessary tokens, only needing a logarithmic number of tokens log(N) << N, making it efficient regardless of the encoder used.", "weaknesses": "1. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)