# Treeformer: Dense Gradient Trees for Efficient Attention Computation

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 36 references that are in the map.
- First posted on arXiv 18 Aug 2022.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [Multilingual Denoising Pre-training for Neural Machine Translation](https://mlsanity.com/paper/2001.08210/) (2020)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [Efficient Content-Based Sparse Attention with Routing Transformers](https://mlsanity.com/paper/2003.05997/) (2020)
- [A Simple Method for Commonsense Reasoning](https://mlsanity.com/paper/1806.02847/) (2018)
- [Scatterbrain: Unifying Sparse and Low-rank Attention Approximation](https://mlsanity.com/paper/2110.15343/) (2021)
- [O(n) Connections are Expressive Enough: Universal Approximability of Sparse Transformers](https://mlsanity.com/paper/2006.04862/) (2020)
- [Memory-efficient Transformers via Top-k Attention](https://mlsanity.com/paper/2106.06899/) (2021)
- [SparseBERT: Rethinking the Importance Analysis in Self-attention](https://mlsanity.com/paper/2102.12871/) (2021)
- [Learning Accurate Decision Trees with Bandit Feedback via Quantized Gradient Descent](https://mlsanity.com/paper/2102.07567/) (2021)

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

- [Training-free Diffusion Model Adaptation for Variable-Sized Text-to-Image Synthesis](https://mlsanity.com/paper/2306.08645/) (2023)
- [Log-Linear Attention](https://mlsanity.com/paper/2506.04761/) (2025)
- [Efficient Attention via Control Variates](https://mlsanity.com/paper/2302.04542/) (2023)
- [HiRE: High Recall Approximate Top-k Estimation for Efficient LLM Inference](https://mlsanity.com/paper/2402.09360/) (2024)
- [Tree Cross Attention](https://mlsanity.com/paper/2309.17388/) (2023)
- [Transparent Networks for Multivariate Time Series](https://mlsanity.com/paper/2410.10535/) (2024)
- [Multiple Instance Learning with Trainable Decision Tree Ensembles](https://mlsanity.com/paper/2302.06601/) (2023)
- [A Faster Generalized Two-Stage Approximate Top-K](https://mlsanity.com/paper/2506.04165/) (2025)
- [FreeSolo: Linear-time Sparse Attention with Zero-Shot Length Extrapolation](https://mlsanity.com/paper/2610.17191/) (2026)

## Abstract (arXiv)

Standard inference and training with transformer based architectures scale quadratically with input sequence length. This is prohibitively large for a variety of applications especially in web-page translation, query-answering etc. Consequently, several approaches have been developed recently to speedup attention computation by enforcing different attention structures such as sparsity, low-rank, approximating attention using kernels. In this work, we view attention computation as that of nearest neighbor retrieval, and use decision tree based hierarchical navigation to reduce the retrieval cost per query token from linear in sequence length to nearly logarithmic. Based on such hierarchical navigation, we design Treeformer which can use one of two efficient attention layers -- TF-Attention and TC-Attention. TF-Attention computes the attention in a fine-grained style, while TC-Attention is a coarse attention layer which also ensures that the gradients are "dense". To optimize such challenging discrete layers, we propose a two-level bootstrapped training method. Using extensive experiments on standard NLP benchmarks, especially for long-sequences, we demonstrate that our Treeformer architecture can be almost as accurate as baseline Transformer while using 30x lesser FLOPs in the attention layer. Compared to Linformer, the accuracy can be as much as 12% higher while using similar FLOPs in the attention layer.