Is Flash Attention Stable?

• URL: http://arxiv.org/abs/2405.02803v1
• Date: Sun, 5 May 2024 03:25:25 GMT
• Title: Is Flash Attention Stable?
• Abstract summary: Many organizations training state-of-the-art Generative AI models have reported isolated cases of instability during training. We develop a principled approach to understanding the effects of numeric deviation, and construct proxies to put observations into context. We find that Flash Attention sees roughly an order magnitude of more numeric deviation as compared to Baseline Attention at BF16 when measured during a forward pass.
• Score: 12.827526286642282
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Training large-scale machine learning models poses distinct system challenges, given both the size and complexity of today's workloads. Recently, many organizations training state-of-the-art Generative AI models have reported cases of instability during training, often taking the form of loss spikes. Numeric deviation has emerged as a potential cause of this training instability, although quantifying this is especially challenging given the costly nature of training runs. In this work, we develop a principled approach to understanding the effects of numeric deviation, and construct proxies to put observations into context when downstream effects are difficult to quantify. As a case study, we apply this framework to analyze the widely-adopted Flash Attention optimization. We find that Flash Attention sees roughly an order of magnitude more numeric deviation as compared to Baseline Attention at BF16 when measured during an isolated forward pass. We then use a data-driven analysis based on the Wasserstein Distance to provide upper bounds on how this numeric deviation impacts model weights during training, finding that the numerical deviation present in Flash Attention is 2-5 times less significant than low-precision training.

Related papers

• Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention [37.00916406627862]
  This paper provides the first mechanistic explanation for a long-standing and unresolved failure case where training with flash attention in low-precision settings leads to catastrophic loss explosion.<n>Our in-depth analysis reveals that the failure is not a random artifact but caused by two intertwined phenomena: the emergence of similar low-rank representations within the attention mechanism and the compounding effect of biased rounding errors inherent in low-precision arithmetic.
  arXiv  Detail & Related papers  (2025-10-05T14:01:24Z)
• An Active Learning-Based Streaming Pipeline for Reduced Data Training of Structure Finding Models in Neutron Diffractometry [1.3083205962260995]
  We introduce a novel batch-mode active learning (AL) policy that uses uncertainty sampling to simulate training data drawn from a probability distribution.<n>We confirm its efficacy in training the same models with about 75% less training data while improving the accuracy.<n>We then discuss the design of an efficient stream-based training workflow that uses this AL policy and present a performance study on two heterogeneous platforms.
  arXiv  Detail & Related papers  (2025-06-06T15:48:22Z)
• A Cost-Aware Approach to Adversarial Robustness in Neural Networks [1.622320874892682]
  We propose using accelerated failure time models to measure the effect of hardware choice, batch size, number of epochs, and test-set accuracy. We evaluate several GPU types and use the Tree Parzen Estimator to maximize model robustness and minimize model run-time simultaneously.
  arXiv  Detail & Related papers  (2024-09-11T20:43:59Z)
• Data Attribution for Diffusion Models: Timestep-induced Bias in Influence Estimation [53.27596811146316]
  Diffusion models operate over a sequence of timesteps instead of instantaneous input-output relationships in previous contexts. We present Diffusion-TracIn that incorporates this temporal dynamics and observe that samples' loss gradient norms are highly dependent on timestep. We introduce Diffusion-ReTrac as a re-normalized adaptation that enables the retrieval of training samples more targeted to the test sample of interest.
  arXiv  Detail & Related papers  (2024-01-17T07:58:18Z)
• Can Active Sampling Reduce Causal Confusion in Offline Reinforcement Learning? [58.942118128503104]
  Causal confusion is a phenomenon where an agent learns a policy that reflects imperfect spurious correlations in the data. This phenomenon is particularly pronounced in domains such as robotics. In this paper, we study causal confusion in offline reinforcement learning.
  arXiv  Detail & Related papers  (2023-12-28T17:54:56Z)
• Small-scale proxies for large-scale Transformer training instabilities [69.36381318171338]
  We seek ways to reproduce and study training stability and instability at smaller scales. By measuring the relationship between learning rate and loss across scales, we show that these instabilities also appear in small models when training at high learning rates. We study methods such as warm-up, weight decay, and the $mu$Param to train small models that achieve similar losses across orders of magnitude of learning rate variation.
  arXiv  Detail & Related papers  (2023-09-25T17:48:51Z)
• Exploring Weight Balancing on Long-Tailed Recognition Problem [32.01426831450348]
  Recognition problems in long-tailed data, in which the sample size per class is heavily skewed, have gained importance. Weight balancing, which combines classical regularization techniques with two-stage training, has been proposed. We analyze weight balancing by focusing on neural collapse and the cone effect at each training stage.
  arXiv  Detail & Related papers  (2023-05-26T01:45:19Z)
• PLATON: Pruning Large Transformer Models with Upper Confidence Bound of Weight Importance [114.1541203743303]
  We propose PLATON, which captures the uncertainty of importance scores by upper confidence bound (UCB) of importance estimation. We conduct extensive experiments with several Transformer-based models on natural language understanding, question answering and image classification.
  arXiv  Detail & Related papers  (2022-06-25T05:38:39Z)
• An Empirical Investigation of the Role of Pre-training in Lifelong Learning [21.995593026269578]
  We show that generic pre-training implicitly alleviates the effects of catastrophic forgetting when learning multiple tasks sequentially. We study this phenomenon by analyzing the loss landscape, finding that pre-trained weights appear to ease forgetting by leading to wider minima.
  arXiv  Detail & Related papers  (2021-12-16T19:00:55Z)
• On the Impact of Hard Adversarial Instances on Overfitting in Adversarial Training [70.82725772926949]
  Adversarial training is a popular method to robustify models against adversarial attacks. In this work, we investigate this phenomenon from the perspective of training instances. We show that the decay in generalization performance of adversarial training is a result of fitting hard adversarial instances.
  arXiv  Detail & Related papers  (2021-12-14T12:19:24Z)
• On the Robustness of Pretraining and Self-Supervision for a Deep Learning-based Analysis of Diabetic Retinopathy [70.71457102672545]
  We compare the impact of different training procedures for diabetic retinopathy grading. We investigate different aspects such as quantitative performance, statistics of the learned feature representations, interpretability and robustness to image distortions. Our results indicate that models from ImageNet pretraining report a significant increase in performance, generalization and robustness to image distortions.
  arXiv  Detail & Related papers  (2021-06-25T08:32:45Z)
• Precise Tradeoffs in Adversarial Training for Linear Regression [55.764306209771405]
  We provide a precise and comprehensive understanding of the role of adversarial training in the context of linear regression with Gaussian features. We precisely characterize the standard/robust accuracy and the corresponding tradeoff achieved by a contemporary mini-max adversarial training approach. Our theory for adversarial training algorithms also facilitates the rigorous study of how a variety of factors (size and quality of training data, model overparametrization etc.) affect the tradeoff between these two competing accuracies.
  arXiv  Detail & Related papers  (2020-02-24T19:01:47Z)

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.