Tree Attention: Topology-aware Decoding for Long-Context Attention on GPU clusters

• URL: http://arxiv.org/abs/2408.04093v3
• Date: Wed, 14 Aug 2024 12:47:31 GMT
• Title: Tree Attention: Topology-aware Decoding for Long-Context Attention on GPU clusters
• Authors: Vasudev Shyam, Jonathan Pilault, Emily Shepperd, Quentin Anthony, Beren Millidge,
• Abstract summary: Self-attention is a significant computational bottleneck due to its complexity in the sequence length. In this work, we derive the scalar energy function whose gradient computes the self-attention block. Our formulation reveals that the reduction across the sequence axis can be efficiently computed in parallel through a tree reduction.
• Score: 10.403248386029407
• License: http://creativecommons.org/licenses/by/4.0/
• 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 8x faster in our experiments) than alternative approaches such as Ring Attention, while also requiring significantly less communication volume and incurring 2x less peak memory. Our code is publicly available here: \url{https://github.com/Zyphra/tree_attention}.

Related papers

• Lean Attention: Hardware-Aware Scalable Attention Mechanism for the Decode-Phase of Transformers [4.674454841332859]
  Transformer-based models have emerged as one of the most widely used architectures for natural language processing. These huge models are memory hungry and incur significant inference latency even on cutting edge AI-accelerators. We propose LeanAttention, a scalable technique of computing self-attention for the token-generation phase.
  arXiv  Detail & Related papers  (2024-05-17T00:52:39Z)
• Simulation of Graph Algorithms with Looped Transformers [6.0465914748433915]
  We study the ability of transformer networks to simulate algorithms on graphs from a theoretical perspective. We prove by construction that this architecture can simulate individual algorithms such as Dijkstra's shortest path. We show a Turing Completeness result with constant width when the extra attention heads are utilized.
  arXiv  Detail & Related papers  (2024-02-02T02:48:03Z)
• Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models [20.78813311569383]
  We present Lightning Attention, the first linear attention implementation that enables linear attention to realize its theoretical computational benefits. Specifically, we utilize the conventional attention mechanism for the intra-blocks and apply linear attention kernel tricks for the inter-blocks. Various experiments are conducted on different model sizes and sequence lengths.
  arXiv  Detail & Related papers  (2024-01-09T16:27:28Z)
• Memory-aware Scheduling for Complex Wired Networks with Iterative Graph Optimization [4.614780125575351]
  We propose an efficient memory-aware scheduling framework based on iterative graph optimization. Our framework features an iterative graph fusion algorithm that simplifies the graph while preserving the scheduling optimality.
  arXiv  Detail & Related papers  (2023-08-26T14:52:02Z)
• UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation [93.88170217725805]
  We propose a 3D medical image segmentation approach, named UNETR++, that offers both high-quality segmentation masks as well as efficiency in terms of parameters, compute cost, and inference speed. The core of our design is the introduction of a novel efficient paired attention (EPA) block that efficiently learns spatial and channel-wise discriminative features. Our evaluations on five benchmarks, Synapse, BTCV, ACDC, BRaTs, and Decathlon-Lung, reveal the effectiveness of our contributions in terms of both efficiency and accuracy.
  arXiv  Detail & Related papers  (2022-12-08T18:59:57Z)
• Treeformer: Dense Gradient Trees for Efficient Attention Computation [24.045251327736814]
  We show how to speed up attention computation by enforcing different attention structures such as sparsity, low-rank, approximating attention using kernels. Based on such hierarchical navigation, we design Treeformer which can use one of two efficient attention layers -- TF-Attention and TC-Attention. We demonstrate that our Treeformer architecture can be almost as accurate as baseline Transformer while using 30x lesser FLOPs in the attention layer.
  arXiv  Detail & Related papers  (2022-08-18T18:31:40Z)
• A Unified Framework for Implicit Sinkhorn Differentiation [58.56866763433335]
  We propose an algorithm that obtains analytical gradients of a Sinkhorn layer via implicit differentiation. We show that it is computationally more efficient, particularly when resources like GPU memory are scarce.
  arXiv  Detail & Related papers  (2022-05-13T14:45:31Z)
• Instant Neural Graphics Primitives with a Multiresolution Hash Encoding [67.33850633281803]
  We present a versatile new input encoding that permits the use of a smaller network without sacrificing quality. A small neural network is augmented by a multiresolution hash table of trainable feature vectors whose values are optimized through a gradient descent. We achieve a combined speed of several orders of magnitude, enabling training of high-quality neural graphics primitives in a matter of seconds.
  arXiv  Detail & Related papers  (2022-01-16T07:22:47Z)
• Combiner: Full Attention Transformer with Sparse Computation Cost [142.10203598824964]
  We propose Combiner, which provides full attention capability in each attention head while maintaining low computation complexity. We show that most sparse attention patterns used in existing sparse transformers are able to inspire the design of such factorization for full attention. An experimental evaluation on both autoregressive and bidirectional sequence tasks demonstrates the effectiveness of this approach.
  arXiv  Detail & Related papers  (2021-07-12T22:43:11Z)
• VersaGNN: a Versatile accelerator for Graph neural networks [81.1667080640009]
  We propose textitVersaGNN, an ultra-efficient, systolic-array-based versatile hardware accelerator. textitVersaGNN achieves on average 3712$times$ speedup with 1301.25$times$ energy reduction on CPU, and 35.4$times$ speedup with 17.66$times$ energy reduction on GPU.
  arXiv  Detail & Related papers  (2021-05-04T04:10:48Z)
• Fast Graph Attention Networks Using Effective Resistance Based Graph Sparsification [70.50751397870972]
  FastGAT is a method to make attention based GNNs lightweight by using spectral sparsification to generate an optimal pruning of the input graph. We experimentally evaluate FastGAT on several large real world graph datasets for node classification tasks.
  arXiv  Detail & Related papers  (2020-06-15T22:07:54Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.