Tree Attention: Topology-Aware Decoding for Long-Context Attention

## Abstract

Abstract Self-attention is the core mathematical operation of modern transformer architectures and is also a significant computational bottleneck due to its quadratic complexity in the sequence length.
In this work, we derive the scalar energy function whose gradient computes the self-attention block, thus elucidating the theoretical underpinnings of self-attention, providing a Bayesian interpretation of the operation and linking it closely with energy-based models such as Hopfield Networks.
Our formulation reveals that the reduction across the sequence axis can be efficiently computed in parallel through a tree reduction.
Our algorithm, for parallelizing attention computation across multiple GPUs enables cross-device decoding to be performed asymptotically faster (up to 8 × 8\times faster in our experiments) than alternative approaches such as Ring Attention, while also requiring significantly less communication volume and incurring 2 × 2\times less peak memory.
Our code is publicly available here: https://github.com/Zyphra/tree_attention .

## 1.  Introduction

The self-attention operation is the core computational building block of the transformer architecture Bahdanau *et al.* (); Vaswani *et al.* (), which has become an ubiquitous and highly effective workhorse architecture currently applied at scale to language Brown *et al.* (); Kaplan *et al.* (); Hoffmann *et al.* (); Team *et al.* (); Achiam *et al.* (), vision Dosovitskiy *et al.* (), audio Betker (), and decision-making Chen *et al.* (); Reed *et al.* ().
Nonetheless, the quadratic time complexity of self-attention means that significant resources are required to train and generate from transformer-based Large Language Models (LLMs), especially for models with large context lengths.

During inference, the attention block largely determines the computational and memory requirements, which become more demanding as the input sequence length increases.
Although LLMs generate one token at a time, the entire sequence of past tokens must still be stored in memory and used to compute attention scores during generation.
Since attention performs a similarity matching of every token representation with every other, it incurs quadratic computational complexity in terms of flops.

There have been recent advances in training LLMs to handle extremely long contexts (up to 1M tokens)Chen *et al.* (); kai (); Peng *et al.* ().Such models attain qualitatively new capabilities such as extremely large-scale in-context learning of entire small datasets held in the prompt Reid *et al.* (); Lee *et al.* (); Bertsch *et al.* (). They can also avoid putting multi-modal continuous data through a lossy tokenization scheme Reid *et al.* (); Team () by directly operating at the byte level Xue *et al.* (); Wu *et al.* (). The issue however is that performing inference on such long contexts is very expensive.

To speed up inference and alleviate memory requirements, recent works have attempted to alter the attention mechanism itself, either by linearizing it Katharopoulos *et al.* (), or approximating it by a kernel map Choromanski *et al.* (); Peng *et al.* (); Arora *et al.* (), which reduces the complexity to linear at the cost of reduced expressiveness.
Others have invented alternative sequence mixing architectures such as state-space models which are designed to be efficiently computable in linear time and constant memory Gu and Dao (); Dao and Gu (); Katsch (); Sun *et al.* (); Glorioso *et al.* ().
Other approaches utilize efficient algorithms to reduce the computational burden of attention while keeping the core computation the same.
These include memory-efficient attention Rabe and Staats (), Flash Attention Dao *et al.* () and Flash Decoding fla (), which provide a set of IO-aware kernels to map the attention operation to the GPU hardware resources in an extremely efficient way, significantly reducing the memory overhead required.
Further works Character AI (); Kang *et al.* (); Liu *et al.* (); Nawrot *et al.* () explore compressing or otherwise reducing the KV cache required in generation.
Finally, Ring Attention Liu *et al.* () proposes a way to parallelize the attention computation across the sequence axis between GPUs, thus enabling significantly longer contexts than can be served on a single GPU.
This is the regime of primary interest of this paper.
By leveraging the exact energy
function for the self-attention block, we develop a method to speed up inference for long context use-cases when keys and values are sharded across multiple GPUs along the sequence axis.

Our proposed algorithm for computing attention via the gradient of the energy function is built on top of an efficient parallel computation and tree reduction communication strategy.
In particular, this formulation lets us devise an asymptotically faster algorithm for performing decoding in which the number of communication steps scales logarithmically with the number of devices, instead of linearly in alternatives such as Ring Attention Liu *et al.* ().
Our topology-aware approach illustrated in Fig. 1. significantly outperforms leading attention parallelization methods such as Ring Attention on multiple devices.

Figure: (a) Ring and Tree Attention Topologies. Due to the associative properties of the logsumexp and max operations of Tree Attention (Fig. LABEL:fig:tree), is possible to structure the reduction across the sequence as a tree, requiring asymptotically fewer communication steps than Ring Attention (Fig. LABEL:fig:ring) as well as less memory and communications volume.
[Uncaptioned image]: fig/legend.png

## 2.  Related Work

Following the ubiquitous success of the transformer architecture, there has been significant effort to mathematically understand the nature and meaning of the attention operation and link it to related architectures, such as recurrent networks and Hopfield Networks.
A number of recent works have attempted to study self-attention mathematically through the lens of energy functions.
Ramsauer *et al.* () pioneered this field by performing a similar but distinct analysis to relate self-attention with the modern Hopfield networks Krotov and Hopfield (), providing a novel and insightful interpretation of self-attention as performing hetero-associative memory lookups using a high-powered nonlinear similarity function Krotov (); Millidge *et al.* ().
This work was later extended by Hoover *et al.* (), who derived a modified version of the transformer based off an energy function.
Beyond this, other recent work attempted a Bayesian reformulation of attention by deriving a probabilistic generative model which matches the operations performed in a self-attention operation Singh and Buckley ().

However, while it has long been known that the softmax operation can be derived as the gradient of the following scalar function:

$$ $\partial_{z_{j}}\log\sum^{n}_{a=1}\exp(z_{a})=\frac{e^{z_{j}}}{\sum^{n}_{a=1}e^{z_{a}}}=\textrm{softmax}(z_{j}),$ (1) $$

known as the log-sum-exp, an equivalent function for the self-attention block has not yet been derived.

In particular, consider the energy function proposed in Hoover *et al.* (), which is the logsumexp.
Since the gradient in the update rule of that paper is taken with respect to the input to the block, the resulting function is a modified version of the self-attention operation.
Similarly, the update rule in Ramsauer *et al.* () requires the tying of certain weights (K and V) in the attention operation.
This restricts the Hopfield derivation to modelling auto-associative retrieval networks, while transformer attention is hetero-associative.

