Efficiently Dispatching Flash Attention For Partially Filled Attention Masks

• URL: http://arxiv.org/abs/2409.15097v2
• Date: Tue, 24 Sep 2024 12:56:13 GMT
• Title: Efficiently Dispatching Flash Attention For Partially Filled Attention Masks
• Authors: Agniv Sharma, Jonas Geiping,
• Abstract summary: Transformers are widely used across various applications, many of which yield sparse or partially filled attention matrices. We introduce Binary Block Masking, a highly efficient modification that enhances Flash Attention by making it mask-aware. Our experiments on attention masks derived from real-world scenarios demonstrate up to a 9x runtime improvement.
• Score: 29.36452085947087
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Transformers are widely used across various applications, many of which yield sparse or partially filled attention matrices. Examples include attention masks designed to reduce the quadratic complexity of attention, sequence packing techniques, and recent innovations like tree masking for fast validation in MEDUSA. Despite the inherent sparsity in these matrices, the state-of-the-art algorithm Flash Attention still processes them with quadratic complexity as though they were dense. In this paper, we introduce Binary Block Masking, a highly efficient modification that enhances Flash Attention by making it mask-aware. We further propose two optimizations: one tailored for masks with contiguous non-zero patterns and another for extremely sparse masks. Our experiments on attention masks derived from real-world scenarios demonstrate up to a 9x runtime improvement. The implementation will be publicly released to foster further research and application.

Related papers

• FlashMask: Efficient and Rich Mask Extension of FlashAttention [22.810595298076866]
  FlashMask is an extension of FlashAttention that introduces a column-wise sparse representation of attention masks. By adopting this novel representation, FlashMask achieves linear memory complexity $O(N)$, suitable for modeling long-context sequences. We evaluate FlashMask's performance in fine-tuning and alignment training of LLMs such as SFT, LoRA, DPO, and RM.
  arXiv  Detail & Related papers  (2024-10-02T09:17:26Z)
• ColorMAE: Exploring data-independent masking strategies in Masked AutoEncoders [53.3185750528969]
  Masked AutoEncoders (MAE) have emerged as a robust self-supervised framework. We introduce a data-independent method, termed ColorMAE, which generates different binary mask patterns by filtering random noise. We demonstrate our strategy's superiority in downstream tasks compared to random masking.
  arXiv  Detail & Related papers  (2024-07-17T22:04:00Z)
• Downstream Task Guided Masking Learning in Masked Autoencoders Using Multi-Level Optimization [42.82742477950748]
  Masked Autoencoder (MAE) is a notable method for self-supervised pretraining in visual representation learning. We introduce the Multi-level Optimized Mask Autoencoder (MLO-MAE), a novel framework that learns an optimal masking strategy during pretraining. Our experimental findings highlight MLO-MAE's significant advancements in visual representation learning.
  arXiv  Detail & Related papers  (2024-02-28T07:37:26Z)
• DynaMask: Dynamic Mask Selection for Instance Segmentation [21.50329070835023]
  We develop a Mask Switch Module (MSM) with negligible computational cost to select the most suitable mask resolution for each instance. The proposed method, namely DynaMask, brings consistent and noticeable performance improvements over other state-of-the-arts at a moderate computation overhead.
  arXiv  Detail & Related papers  (2023-03-14T13:01:25Z)
• MP-Former: Mask-Piloted Transformer for Image Segmentation [16.620469868310288]
  Mask2Former suffers from inconsistent mask predictions between decoder layers. We propose a mask-piloted training approach, which feeds noised ground-truth masks in masked-attention and trains the model to reconstruct the original ones.
  arXiv  Detail & Related papers  (2023-03-13T17:57:59Z)
• Bi-directional Masks for Efficient N:M Sparse Training [64.9617631724811]
  We present a novel method of Bi-directional Masks (Bi-Mask) with its two central innovations. It disentangles the forward and backward weight sparsity and overcomes the very dense gradient. Compared with existing uni-directional scenario that applies a transposable mask and enables backward acceleration, our Bi-Mask is experimentally demonstrated to be more superior in performance.
  arXiv  Detail & Related papers  (2023-02-13T02:32:02Z)
• Mask Transfiner for High-Quality Instance Segmentation [95.74244714914052]
  We present Mask Transfiner for high-quality and efficient instance segmentation. Our approach only processes detected error-prone tree nodes and self-corrects their errors in parallel. Our code and trained models will be available at http://vis.xyz/pub/transfiner.
  arXiv  Detail & Related papers  (2021-11-26T18:58:22Z)
• Image Inpainting by End-to-End Cascaded Refinement with Mask Awareness [66.55719330810547]
  Inpainting arbitrary missing regions is challenging because learning valid features for various masked regions is nontrivial. We propose a novel mask-aware inpainting solution that learns multi-scale features for missing regions in the encoding phase. Our framework is validated both quantitatively and qualitatively via extensive experiments on three public datasets.
  arXiv  Detail & Related papers  (2021-04-28T13:17:47Z)
• DCT-Mask: Discrete Cosine Transform Mask Representation for Instance Segmentation [50.70679435176346]
  We propose a new mask representation by applying the discrete cosine transform(DCT) to encode the high-resolution binary grid mask into a compact vector. Our method, termed DCT-Mask, could be easily integrated into most pixel-based instance segmentation methods.
  arXiv  Detail & Related papers  (2020-11-19T15:00:21Z)
• Ternary Feature Masks: zero-forgetting for task-incremental learning [68.34518408920661]
  We propose an approach without any forgetting to continual learning for the task-aware regime. By using ternary masks we can upgrade a model to new tasks, reusing knowledge from previous tasks while not forgetting anything about them. Our method outperforms current state-of-the-art while reducing memory overhead in comparison to weight-based approaches.
  arXiv  Detail & Related papers  (2020-01-23T18:08:37Z)

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.