Research map / Architectures and efficiency / Pruning and efficient tuning
S-STE: Continuous Pruning Function for Efficient 2:4 Sparse Pre-training
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 S-STE, a continuous 2:4 soft-thresholding pruning function with fixed MSE-optimal per-tensor rescaling, plus MVUE and FP8, for efficient 2:4 sparse pre-training of transformers.
Problem
STE-based 2:4 sparse pre-training uses a discontinuous hard-thresholding pruning function, causing incorrect descent directions, unpredictable loss reduction, and sparse-mask oscillation, leading to accuracy gaps versus dense training.
Approach
Replace hard-thresholding with a continuous 2:4 soft-thresholding projection (subtract the third-largest magnitude from the top-two elements), rescale with a per-tensor factor computed once at initialization to minimize MSE, and combine with MVUE on output gradients and FP8 training.
Key results
- On Transformer-base WMT14 En-De, S-STE reaches 26.11 test BLEU vs 25.84 for SR-STE and 25.52 for STEP (dense 26.42).
- On DeiT-small ImageNet-1K, S-STE achieves 78.5% top-1 vs 75.7% for SR-STE and 77.6% for Bi-Mask (dense 79.9%).
- On GPT-2 774M, S-STE reaches 78.6 GLUE average and 75.5/85.5 SQuAD EM/F1, surpassing SR-STE+DF-style baseline (77.1 GLUE) and approaching dense (76.2 GLUE).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use S-STE instead of SR-STE | Transformer-base · WMT14 En-De test | BLEU: 25.84 -> 26.11 | Helped |
| use S-STE instead of STEP | Transformer-base · WMT14 En-De test | BLEU: 25.52 -> 26.11 | Helped |
| use S-STE instead of SR-STE | DeiT-small · ImageNet-1k | top-1 accuracy: 75.7 -> 78.5 | Helped |
| use S-STE instead of Bi-Mask | DeiT-small · ImageNet-1k | top-1 accuracy: 77.6 -> 78.5 | Helped |
| use S-STE instead of T-SR-STE+DF | GPT-2 774M · GLUE | average score: $77.1 \pm 0.4$ -> $78.6 \pm 0.8$ | Helped |
| use S-STE instead of T-SR-STE+DF | GPT-2 774M · SQuAD | F1: 84.6 -> 85.5 | Helped |
| remove weight rescaling (ablation) | S-STE on DeiT-small · ImageNet-1K | top-1 accuracy: 78.8 -> 78.5 | Hurt |
| add FP8 training (ablation) | S-STE on DeiT-small · ImageNet-1K | top-1 accuracy: 78.8 -> 78.6 | No clear effect |
| add MVUE on activation gradients (ablation) | S-STE on DeiT-small · ImageNet-1K | top-1 accuracy: 78.8 -> 78.9 | No clear effect |
| add FP8 and MVUE on activation gradients (ablation) | S-STE on DeiT-small · ImageNet-1K | top-1 accuracy: 78.8 -> 78.5 | Hurt |
| set gamma to 0 instead of 0.33 (ablation) | S-STE on Transformer-base · WMT14 En-De test | BLEU: 26.01 -> 26.30 | Helped |
| set gamma to 0 instead of 0.67 (ablation) | S-STE on Transformer-base · WMT14 En-De test | BLEU: 26.16 -> 26.30 | Helped |
Surprising findings
- Dynamically recomputing the scaling factor each iteration makes it grow over training and increases flip rate; fixing it after the first iteration aligns flip rate with dense training.
- Sparsifying the weight matrix with MVUE in the backward pass accumulates errors and hurts accuracy, so only the output gradient is sparsified.
Limitations
- No end-to-end wall-clock speedup is demonstrated; acceleration is measured only on FFN blocks and GPT-2 blocks on RTX 3090 with FP16, and FP8 2:4 kernels on H100 were unsatisfactory.
- Only FFN linear layers are sparsified; QKV projections are left for future work.
- The continuous pruning function is one possible solution; smoother alternatives may be needed.
Experimental setup
Datasets: WMT 14 En-De ImageNet-1K OpenWebText GLUE SQuAD
Benchmarks: WMT 14 En-De ImageNet-1K GLUE SQuAD
Metrics: BLEU top-1 accuracy top-5 accuracy validation loss GLUE average SQuAD EM SQuAD F1
Baselines: Dense SR-STE STEP Bi-Mask T-SR-STE+DF T-SR-STE
Builds on: STE SR-STE MVUE Bi-Mask STEP
Closest papers by meaning
- Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch (ICLR 2021 · similarity 0.89)
- Accelerating Transformer Pre-training with 2:4 Sparsity (ICML 2024 · similarity 0.87)
- Accelerated Sparse Neural Training: A Provable and Efficient Method to Find N:M Transposable Masks (NeurIPS 2021 · similarity 0.87)
- A Single-Step, Sharpness-Aware Minimization is All You Need to Achieve Efficient and Accurate Sparse Training (NeurIPS 2024 · similarity 0.86)
- Sparse-IFT: Sparse Iso-FLOP Transformations for Maximizing Training Efficiency (ICML 2024 · similarity 0.84)
- Dynamic Sparse Training with Structured Sparsity (ICLR 2024 · similarity 0.84)
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.