To advance beyond this, our energy function depends on an auxiliary source vector $\zeta$ (see section 4.). Taking the gradient with respect to $\zeta$ yields the exact self-attention operation, as we show below. Consequently, this formulation enables us to express self-attention with no weight tying.

Another notable related work is Feng *et al.* (), where the authors made similar observations as we do in section 5. about how the associative operations within the attention computation can be efficiently parallelized to motivate an attention-based modified RNN architecture for sequence modeling.

While this energy function by itself is primarily a mathematical and theoretical curiosity, we demonstrate below that when combined with automatic differentiation, our formulation naturally leads to highly efficient parallel algorithms for computing attention and performing decoding, especially across multiple devices.

## 3.  Self-Attention

The self-attention operation can be represented as a set of dot product similarity searches between queries and keys.
These similarity scores are then reduced along the sequence axis and softmaxed, so that for a given query, there is a probability distribution of the similarities of each given key. We then take the expectation of the value vectors against this distribution.
We denote the queries assigned to a sequence of length $N$ as $\{q_{a},a=1,\cdots N\}$, where each query is a vector of size $d$ that stands for hidden dimension, $q_{a}\in\mathbb{R}^{d}$, and similarly the keys and values $\{(k_{a},v_{a}),a=1,\cdots N\}$.
Attention can be written as

$$ $\displaystyle z^{a}=\sum^{N}_{i=1}\text{softmax}(q_{a}\cdot k_{i}^{T})v_{i}\,.$ $$

The dot product, more explicitly, reads:

$$ $\displaystyle q_{a}\cdot k_{i}=\sum^{d}_{A=1}q_{a,A}k_{i,A},\forall A\in\{1,\cdots,d_{h}\}.$ $$

The capital Latin indices ($A$) span the hidden dimension $d_{h}$.
Naively computing attention in this way requires materializing the $qk$ matrix with computational and memory cost quadratic in the sequence length.

Memory-efficient attention Rabe and Staats () is an iterative way to compute the softmax similarities without ever having to materialize the full attention matrix.
It performs the following operations, one query (or a chunk of queries) at a time:

$$ $\displaystyle s^{(j)}_{i}$ $\displaystyle=\text{exp}(q_{j}\cdot k_{i})$ (2) $\displaystyle n^{(j)}_{i}$ $\displaystyle=n^{(j)}_{i-1}+v_{i}s^{(j)}_{i}$ (3) $\displaystyle d^{(j)}_{i}$ $\displaystyle=d^{(j)}_{i-1}+s^{(j)}_{i}$ (4) $$

Then, once the values $v$ and softmax denominator $d$ are computed, we divide to get the final softmaxed scores $z^{(j)}=\frac{n^{(j)}}{d^{(j)}}$ for every query index $j$.
Computing attention in this iterative manner significantly reduces the required memory.

Flash attention Dao *et al.* () utilizes a similar approach to reduce the memory and computational cost of attention, but the algorithm is not adapted for multi-GPU computation.
Flash attention performs the iterative algorithm of Rabe and Staats () in a blockwise manner, utilizing the block-parallel computational primitives available inside single GPU tensor cores.
Additionally, it precisely sizes the blocks such that they can fit into the SRAM of the GPU for the entire attention computation, effectively performing kernel fusion and preventing many unnecessary IO operations.

## 4.  Self-Attention as a Gradient

Here, we show how the self-attention operation can be written as the gradient of a certain scalar function. In particular, we define a scalar function that depends on the keys, queries, values and additionally on an auxiliary vector that we refer to as the source $\zeta$.
The source is the parameter with respect to which we compute the gradient of the scalar function to obtain the self-attention operation.
We need the source in order to write down the generating function of the moments of the distribution above.
It is also the variable with respect to which we can Taylor-expand the generating function and extract the moments as the coefficients of the monomials of $\zeta$ appearing in the Taylor series.
Explicitly, we want to find a function $F(q,k,v,\zeta)$ such that:

$$ $\sum^{N}_{a=1}\textrm{softmax}(q\cdot k_{a})v_{a}=\frac{\partial F}{\partial\zeta}\bigg|_{\zeta=0}.$ (5) $$

This terminology is inspired by work on energy-based models in machine learning Beal (); LeCun *et al.* (); Song and Kingma ().
A summary of variables and indices is provided in appendix  12.

### 4.1.  The energy function

Here, we show how the energy function is given by the cumunlant-generating function associated to the distribution given by attention scores.
Taking inspiration from statistical mechanics, where an analogous cumulant-generating function defines the Helmholtz Free energy Wikipedia (), we dub our cumulant-generating function the energy function for self-attention.

Let us focus on the case with a single query.
As noted above, we leverage the fact that the attention operation can be seen as the computation of the expectation value of the vectors $v$ in the distribution set by the attention scores $z$:

$$ $z=\langle v\rangle=\sum^{N}_{a=1}P_{a}v_{a}=\frac{\sum^{N}_{a=1}e^{q\cdot k^{T}_{a}}v_{a}}{\sum^{N}_{i=1}e^{q\cdot k^{T}_{i}}}.$ (6) $$

The probability density is given by:

$$ $P_{a}=\frac{e^{q\cdot k^{T}_{a}}}{\sum^{N}_{i=1}e^{q\cdot k^{T}_{i}}},$ (7) $$

and satisfies:

$$ $\sum^{N}_{a=1}P_{a}=1.$ (8) $$

Typically, the denominator or normalization factor is identified with the so-called partition function:

$$ $Z=\sum^{N}_{a=1}e^{q\cdot k^{T}_{a}}.$ (9) $$

We can now compute the first moment of the probability distribution given above by introducing a source, $\zeta\in\mathbb{R}^{d}$.
In our case, with $\zeta,$ we can extend the partition function to the function:

$$ $Z(\zeta)=\sum^{N}_{a=1}e^{q\cdot k^{T}_{a}+\zeta\cdot v^{T}_{a}}.$ (10) $$

Now, we can compute any moment of the distribution as the $n$-th Taylor coefficient of $Z(\zeta)$ $\forall A_{1},A_{2},\cdots\in\{1,\cdots,d_{h}\}$ :

