Research map / Architectures and efficiency / Attention and transformers
Tree Cross Attention
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Tree Cross Attention (TCA), which organizes context tokens in a k-d tree and uses an RL-learned policy to retrieve only O(log N) nodes for cross attention, and ReTreever, an encoder plus TCA architecture for token-efficient inference.
Problem
Cross attention at inference scans all O(N) context tokens, while latent-bottleneck methods like Perceiver IO compress information into a fixed number of latents that may not suffice for high-intrinsic-dimensionality problems.
Approach
Organize context tokens as leaves of a balanced k-d tree with aggregated internal nodes, and learn a REINFORCE-trained tree-search policy (parameterized as attention over children) that selects a logarithmic set of nodes for cross attention, allowing non-differentiable rewards such as accuracy.
Key results
- On the Copy Task at N=1024, TCA reaches 99.6% accuracy using 2.0% of tokens versus Perceiver IO's 11.6%.
- On 1-D GP regression with RBF kernel, ReTreever scores 1.25 log-likelihood at 14.9% tokens versus Perceiver IO's 1.06 and Transformer+CA's 1.35 at 100%.
- On Human Activity, ReTreever reaches 88.9% accuracy using 14% of tokens, close to Transformer+CA's 89.1% and above Perceiver IO's 87.6%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace Perceiver IO with TCA using the same token budget | Copy Task · Copy Task N=256 | accuracy: 15.2 ± 0.0 -> 100.0 ± 0.0 | Helped |
| replace Perceiver IO with ReTreever using the same token budget | GP Regression · GP Regression RBF | log-likelihood: 1.06 ± 0.05 -> 1.25 ± 0.02 | Helped |
| replace Perceiver IO with ReTreever using the same token budget | ReTreever · CelebA | log-likelihood: 3.20 ± 0.01 -> 3.52 ± 0.01 | Helped |
| replace Perceiver IO with ReTreever using the same token budget | ReTreever · EMNIST | log-likelihood: 1.25 ± 0.01 -> 1.30 ± 0.01 | Helped |
| replace Perceiver IO with ReTreever using the same token budget | ReTreever · Human Activity | accuracy: 87.6 ± 0.3 -> 88.9 ± 0.4 | Helped |
| use accuracy rather than negative cross entropy as the RL reward (ablation) | TCA · Copy Task N=1024 | accuracy: 80.3 ± 14.8 -> 99.6 ± 0.6 | Helped |
| organize the tree by image structure rather than randomly (ablation) | ReTreever · CelebA | log-likelihood: 3.12 ± 0.02 -> 3.52 ± 0.01 | Helped |
| retrieve all tree leaves instead of a policy-selected subset (ablation) | ReTreever · CelebA | log-likelihood: 3.52 ± 0.01 -> 3.77 ± 0.01 | Helped |
| retrieve all tree leaves instead of a policy-selected subset (ablation) | ReTreever · GP Regression RBF | log-likelihood: 1.25 ± 0.02 -> 1.33 ± 0.02 | Helped |
| set the RL loss weight to 0.0 or 0.1 rather than 1.0 (ablation) | ReTreever · Copy Task | accuracy | Hurt |
| set the RL loss weight to 10.0 rather than 1.0 (ablation) | ReTreever · Copy Task | accuracy | Hurt |
| set the cross-attention loss weight to 0.0 or 0.1 rather than 1.0 (ablation) | ReTreever · Copy Task | training stability | Hurt |
Surprising findings
- Perceiver IO with 100% of tokens on the Copy Task at N=1024 ran out of memory, and Perceiver IO performance degraded as sequence length grew even with more latents.
- Using accuracy as the RL reward outperformed negative cross entropy on the Copy Task (80.3% vs 99.6% at N=1024 for the CE reward variant).
Limitations
- Wall-clock inference time of TCA is higher than standard cross attention in the reported measurements.
- ReTreever's performance depends on the number of selected nodes and falls short of ReTreever-Full on GP regression and CelebA.
- Tree construction uses hand-designed k-d tree heuristics rather than learned organization.
Experimental setup
Datasets: Copy Task (synthetic palindromes) GP regression (synthetic RBF/Matern GP functions) CelebA EMNIST Human Activity
Benchmarks: Copy Task GP Regression (RBF, Matern 5/2) CelebA image completion EMNIST image completion Human Activity
Metrics: accuracy log-likelihood percent tokens used CPU/GPU runtime memory usage
Baselines: Perceiver IO Transformer + Cross Attention TNP-D LBANP CMANP mTAND-Enc
Builds on: Perceiver IO REINFORCE k-d tree
Closest papers by meaning
- DeFT: Decoding with Flash Tree-attention for Efficient Tree-structured LLM Inference (ICLR 2025 · similarity 0.85)
- Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction (ICLR 2025 · similarity 0.84)
- Learning to Focus: Causal Attention Distillation via Gradient‐Guided Token Pruning (NeurIPS 2025 · similarity 0.84)
- Efficient Length-Generalizable Attention via Causal Retrieval for Long-Context Language Modeling (ICML 2025 · similarity 0.84)
- CLOVER: Cross-Layer Orthogonal Vectors Pruning (ICML 2025 · similarity 0.84)
- OrthoRank: Token Selection via Sink Token Orthogonality for Efficient LLM inference (ICML 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.