Research map / Architectures and efficiency / Pruning and efficient tuning
Accelerating Transformer Pre-training with 2:4 Sparsity
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes an end-to-end 2:4 fully sparse training method for transformer pre-training (FFN layers only), with flip-rate-guided masked decay on gradients, fast decay-factor selection, dense fine-tuning, and kernel-level speedups (convolution-based transposable mask search, GEGLU kernel).
Problem
Pre-training large transformers is computationally expensive, and prior 2:4 sparse training methods were designed for CNNs, lost accuracy on transformers, and reported only simulated rather than actual speedups.
Approach
Monitor mask stability with a flip rate, apply masked decay on gradients (normalized by Adam's second moment), choose the decay factor by a fast warm-up flip-rate ratio test, move the dense steps to a final fine-tuning phase, and accelerate transposable mask search and gated activations with GPU kernels.
Key results
- On BERT-base, the proposed method reaches 79.6 GLUE average vs 79.8 for dense, outperforming STEP (77.7) and Bi-Mask (77.7).
- On GPT-2 774M, the method matches dense GLUE average (77.1 vs 76.2) with 1.2x end-to-end training speedup on RTX3090.
- A single FFN layer is accelerated up to 1.7x and a transformer block up to 1.3x; the transposable mask search is up to about 5x faster than the 2-approximation.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use 2:4 sparse training with the proposed techniques instead of dense training | BERT-base · GLUE (WNLI excluded) | Avg score: 79.8 ± 0.4 -> 79.6 ± 0.6 | No clear effect |
| Use 2:4 sparse training instead of STEP | BERT-base · GLUE (WNLI excluded) | Avg score: 77.7 ± 0.1 -> 79.6 ± 0.6 | Helped |
| Use 2:4 sparse training instead of a half-width dense FFN | BERT-base · GLUE (WNLI excluded) | Avg score: 77.9 ± 0.4 -> 79.6 ± 0.6 | Helped |
| Use 2:4 sparse training instead of dense training | GPT-2 350M · SQuAD | F1: 83.6 -> 82.4 | Hurt |
| Use 2:4 sparse training instead of dense training | Transformer-base · WMT 14 En-De | Test BLEU: 26.15 -> 26.48 | Helped |
| Use 2:4 sparse training | GPT-2 774M · GPT-2 pre-training | Speedup: 1.21 | Helped |
| Add masked decay (ablation) | BERT-base 2:4 training · GLUE | Avg score: 77.6 ± 0.2 -> 79.2 ± 0.2 | Helped |
| Add MVUE (ablation) | BERT-base 2:4 training with masked decay · GLUE | Avg score: 79.2 ± 0.2 -> 78.4 ± 0.3 | Hurt |
| Add dense fine-tuning (ablation) | BERT-base 2:4 training with masked decay · GLUE | Avg score: 79.2 ± 0.2 -> 79.4 ± 0.2 | Helped |
| Add MVUE to masked decay and dense fine-tuning (ablation) | BERT-base 2:4 training · GLUE | Avg score: 79.4 ± 0.2 -> 79.6 ± 0.6 | No clear effect |
Surprising findings
- Simply halving the FFN inner dimension (Half model) outperforms most prior 2:4 sparse training methods at identical FLOPs.
- Applying masked decay on weights rather than gradients has no effect on suppressing flip rate explosion for STE.
Limitations
- Sparsity is applied only to FFN layers; attention and other GEMMs remain dense.
- Speedups measured only on RTX3090 with GPT-2 models, and end-to-end speedup is modest (1.18-1.21x).
Experimental setup
Datasets: C4 OpenWebText ImageNet-1K WMT 14 En-De
Benchmarks: GLUE SQuAD ImageNet-1K WMT 14 En-De BLEU
Metrics: GLUE average score EM/F1 Acc@1 Acc@5 BLEU validation loss training speedup
Baselines: Dense Half STEP Bi-Mask Original DeiT
Builds on: SR-STE STEP Bi-Mask MVUE Cramming
Closest papers by meaning
- S-STE: Continuous Pruning Function for Efficient 2:4 Sparse Pre-training (NeurIPS 2024 · similarity 0.87)
- Two Heads are Better than One: Simulating Large Transformers with Small Ones (NeurIPS 2025 · similarity 0.85)
- CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs (NeurIPS 2026 · similarity 0.85)
- Jetfire: Efficient and Accurate Transformer Pretraining with INT8 Data Flow and Per-Block Quantization (ICML 2024 · similarity 0.85)
- Dense Structural Compression of Transformers via Gauge-Correct Channel Removal (NeurIPS 2026 · similarity 0.85)
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.