$$ $\langle v_{A_{1}}\cdots v_{A_{n}}\rangle=\frac{1}{Z}\frac{\partial^{n}Z(\zeta)}{\partial\zeta_{A_{1}}\cdots\partial\zeta_{A_{n}}}\bigg|_{\zeta=0}.$ (11) $$

In other words, we can write $Z(\zeta)$
as:

$$ $Z(\zeta)=Z\left(1+\langle v\rangle\zeta+\frac{1}{2!}\langle v_{A_{1}}v_{A_{2}}\rangle\zeta_{A_{1}}\zeta_{A_{2}}+\cdots\right)$ (12) $$

Therefore, the first moment can be written as:

$$ $\langle v\rangle=\frac{1}{Z}\frac{\partial Z}{\partial\zeta}\bigg|_{\zeta=0},$ (13) $$

which can be written as the gradient of the log of $Z(\zeta)$:

$$ $\langle v\rangle=\frac{\partial}{\partial\zeta}\log Z(\zeta)\bigg|_{\zeta=0}.$ (14) $$

This quantity is the generating function, a.k.a. the free energy:

$$ $F=\log\sum_{a}\exp\left(q\cdot k_{a}^{T}+\zeta\cdot v^{T}_{a}\right).$ (15) $$

To compute causal self-attention, we introduce $N$ sources $\zeta^{i}$ each $\in\mathbb{R}^{d}$ and take

$$ $F_{tot}=\sum^{N}_{i=1}F_{i}=\sum^{N}_{i=1}\log\sum^{i}_{a=1}\exp(q_{i}\cdot k^{T}_{a}+\zeta_{i}\cdot v^{T}_{a}).$ (16) $$

The truncation of the inner sum up to index $i$ is due to causal masking.

Now, in order to compute the $i$-th element of causal self-attention, we differentiate with respect to $\zeta_{i}$ and set it to zero:

$$ $\frac{\partial F_{tot}}{\partial\zeta_{i,A}}\bigg|_{\zeta_{i}=0,\forall i}=\frac{\sum^{i}_{a=1}\exp(q^{i}\cdot k^{T}_{a})v_{a,A}}{\sum^{i}_{a=1}\exp(q^{i}\cdot k^{T}_{a})}.$ (17) $$

The generalization to the multi-head attention case is straightforward.
In this case, there is one key, query and value per head.
For $n_{h}$ total heads, the generating function takes the form:

$$ $F_{tot}=\sum^{N}_{i=1}\sum^{n_{h}}_{h=1}F^{i,h},$ (18) $$

where

$$ $F_{i,h}=\log\sum^{i}_{a=1}\exp\left(q_{i,h}\cdot k_{h,a}^{T}+\zeta_{h,i}\cdot v^{T}_{h,a}\right).$ (19) $$

The output projection weight is included in the definition of $v_{j}$ here, meaning that

$$ $v_{b,A}=x_{b,\bar{B}}(W_{O}W_{V})_{A,\bar{B}}$ (20) $$

where $W_{O}\in\mathbb{R}^{d_{h}}\times\mathbb{R}^{d_{emb}}$ denotes a head size slice of the output projection weight and $\bar{B}\in\{1,\cdots,d_{h}\}$ spans the intra-head indices.
In the index notation above, the head indices are barred whereas the embedding space indices are unbarred.
We proceed focusing on the single-head case, as it makes the presentation simpler, and the multi-head generalization is immediate.
Note that we demonstrate that our energy function approach also can account for safe softmax in Appendix 11.

### 4.2.  Bayesian interpretation

The fact that it is possible to derive the self-attention operation as the minimization of an energy function implies that it is possible to provide a Bayesian gloss on self-attention by identifying a likelihood function and showing that we can obtain the forward pass of the attention block from computing the maximum a posteriori estimate of this likelihood.

In particular, we propose the following for the log-likelihood function:

$$ $\Gamma(\zeta,z)=\sum^{N}_{i=1}\sum^{d}_{A=1}\bigg(z_{i,A}\zeta_{i,A}-F(\zeta,x)\bigg).$ (21) $$

We denote by $x$ the input to the self-attention block from which we obtain $q,k,v$ from multiplying it by the weights $W_{Q},W_{K},W_{V}$ respectively.
Let us minimize the above with respect to $\zeta$ and $z$ simultaneously:

$$ $\displaystyle\frac{\partial\Gamma}{\partial\zeta_{i,A}}=0,$ (22) $\displaystyle\frac{\partial\Gamma}{\partial z_{i,A}}=0.$ (23) $$

These conditions written explicitly read

$$ $\zeta_{i,A\,*}=0,\,\,\,\,z_{i,A\,*}=\frac{\partial F}{\partial\zeta_{i,A}}.$ (24) $$

Plugging the first condition into the second leads to the attention forward pass:

$$ $z_{i\,*,A}=\frac{\sum^{i}_{a=1}e^{q_{i}\cdot k^{T}_{a}}v_{a,A}}{\sum^{i}_{b=1}e^{q_{i}\cdot k^{T}_{b}}}.$ (25) $$

In all, this means we can obtain the gradient w.r.t. $\zeta$ from MAP estimation of the following likelihood:

$$ $z_{i\,*,A},\zeta_{i\,*,A}=\textrm{argmax}_{\zeta,z}e^{-\Gamma(\zeta,z)}.$ (26) $$

Moreover, such a procedure enables us to identify the energy-based model associated with the self-attention function.

## 5.  Tree Attention

In this section we show how the formulation of the attention operation as the gradient of an energy function suggests an efficient parallel strategy for computing it. The key insight is to leverage an efficient algorithm to compute the energy, and then differentiate it in order to obtain an efficient algorithm to compute attention.

### 5.1.  Computing the energy efficiently

Let us focus on the case of decoding with a KV cache in a causal language model where we have one query and $N$ keys and values. In this case, the energy function is:

$$ $\displaystyle F_{dec}=\log\sum^{N}_{a=1}\exp(q\cdot k_{a}^{T}+\zeta\cdot v^{T}_{a})$ (27) $\displaystyle\equiv\textrm{logsumexp}_{a}(\{q\cdot k^{T}_{a}+\zeta\cdot v^{T}_{a},a=1,\cdots,N\}).$ (28) $$

A crucial fact is that both $\textrm{logsumexp}_{a}$ and $\max_{a}$ are associative operations:

