# Tree Cross Attention

arXiv: https://arxiv.org/abs/2309.17388 · HTML: https://mlsanity.com/paper/2309.17388/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 24 references that are in the map.
- First posted on arXiv 29 Sep 2023.

## Builds on (most cited first, 18 of 18)

- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [A Comprehensive Survey on Graph Neural Networks](https://mlsanity.com/paper/1901.00596/) (2019)
- [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://mlsanity.com/paper/2107.14795/) (2021)
- [Recurrent Neural Networks for Multivariate Time Series with Missing Values](https://mlsanity.com/paper/1606.01865/) (2016)
- [Conditional Neural Processes](https://mlsanity.com/paper/1807.01613/) (2018)
- [Neural Processes](https://mlsanity.com/paper/1807.01622/) (2018)
- [Improved Semantic Representations From Tree-Structured Long Short-Term Memory Networks](https://mlsanity.com/paper/1503.00075/) (2015)
- [Attentive Neural Processes](https://mlsanity.com/paper/1901.05761/) (2019)
- [Multi-Time Attention Networks for Irregularly Sampled Time Series](https://mlsanity.com/paper/2101.10318/) (2021)
- [Set Functions for Time Series](https://mlsanity.com/paper/1909.12064/) (2019)
- [Transformer Neural Processes: Uncertainty-Aware Meta Learning Via Sequence Modeling](https://mlsanity.com/paper/2207.04179/) (2022)
- [Interpolation-Prediction Networks for Irregularly Sampled Time Series](https://mlsanity.com/paper/1909.07782/) (2019)
- [Tree Transformer: Integrating Tree Structures into Self-Attention](https://mlsanity.com/paper/1909.06639/) (2019)
- [Bootstrapping Neural Processes](https://mlsanity.com/paper/2008.02956/) (2020)
- [Tree-structured Attention with Hierarchical Accumulation](https://mlsanity.com/paper/2002.08046/) (2020)
- [Latent Bottlenecked Attentive Neural Processes](https://mlsanity.com/paper/2211.08458/) (2022)
- [Graph Neural Networks Designed for Different Graph Types: A Survey](https://mlsanity.com/paper/2204.03080/) (2022)
- [Treeformer: Dense Gradient Trees for Efficient Attention Computation](https://mlsanity.com/paper/2208.09015/) (2022)

## Built on by (most cited first, 2 of 2)

- [Denoising Diffusion Generative Models Secretly Calculate Attentions](https://mlsanity.com/paper/2609.00885/) (2026)
- [FLaG: Frequency-Domain Latent-attention Gated Pooling for Token Aggregation](https://mlsanity.com/paper/2609.00831/) (2026)

## Abstract (arXiv)

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