AdaSplash: Adaptive Sparse Flash Attention

• URL: http://arxiv.org/abs/2502.12082v1
• Date: Mon, 17 Feb 2025 17:56:23 GMT
• Title: AdaSplash: Adaptive Sparse Flash Attention
• Authors: Nuno Gonçalves, Marcos Treviso, André F. T. Martins,
• Abstract summary: We propose AdaSplash, which combines the efficiency of GPU-optimized algorithms with the sparsity benefits of $alpha$-entmax.<n>AdaSplash achieves substantial improvements in runtime and memory efficiency compared to existing $alpha$-entmax implementations.
• Score: 20.28859850361068
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: The computational cost of softmax-based attention in transformers limits their applicability to long-context tasks. Adaptive sparsity, of which $\alpha$-entmax attention is an example, offers a flexible data-dependent alternative, but existing implementations are inefficient and do not leverage the sparsity to obtain runtime and memory gains. In this work, we propose AdaSplash, which combines the efficiency of GPU-optimized algorithms with the sparsity benefits of $\alpha$-entmax. We first introduce a hybrid Halley-bisection algorithm, resulting in a 7-fold reduction in the number of iterations needed to compute the $\alpha$-entmax transformation. Then, we implement custom Triton kernels to efficiently handle adaptive sparsity. Experiments with RoBERTa and ModernBERT for text classification and single-vector retrieval, along with GPT-2 for language modeling, show that our method achieves substantial improvements in runtime and memory efficiency compared to existing $\alpha$-entmax implementations. It approaches -- and in some cases surpasses -- the efficiency of highly optimized softmax implementations like FlashAttention-2, enabling long-context training while maintaining strong task performance.

Related papers

• Second-order Optimization of Gaussian Splats with Importance Sampling [51.95046424364725]
  3D Gaussian Splatting (3DGS) is widely used for novel view rendering due to its high quality and fast inference time. We propose a novel second-order optimization strategy based on Levenberg-Marquardt (LM) and Conjugate Gradient (CG) Our method achieves a $3times$ speedup over standard LM and outperforms Adam by $6times$ when the Gaussian count is low.
  arXiv  Detail & Related papers  (2025-04-17T12:52:08Z)
• VEXP: A Low-Cost RISC-V ISA Extension for Accelerated Softmax Computation in Transformers [13.984340807378457]
  Accelerating Softmax is challenging due to its non-pointwise, non-linear nature, with exponentiation as the most demanding step. We design a custom arithmetic block for Bfloat16 exponentiation leveraging a novel approximation algorithm based on Schraudolph's method. We execute Softmax with 162.7$times$ less latency and 74.3$times$ less energy compared to the baseline cluster.
  arXiv  Detail & Related papers  (2025-04-15T14:28:48Z)
• Simultaneous Computation and Memory Efficient Zeroth-Order Optimizer for Fine-Tuning Large Language Models [33.911521719528686]
  Fine-tuning is powerful for adapting large language models to downstream tasks, but it often results in huge memory usages. A promising approach is using Zeroth-Order (ZO) gradients, which estimates to replace First-Order (FO) gradients. We introduce a novel layer-wise sparse computation and memory efficient ZO, named LeZO.
  arXiv  Detail & Related papers  (2024-10-13T12:47:37Z)
• CAME: Confidence-guided Adaptive Memory Efficient Optimization [20.009302737137787]
  Adaptive gradient methods have demonstrated excellent performance in the training of large language models. The need for maintaining second-moment estimates requires maintaining a high cost of extra memory overheads. Several memory-efficients have been proposed to obtain a drastic reduction in auxiliary memory usage, but with a performance penalty. We propose CAME to simultaneously achieve two goals: fast convergence as in traditional adaptive methods, and low memory usage as in memory-efficient methods.
  arXiv  Detail & Related papers  (2023-07-05T06:05:36Z)
• Accelerated First-Order Optimization under Nonlinear Constraints [73.2273449996098]
  We exploit between first-order algorithms for constrained optimization and non-smooth systems to design a new class of accelerated first-order algorithms. An important property of these algorithms is that constraints are expressed in terms of velocities instead of sparse variables.
  arXiv  Detail & Related papers  (2023-02-01T08:50:48Z)
• Softmax-free Linear Transformers [90.83157268265654]
  Vision transformers (ViTs) have pushed the state-of-the-art for visual perception tasks. Existing methods are either theoretically flawed or empirically ineffective for visual recognition. We propose a family of Softmax-Free Transformers (SOFT)
  arXiv  Detail & Related papers  (2022-07-05T03:08:27Z)
• Two-step Lookahead Bayesian Optimization with Inequality Constraints [21.703234193908038]
  We propose a two-step lookahead constrained Bayesian optimization acquisition function (2-OPT-C) supporting both sequential and batch settings. In numerical experiments, 2-OPT-C typically improves query efficiency by 2x or more over previous methods, and in some cases by 10x or more.
  arXiv  Detail & Related papers  (2021-12-06T07:40:54Z)
• Adapting to Misspecification in Contextual Bandits [82.55565343668246]
  We introduce a new family of oracle-efficient algorithms for $varepsilon$-misspecified contextual bandits. We obtain the first algorithm that achieves the optimal $O(dsqrtT + varepsilonsqrtdT)$ regret bound for unknown misspecification level.
  arXiv  Detail & Related papers  (2021-07-12T21:30:41Z)
• Minimax Optimization with Smooth Algorithmic Adversaries [59.47122537182611]
  We propose a new algorithm for the min-player against smooth algorithms deployed by an adversary. Our algorithm is guaranteed to make monotonic progress having no limit cycles, and to find an appropriate number of gradient ascents.
  arXiv  Detail & Related papers  (2021-06-02T22:03:36Z)
• Kernel methods through the roof: handling billions of points efficiently [94.31450736250918]
  Kernel methods provide an elegant and principled approach to nonparametric learning, but so far could hardly be used in large scale problems. Recent advances have shown the benefits of a number of algorithmic ideas, for example combining optimization, numerical linear algebra and random projections. Here, we push these efforts further to develop and test a solver that takes full advantage of GPU hardware.
  arXiv  Detail & Related papers  (2020-06-18T08:16:25Z)

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.