$$ $\displaystyle\textrm{logsumexp}_{a}(\{T_{a},\textrm{logsumexp}_{a}(\{R_{a},S_{a}\})\})=$ $\displaystyle\textrm{logsumexp}_{a}(\{\textrm{logsumexp}_{a}(\{T_{a},R_{a}\}),S_{a}\}),$ $$

$$ $\displaystyle\max_{a}(\{\max_{a}(\{T_{a},R_{a}\}),S_{a}\})=$ $\displaystyle=\max_{a}(\{T_{a},\max_{a}(\{R_{a},S_{a}\})\}).$ $$

We can prove that this associative property allows these reductions to be performed efficiently in parallel with logarithmic time complexity, provided we have adequately many parallel workers:

###### Theorem 1 .

The time complexity of a reduction operation involving an associative function, such as $\textrm{logsumexp}_{a}$ or $\max_{a}$, over an array of size $N$ using $p$ parallel processors is $O\left(\frac{N}{p}+\log p\right)$.
When the number of processors $p$ is equal to $N$, the time complexity is reduced to $O(\log N)$.

The proof of Theorem 1 is in Appendix 10..
Putting this result together, and for $\hat{a},\hat{b}\in\{1,\cdots,t\}$ intra-chunk indices, we get the following highly parallel Algorithm 1:

Figure: Algorithm 1 Single Query Energy Forward

###### Theorem 1 .

### 5.2.  Efficient parallel decoding

One of the core insights of automatic differentiation is that the gradient of a function $\nabla_{x}f(x)$ can be computed with the same time complexity as computing $f(x)$ Vieira ().
The caveat however is that if the function has a deep computational graph, then the memory footprint of computing the gradient grows with that depth as backpropagation requires storing the values of the intermediate tensors.
In our case, the computational graph involved in computing the energy is shallow and therefore the memory overhead is negligible.
This means that if we can compute the energy efficiently, we obtain an efficient algorithm for computing its gradient (i.e. the self-attention operation) automatically.

In our case, we want to compute the gradient of the energy function with respect to $\zeta_{A}$ and then set it to zero.
This can be done with automatic differentiation engines having set $\zeta$ to be a tensor of zeros from the very outset.
We can however manually implement a gradient with respect to $\zeta$ pass of the above Algorithm  1 that doesn’t materialize $\zeta$ in Algorithm 2 below.
Note in particular that when we set $\zeta_{A}=0,\,\,\forall A\in\{1,\cdots,d_{h}\}$ then $O$ involves only the logsumexp of the dot product between queries and keys.

Figure: Algorithm 2 Single Query Energy Gradient w.r.t $\zeta$

Notice here that by storing $O,m$ for backward, the only remaining reduction operation that needs to be performed is the one in line 5 of the above algorithm. This single reduction takes $O(N/p)$ time to compute the local sums on each device and $\log p$ time to communicate and combine partial results, and therefore we get the same asymptotic complexity as the logsumexp calculation.

In practice, we implement the forward and gradient w.r.t. $\zeta$ in a single function which returns both the value and the gradient of the energy function.
We can therefore put together Algorithms 1 and 2 into the following efficient parallel decoding Algorithm 3:

Figure: Algorithm 3 Tree Decoding

This algorithm requires three Allreduce operations in total, meaning that the required time complexity is $O(3(N/p+\log p))$.

### 5.3.  Implementation and topology-awareness

Figure: Figure 2: NCCL Send/Recv between two H100 GPUs intra-node and inter-node. GPU clusters offer a two-tier topology whereby intra-node bandwidth is significantly higher than inter-node. Algorithms such as Tree Attention exploit this topology by reducing inter-node communication requirements, enabling better overlap of communication with computation.
Refer to caption: fig/nccl-lat-bw.png

Figure: (a) Execution time of 16-head Tree Attention vs Ring Attention for different sizes of GPU cluster (from 1 to 16 H100 DGX nodes). Relative execution times are indexed to the Ring Attention times at a sequence length of 80k tokens.

## 6.  Results

Similar to Ring Attention, Tree Attention is an exact computation of attention.
Since training and evaluation metrics are the same as for attention, our experimental results are focused primarily on latency in section 6.1., peak memory usage in section 6.2. and communication volumes in section 6.3..
Since our algorithm computes numerically identical results as the forward pass of standard attention, our performance results transfer seamlessly to transformer architectures.

We performed all experiments on a DGX H100 cluster consisting of 16 nodes, each containing 8 H100 GPUs. All GPUs within the node are connected via an all-to-all NVLINK 4.0 (900GBps) topology. Nodes are connected to each other via 8 InfiniBand NDR interconnects per node (1 per GPU), each of which provides 400 Gbps (leading to an aggregate 3.2 Tbps node injection bandwidth).

### 6.1.  Latency

In terms of practical usefulness, our study of the energy function brought to light a previously unnoted parallelizability inside the attention computation – that of the reduction of the logsumexp across the sequence dimension, which can be implemented as a parallel Allreduce.
As stated in Theorem 1, it becomes theoretically possible to implement attention, per query as an $N/p+\log(p)$ parallel operations rather than $N$, where the logarithmic term is proportional to the number of devices available for parallelization.
When the attention is sharded across multiple devices, this asymptotic speedup creates a considerable speedup over alternative methods for decoding.

To empirically test the theoretical benefits of our Tree Attention method, we compute latency by measuring the time required to perform decoding for different sequence lengths and varying number of H100 nodes.
We compare Tree Attention to our own Ring Attention execution times in Fig. 4(a).
Both methods use Flash Attention 2 Dao () for the individual-GPU attention computation.
For our experiments, we benchmark on a standard attention block consisting of 16 heads of dimension 128 across different sequence lengths.

Our latency results shows how Tree Attention improves over Ring Attention as we increase the sequence length in Fig. LABEL:fig:relative_times and increase the number of GPUs in Fig. LABEL:fig:absolute_times.
As the plots demonstrate, as we scale the sequence length or the number of GPUs, the gap between Tree Attention and Ring Attention execution time widens *asymptotically*.
Remarkably, Tree Attention achieves close $\times 8$ speedups when we use 128 GPUs on a sequence length of 5.12M.
We expect this trend to continue for larger sequence lengths.

### 6.2.  Memory cost

