FlashBoB cuts memory needs for exact gradients in long softmax attention
FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention
Machine LearningArtificial Intelligence
Summary
Working with transformer models that process very long sequences is slow and uses a lot of memory due to the attention mechanism. The authors developed FlashBoB, a method that computes exact gradients during training more efficiently by avoiding large memory use and heavy data movement. This lets models handle much longer sequences on a single GPU than before, speeding up certain advanced training techniques that require precise gradient calculations. FlashBoB achieves this by cleverly reusing small internal computations and reducing the need to store big intermediate data.
What this means in practice
- •For machine learning engineers: Train transformer models with very long input sequences using exact second-order gradients on a single GPU without running out of memory.
- •For ai model trainers: Perform test-time training and meta-learning that depend on accurate backward-over-backward gradients efficiently at long sequence lengths.
Authors
Anthony Givans, Michael Crawshaw, Mingrui Liu
Abstract
Transformer models built on the attention mechanism have become a central building block in modern deep learning, yet softmax attention remains a major bottleneck for long-context workloads. While FlashAttention makes the forward and first backward passes I/O-efficient, it does not support backward-over-backward (BoB), which enables exact differentiation through the backward pass for applications such as second-order optimization, test-time training, gradient-based memory, and meta-learning. Existing BoB implementations either materialize large intermediate tensors or exhaust GPU memory at long sequence lengths. We present FlashBoB, an exact, I/O-efficient algorithm for BoB in softmax attention that keeps computation within on-chip tiles and avoids all $N \times N$ intermediate tensors, where $N$ is the sequence length. The key insight is a hierarchical affine structure in the softmax double backward: two row-wise scalars determine all outputs through affine transformations. This yields a two-pass schedule with bounded on-chip static random-access memory (SRAM) usage and minimal off-chip high-bandwidth memory (HBM) traffic. FlashBoB achieves $Θ(N^2 d^2/M)$ HBM traffic ($d$ is the head dimension and $M$ is the memory size) and, within the standard FlashAttention-style score-recomputation model, matches the inherited large-cache lower bound for exact forward attention. Empirically, it scales exact attention BoB to $N=262\text{K}$ on a single A100 80GB GPU, where prior PyTorch exact baselines fail by $N=16\text{K}$, and is up to $6.3\times$ faster than FlashBack. These results make exact second-order attention practical at long-context sequence lengths where prior implementations cannot run efficiently.