Figure: (b) Peak memory usage of a single attention block with Tree Attention vs Ring Attention when sharded between two RTX 4090s. Results taken using the JAX memory profiler on one GPU. The difference in peak memory scales with hidden size and sequence length.

### 6.3.  Communication volume

For Ring Attention’s P2P communication strategy, the total volume of data being communicated between devices (in units of number of tensor elements) per iteration is given by:

$$ $V_{ring}=2btd\times p$ (31) $$

where $p$ is the number of devices. The first factor comes from counting the total number of communicated elements corresponding to $\{(\mathbf{k}_{\hat{a}},\mathbf{v}_{\hat{a}}),\hat{a}=1,\cdots,t\}$, i.e.

$$ $\textrm{numel}\left(\{(\mathbf{k}_{\hat{a}},\mathbf{v}_{\hat{a}}),\hat{a}=1,\cdots,t\}\right)=2btd.$ (32) $$

The Allreduce strategy we use in Tree Decoding requires the following volume Anthony *et al.* ():

$$ $V_{\texttt{Allreduce}}=2\times\frac{p-1}{p}\times\textrm{numel}.$ (33) $$

We communicate a shard of the numerator, denominator and max, requiring:

$$ $\textrm{numel}\left(\textrm{n},\textrm{d},m\right)=bd+2bn_{h}.$ (34) $$

Note that we first perform on device the local reductions to obtain the local numerator and denominator on each device which consequently makes it so that $t$, i.e. the size of the local sequence chunk doesn’t appear in the above expression. We then obtain:

$$ $V_{Tree}=2\frac{p-1}{p}\times\left(bd+2bn_{h}\right).$ (35) $$

Our theoretical analysis shows that per iteration our algorithm maintains a lower communication volume than Ring Attention. Note however that Ring Attention when performed in the training setting with many queries overlaps communication and computation so as to hide it’s communication costs.
However, overlapping communication and computation in the decoding case is infeasible because of how fast the attention computation on a single GPU is relative to how long it takes to communicate the chunk of keys and values between two devices.

Concretely, let us take the example of decoding from a context of length 640000 split between $8$ GPUs within one node. Let us take a hidden size of 2048 and fix our data type to be bfloat16. Each device for decoding takes $O(10^{-5})$ seconds to perform
the Flash Attention computation.
The time it takes to move the keys and values of the corresponding size between adjacent GPUs as per Fig. 2 is roughly $O(10^{-3})$ seconds. The latency incurred between nodes is even greater and therefore overlapping isn’t feasible due to this disparity in timescales.

## 7.  Discussion and Conclusion

In this paper, we have derived the energy function for self-attention and demonstrated how the computation of the derivative of this function provides a novel and efficient method for computing attention in parallel.
This advantage is especially apparent when performing decoding across multiple devices, in which case our Tree Attention enables us to substantially outperform Ring Attention with an *asymptotically* superior algorithm, with $\times 8$ speedups when we use 128 GPUs on a sequence length of 5.12M.
We also see that the AllReduce operation that we use involves sending partially reduced objects, which greatly reduces the volume of communicated data as well as the peak memory requirement.

Our introduction of a unique energy function for self-attention develops interesting connections between attention and other related models such as Hopfield networks and the general notion of associative memories.
Given this energy function, it is possible to mathematically analyze its loss landscape and dynamics, as well as begin to understand how the attention operation could potentially be improved.
Additionally, since the energy function can be framed as a variational free energy, it becomes possible to interpret the attention operation through a Bayesian lens and understand the generative model that is implicit in its operation.
We hope that these mathematical derivations aid research in understanding the properties of transformer networks, as a deep understanding of the Hopfield energy function has done for Hopfield networks.

Theoretical asymptotic improvements should be demonstrable in the single-device case, but in practice, this advantage is not realized. SMs in H100 and older architectures communicate through shared global memory, not in a peer-to-peer fashion like individual GPUs in a multi-device setup, making the reduced communication negligible. However, recent experimental instructions in the H100 have enabled peer-to-peer SM communication  NVIDIA (), suggesting that these instructions could lead to speedups over Flash Attention on a single device.

## 8.  Acknowledgement

We would like to thank Skye Wanderman-Milne for answering our questions on JAX multi-device programming and Anna Golubeva for proof-reading our manuscript. We are also grateful to all those who contributed their expertise and support throughout the course of this work, including our colleagues at Zyphra and friends whose encouragement and insightful feedback were invaluable.

## 9.  Appendix: JAX code

Below is the tree_flash_decode method. Our full code base is available here: [https://github.com/Zyphra/tree_attention](https://github.com/Zyphra/tree_attention).

The function uses Flash Attention 2 Dao () to compute the local numerator and denominator, both of which are accumulated between devices using an Allreduce (which is what psum and pmax call). NCCL determines in what pattern these results are communicated.

## 10.  Theorem 1 Proof

We prove theorem 1 below.

###### Proof.

Sequential Case:
On a single GPU, the reduction operation over an array of size $N$ has a time complexity of $O(N)$ since the processor must sequentially process each element.

Parallel Processing with $p$ Processors:
Divide the array of size $N$ into $p$ chunks, each of size $\frac{N}{p}$.
Each processor performs the reduction operation on its chunk independently.
The time complexity for each processor is $O\left(\frac{N}{p}\right)$.

Combining Partial Results:
The partial results from the $p$ processors need to be combined.
Using a tree pattern for reduction, the partial results can be reduced in $O(\log p)$ steps.
Each step involves combining pairs of results, halving the number of results at each step until only one result remains.

Total Time Complexity:
The total time complexity is the sum of the time complexities for processing the chunks and combining the results:

$$ $O\left(\frac{N}{p}\right)+O(\log p).$ $$

This proves that the time complexity of a reduction involving an associative operation over an array of size $N$ is $O\left(\frac{N}{p}+\log p\right)$ when using $p$ parallel processors, and it reduces to $O(\log N)$ when the number of processors is equal to the size of the array.
∎

###### Proof.

## 11.  Computing safe softmax

While, mathematically, attention utilizes the softmax operation, in practice this is often numerically unstable using relatively low precision operations.
To address this, a mathematically equivalent function, the ‘safe softmax’ is instead used which subtracts all dot products in the exponential by the max. This ensures that all values being exponentiated are less than 1 and hence less likely to explode and cause numerical instability.
Here, we demonstrate that our energy function approach also can account for safe softmax.

Let us suppose we compare our generating function

$$ $F_{tot}=\sum_{i}\log\sum_{a=1}^{i}\exp\left(q_{i}\cdot k_{a}^{T}+\zeta_{a}\cdot v^{T}_{a}\right)$ (36) $$

and a slightly modified one:

$$ $F^{\prime}_{tot}=\sum_{i}\log\sum_{a=1}^{i}\exp\left(q_{i}\cdot k_{a}^{T}+\zeta_{i}\cdot v^{T}_{a}-m_{i}\right).$ (37) $$

When we take the derivative of these two quantities, we see that we get the same result:

$$ $\frac{\partial F_{tot}}{\partial\zeta_{i}}\bigg|_{\zeta_{i}=0}=\frac{\partial F^{\prime}_{tot}}{\partial\zeta_{i}}\bigg|_{\zeta_{i}=0}.$ (38) $$

To see it explicitly:

$$ $\frac{\partial F^{\prime}_{tot}}{\partial\zeta_{i}}\bigg|_{\zeta_{i}=0}=\frac{\sum^{i}_{a=1}\exp(q_{i}\cdot k^{T}_{a}-m_{i})v_{a}}{\sum^{i}_{a=1}\exp(q_{i}\cdot k^{T}_{a}-m_{i})}$ (39) $$

$$ $=\frac{\sum^{i}_{a=1}\exp(q_{i}\cdot k^{T}_{a})v_{a}}{\sum^{i}_{a=1}\exp(q_{i}\cdot k^{T}_{a})}.$ $$

Normally, when computing the softmax in an online fashion, this procedure is performed where $m_{i}$ is the row max of $q\cdot k^{T}$.
This shift makes it so that the sum of exponentials doesn’t lead to overflows.

## 12.  Notations for equations

Here is a summary of the various variables and indices that will be used in the coming sections:

$$ $\displaystyle\begin{array}[]{l|ll}\hline\cr x&\text{Attention Input}\\ q,k,v&\text{Query, key and value vectors }\\ \Gamma&\text{Attention Log-likelihood}\\ \zeta&\text{Source vector }\\ m&\text{Max of $q\cdot k^{T}$ }\\ Z&\text{Partition function }\\ z&\text{Activation vector }\\ \textrm{n}&\text{Attention numerator}\\ \textrm{d}&\text{Attention denominator}\\ \textrm{lse}&\textrm{Attention score logsumexp}\\ F&\text{Generating function }\\ P&\text{Attention score probability density}&\\ \end{array}$ (40) TABLE I: Variable names. $$

$$ $\displaystyle\begin{array}[]{l|ll}\hline\cr N&\text{Sequence length}\\ d&\text{Embedding dimension}\\ d_{h}&\text{Head dimension}\\ p&\text{Number of devices}\\ t&\text{Chunk size N/p}\\ b&\text{Batch size}\\ a,i,j\in\{1,\cdots,N\}&\text{Sequence Indices}\\ A,B\in\{1,\cdots,d\}&\text{Embedding indices}\\ \bar{A},\bar{B}\in\{1,\cdots,d_{h}\}&\text{Intra-head indices }\\ h\in\{1,\cdots,n_{h}\}&\text{Head indices }\\ \hat{a},\hat{b}\in\{1,\cdots,t\}&\text{Intra chunk indices }\\ \end{array}$ (41) TABLE II: Index names and ranges. $$

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: I will now present some comments and questions to the authors, aimed exclusively at improving their interesting work about a more effective computation of self-attention in transformers. Comments and questions are listed here below, divided in Major issues, Minor issues and Typos.

Major issues:

1) From Theorem (1), the computational complexity of self-attention, for a single query, is $O( \frac{N}{p} + \log p )$  where $N$ is the full length of the sequence and $p$ is the number of processors. The total complexity is thus quadratic in $N$. There are examples of methods that simplify the total complexity to $O(N)$ (e.g. BigBird, Longformer) or $O(N \log N)$ (e.g. Reformer) which, by the way, do not rely on a modification of the processors disposition and are not necessarily exact. If we assume the number of GPUs to scale as $p = \alpha N$ where $\alpha > 0$, then the complexity per query would scale as $O(\log N )$, hence the total complexity as $O(N\log N)$ which equals the one obtained by Reformer. Nevertheless, even if this scenario appears difficult to reproduce, since $N$ is large by definition, this limit is mentioned in the text (row 269). Could the authors comment on this? 

2) Regarding the experiments reported in Section 5.1. Physical limitations seem to impose that $p = O(1)$, and $p \ll N$. As a consequence, the effective behaviour of the execution time $T$ with respect to $N$ is $T \simeq \frac{N}{p}$, hence it is linear in $N$. On the other hand, when $N$ is fixed, we have $T \simeq  \frac{N}{p} + log(p)$, it should then scale logarithmically with $p$ for large $p$. Can the authors show these trends more clearly from figure (3) e.g. increasing the number of points (in the intermediate range, i.e. between the current minimum and maximum values on the x-axis) and performing a proper fit of the data-points? Moreover: what is the expected scaling of the computational cost of Ring Attention as a function of $(N,p)$ ? It seems from figure (3,b) that the behaviour of $T_{ring}$ as a function of $p$ is linear, while in principle it should be constant and remain $T_{ring} \gg \log p$. If not, please provide for a theoretical argument to explain this behaviour.

3)May the authors show that the formulation of $o$ in line (2) of Algorithm (3) already normalized by $lse$ is equivalent to implementing the formulation from line (4) in Algo (2)? Expanding the mathematical passages in Algo (3) might useful to visualize this aspect.  

4) The difference between the memory costs of Ring and Tree attentions is $\delta_{M} = Mem_{ring} - Mem_{tree} = 2b(\frac{N}{p}\cdot d - n_{h})$. Could the authors check for the scaling of this quantity with respect to $(N,p)$ to satisfy $\delta_{M} \propto \frac{N}{p}$ ? I believe this would make the analysis more rigorous and reliable. 

5) I propose to the authors to add the most recent attempt to map the self-attention mechanism into a energy-based model, i.e. D’Amico et al. 2024, to the Related Work appendix. Nevertheless, i would spend much more space in the same paragraph to mention previous works in the matter of the computational complexity of computing self-attention, even the cases where multiple processors were not employed. A useful review about this problem can be found in Keles et al. 2023.  

Minor Issues:

6) I suggest the authors to substitute the Wikipedia reference for the definition of Helmholtz free-energy with a more reliable and solid scientific source, among the thousand great statistical mechanics books. One example is  Landau, Lifshitz, "Statistical Physics". 

7) The index A in the second formula in section 2 (i suggest to add numbers to all formulas) seem to be redundant, including the subsequent comment "The capital Latin indices (𝐴) span the hidden dimension $𝑑_ℎ$". I would simply substitute the letter "d" on top of the sum with $d_h$. 

8) What is the difference between "Absolute" and "Relative" execution time in Section 5.1?  

9) Can the authors add a definition for the batches $b$ in Section 5.2 ?

Typos:

10) Row 144: Cumulant.
11) Row 505: Its communication.
soundness: 2
strengths: The main strength of this work lies in the elegant simplicity of the energy-driven theoretical argument, which enables a simplification of the architecture used to compute self-attention coefficients. The paper is well-written and pedagogical, even for readers who are not experts in physics or in the technical aspects of executing self-attention computation on this type of machine.
summary: The manuscript "Tree Attention: Topology-Aware Decoding for Long-Context Attention" proposes a novel strategy for computing self-attention coefficients in Transformers, the fundamental operation behind the functioning of this highly celebrated generative model. The method is based on two fundamental ideas. The first is theoretical, inspired by statistical mechanics: at each step of the data generation process, self-attention can be computed as the gradient of a function that resembles the free energy of a multi-particle system at thermodynamic equilibrium. The second idea, on the other hand, concerns hardware: an ensemble of GPUs can be employed to compute the energy function along a tree-like scheme. Consequently, self-attention is computed in such a way that inter-node communication is avoided, unlike in recent architectures presented in the literature.
The main result of the paper is an algorithm for computing self-attention that is faster, as it scales logarithmically with the number of GPU nodes used in the process; requires less memory for storing useful variables before usage; and reduces inter-node communication volume for shorter execution times.
weaknesses: I believe that the authors should make a stronger effort in validating the theoretical predictions about the improvement in the computational cost of self-attention proposed by the Tree Attention method, which is also their main scientific question. The scaling of the execution time $T$ with respect to variables $N$ and $p$ is reported in figure (3), without any proper experimental analysis of the results, which would effectively validate the theoretical predictions and convince possible future users. While their report only shows that Tree Attention is faster than Ring Attention, a deeper analysis of the experiments would be definitely needed: e.g. one might start from a linear fit of $T$ Vs. $N$ in linear scale and a linear fit of $T$ Vs. $p$ in log scale; a proper expected scaling (to be confirmed by the plots) of Ring Attention is never provided, while the condition $T_{ring} > T_{tree}$ might give useful bounds on the choice of the parameters of the system.  
Moreover, while the paper emphasizes a strong connection with modern Hopfield models (which are even mentioned in the abstract), this aspect is never properly addressed in the manuscript. As a matter of fact, the only—yet fundamental, for the sake of the results—physics-inspired part of the work consists in deriving self-attention by applying the logsumexp function, which satisfies the associative property. On the other hand, self-attention has emerged in previous papers (see: Hoover et al. 2023, Kozachkov et al. 2023, D’Amico et al. 2024) due to its analogy with the gradient descent dynamics over the energy function of the modern Hopfield model. Since both the dynamic aspect/analogy part and the thermodynamic one, which involves considerations about minimization of a free-energy function, are only briefly discussed, I would consider sacrificing some of the discussion about Hopfield models to give more room to a broader test of the analytical predictions.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. What is $d_h$ in line 101? Also there seems to be some notation abuse in the expression since it both sums over $A$ and also is apparently performed for all $A$. I think one of those should be 'a'.

2. How exactly does tree attention save inter node communications and exploit the two-level network topology? Is it by reducing the number of inter-node communications or the data volume of inter-node communications or both?

3. Why must $k_{\hat{a}'}$ , $v_{\hat{a}'} $ from neighboring devices be stored in ring attention (Page 9, lines 451-454)? Since a copy of $q$ is available at all nodes, won't it be possible to perform ring attention by just computing the flash attention computations at each node to its neighbors which can then rescale and update their attention values? This would reduce both the memory overhead as well as the communication volume ('$t$' factor in (28)) of ring attention
soundness: 3
strengths: 1. The proposed algorithm is a novel take on self-attention and the derivation of the energy function and the associated Bayesian interpretation has the potential to contribute to a better theoretical understanding of the efficacy of self-attention in the future.

2. The tree attention algorithm cleverly exploits associativity of $\log\sum\exp$ and $\max$ to create a tree structure for attention computation in decoding that achieves logarithmic complexity as opposed to the linear complexity of prior approaches like ring attention which leads to clear latency improvements in experiments.
summary: The paper introduces Tree Attention, an algorithm for parallel self-attention computation across multiple GPUs to reduce communication volume and memory usage. It derives a scalar energy function for self-attention which leads to a Bayesian interpretation and a tree reduction approach for computing self-attention that is asymptotically faster due to parallelization. Experiments show close to 8× speedups on sequences of 5.12 million tokens using 128 GPUs.
weaknesses: The details of the derivation of the energy function (for e.g. with the Taylor expansion in (10)) as well as the Bayesian interpretation in 3.2 seem unnecessary and may confuse the reader. If the main goal is to introduce the tree attention algorithm and its advantages, I would recommend making Section 3 more compact and moving things like the Bayesian interpretation to the discussion at the end of the paper.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Overall, I think this paper proposed a novel self-attention mechanism that scales logarithmically and reduces memory footprints. I have a few comments and questions as follow.

In section 3, the derivation of energy function for self-attention is interesting. What will be comparative analysis between tree attention's use of energy function and other attention mechanism would use this energy function?

As stated in the cons, the paper showed an impressive result compare to ring-based attention. What would be the testing results compare to other different attention mechanism with more aspects like (different hidden sizes, varying GPU numbers)? Also, what will be the memory usage compare to other methods? 

For large distributed clusters (cloud-based GPUs), communication latency is a major concern, and overlapping communication and computation can be tricky. More insights on how Tree Attention manages or mitigates these latencies would provide better clarity for real-world implementations.

The analysis of communication volume is informative, but more practical scenarios could be discussed. For example, how does the communication overhead change when increasing not only the number of GPUs but also the number of nodes?
soundness: 2
strengths: 1. The proposed method can effectively scale with the number of GPUs and scales logarithmically.

2. The memory usage of the proposed method can significantly reduce the memory usage.

3. The paper innovatively improve self-attention mechanism by using other related field's knowledge, which is interesting.
summary: This paper proposes a new method for improving the performance of self-attention mechanisms used in transformers, especially for long-context sequences. The authors introduce a tree-based method for computing attention, which uses the internal parallelism in the attention operation. This method reduces communication overhead and memory usage across multi-GPU clusters, as compared to existing methods like Ring Attention. The paper establishes a theoretical basis for their approach by deriving an energy function for self-attention and showing how it can be efficiently computed using parallel reductions. Experimental results indicate that Tree Attention achieves up to 8× speedups in long-sequence attention tasks while also halving peak memory usage.
weaknesses: 1. While the paper demonstrates high theoretical and experimental results, it did not provide any real-world deployment or compare with existing frameworks.

2. The system only tests on DGX H100 system, limiting the generalizability to other systems.

3. The paper only compare with self-implemented ring attention, it would be beneficial to compare with more other attention mechanisms like Linformer, Longformer, or performer.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: - What design aspects of Tree Attention allow it to remain exact, despite the use of parallelized operations? Does the tree reduction introduce any approximations, or is it truly an exact method? This clarification would be valuable for understanding its applicability in tasks where exact attention computation is critical.
- The paper reports an 8x speedup compared to Ring Attention. Could the authors break down what aspects of Tree Attention contribute to this speedup? For example, how much of the improvement is due to, e.g., reduced communication volume by tree reduction?
- Given that Ring Attention is already a sequence parallel method for multi-GPU attention parallelism, what specific limitations does Tree Attention address that Ring Attention?
soundness: 2
strengths: The Tree Attention approach accounts for two-level network topology awareness within GPU clusters, which is a practical and well-considered enhancement for real-world distributed computing environments.
summary: This paper addresses the challenge of efficiently computing self-attention in transformer models, which traditionally incurs quadratic complexity in sequence length. The authors propose an approach for parallelizing attention computation using Tree Attention, which leverages the gradient of an energy function. By implementing a tree reduction communication strategy, the proposed approach enables faster parallel decoding and reportedly achieves up to 8x faster performance than Ring Attention, with reduced communication volume and half the memory needed.
weaknesses: - The paper does not clearly explain why using the gradient of the energy function is advantageous for self-attention. This aspect feels insufficiently motivated, making it unclear on the potential contribution. The authors could have strengthened their argument by explaining how this approach uniquely benefits attention computation compared to existing methods.
- The paper claims significant performance improvements over Ring Attention, but it lacks clarity on what specific limitations in Ring Attention are addressed by Tree Attention. Without a clear comparison, it is challenging to understand if Tree Attention fills an existing gap or if it simply provides an alternative approach. 
- The details of the Tree Attention approach are challenging to follow, potentially impacting its accessibility and reproducibility. The method relies on parallel decoding and a tree reduction for attention computation, but the details of these operations are not well explained. Clearer explanations or diagrams could be helpful.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: Please address the issues mentioned in the weaknesses section.
soundness: 2
strengths: 1. This paper includes a microbenchmark of latency and bandwidth between inter-node and intra-node H100 GPUs, highlighting the communication bottleneck.

2. The paper introduces a new approach to designing improved attention mechanisms by considering the system topologies.
summary: This paper proposes a new attention mechanism called tree attention. By leveraging tree reduction, it enables high parallelization of the decoding stage during LLM inference, especially for extremely long contexts.
weaknesses: 1. Although this work is called "tree attention," it is not general enough to handle common attention operations. It can only be used when the query has a single token, which limits its usefulness during training and the prefill stage.

2. The implementation of the ring attention baseline might be incorrect. The authors directly use the ring_fwd function from the flash_attn_jax library, which assumes the query, key, and value matrices are already partitioned [1]. Consequently, the actual calculation in this paper's implementation [2] involves a query matrix with P tokens, where P represents the number of workers.

3. Ring attention is not an appropriate baseline for the scenario considered in this work. The purpose of ring attention is to address long-context training, where the query, key, and value matrices are partitioned across workers. To calculate the complete attention map, these matrices need to be communicated, and ring attention provides a better scheduling mechanism to overlap communication and computation. However, in the scenario discussed here—during the decoding stage—the query has only one token and is copied to all workers. Therefore, only the last row of the attention map needs to be calculated, making ring attention unnecessary.

4. The idea and algorithm presented are straightforward, but the authors appear to complicate the explanation, which adds confusion rather than clarity.

5. The evaluation only demonstrates a single attention module and lacks end-to-end evaluation on a real-world large language model (LLM), which could weaken the authors’ claims about the algorithm’s effectiveness.

6. The algorithm and code do not utilize a tree topology. If they merely rely on the NCCL library, this approach may lack generalizability, reducing this work's contribution.

[1] https://github.com/nshepperd/flash_attn_jax/blob/eab32ee0725f1de66467f35ecd7f6acd15dd3f1f/tests/test_ring.py#L76

[2] https://anonymous.4open.science/r/tree_attention-7C32/ring_shard_test.py

Meta-review: The paper proposes a new method (Tree Attention) for parallel self-attention computation across multiple GPUs. This leads to a reduction in communication overhead and memory usage, as compared to existing approaches such as Ring Attention. The idea is to derive a scalar energy function for self-attention, which is computed efficiently due to parallelization. 

The reviewers appreciated the elegant simplicity of the theoretical argument, the reduction in memory usage and the parallelism of the proposed approach. However, during the review process a number of issues were raised concerning the performance of the method, its evaluation and using ring attention as a baseline. These concerns were not addressed in a satisfactory way during the rebuttal process. In the final discussion between AC and reviewers, reviewers Eqp7, aspw and mxBy agreed that the paper, in its current form, does not provide sufficient evidence of a significant improvement over the state of the art. I agree with this assessment and therefore recommend rejection at this stage. 

The reviewers did praise the novelty of the approach, which I also find valuable. I would thus encourage the authors to resubmit an improved version that addresses the concerns raised during this review process to a future venue.

Decision: Reject