Research map / Learning theory and algorithms / Optimal transport
SWAP: Sparse Entropic Wasserstein Regression for Robust Network Pruning
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 SWAP, a sparse entropic Wasserstein regression formulation for neural network pruning that interpolates gradient neighborhoods via entropic optimal transport, improving robustness to noisy gradients.
Problem
Network pruning relies on empirical Fisher information approximations whose gradients are inaccurate or noisy (noisy data, federated learning, analog memory, adversarial attacks), which degrades which weights are selected for removal.
Approach
Replaces the linear-regression (sparse LR) pruning objective with an entropic Wasserstein regression whose transport plan softly matches each gradient sample to neighborhoods of others, implicitly averaging gradients while preserving covariance; the LR formulation is a special case with a diagonal plan. Pruning is done by projected gradient descent with iterative hard thresholding and Sinkhorn pla
Key results
- On MobileNetV1 at 0.75 sparsity with 25% noise, EWR reduced testing loss by 8.13% relative to LR.
- On ResNet20 CIFAR10 at 0.98 sparsity, EWR reached 69.21% vs 68.15% for LR without noise (Table 1).
- On MobileNetV1 ImageNet at 0.8 sparsity, EWR achieved 66.82% vs 65.72% for LR and 16.38% for CBS.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace LR with entropic Wasserstein regression | ResNet20 pruning · CIFAR10, sparsity 0.95, 10% noisy data, noise = σ | testing loss: 2.83 ($\pm$0.02) -> 2.75 ($\pm$0.02) | Helped |
| Replace LR with entropic Wasserstein regression | MobileNetV1 pruning · ImageNet, sparsity 0.75, 25% noisy data, noise = σ | testing loss: 4.93 ($\pm$0.06) -> 4.53 ($\pm$0.06); 8.13% | Helped |
| Replace LR with entropic Wasserstein regression | MobileNetV1 pruning · ImageNet, sparsity 0.75, 25% noisy data, noise = 2σ | Top-1 accuracy: 13.78 ($\pm$0.21) -> 14.78 ($\pm$0.25); 7.29% | Helped |
| Replace LR with entropic Wasserstein regression | ResNet20 pruning · CIFAR10, sparsity 0.95, 25% noisy data, noise = 2σ | Top-1 accuracy: 11.68 ($\pm$0.18) -> 12.04 ($\pm$0.16); 3.04% | Helped |
| Replace LR with entropic Wasserstein regression | MLPNet pruning · MNIST, sparsity 0.5 | accuracy: 95.26 -> 95.24 | No clear effect |
| Set ε = 0 instead of ε = 1 (ablation) | EWR pruning · CIFAR10, ResNet20, sparsity 0.95, 10% noisy data, noise = σ | testing loss: 2.78 -> 2.77 | Helped |
| Set ε = 2 instead of ε = 1 (ablation) | EWR pruning · CIFAR10, ResNet20, sparsity 0.95, 10% noisy data, noise = σ | testing loss: 2.78 -> 2.78 | No clear effect |
| Set ε = 10 instead of ε = 1 (ablation) | EWR pruning · CIFAR10, ResNet20, sparsity 0.95, 10% noisy data, noise = σ | testing loss: 2.78 -> 2.79 | Hurt |
| Set ε = ∞ instead of ε = 1 (ablation) | EWR pruning · CIFAR10, ResNet20, sparsity 0.95, 10% noisy data, noise = σ | testing loss: 2.78 -> 2.97 | Hurt |
| Increase ε from 1 to ∞ (ablation) | EWR pruning · CIFAR10, ResNet20, sparsity 0.84, 10% noisy data, noise = σ | testing loss: 1.61 -> 2.62 | Hurt |
| Vary ε across 0, 1, 2, 10, and ∞ (ablation) | EWR pruning · CIFAR10, ResNet20, sparsity 0.63 and lower | testing loss: 0.36; 0.36; 0.36; 0.36; 0.37 | No clear effect |
Surprising findings
- With fully large entropic regularization (epsilon to infinity), EWR loss becomes worse than LR at high sparsity.
Limitations
- Gains over prior SoTA on clean data are modest and mostly within confidence intervals on several settings.
- Pruning evaluated without fine-tuning for noise robustness experiments; the approach relies on Gaussian assumptions for a closed-form OT plan.
- Sensitivity to the entropic regularization epsilon; excessive values degrade performance at high sparsity.
Experimental setup
Datasets: MNIST CIFAR10 ImageNet
Benchmarks: MNIST CIFAR10 ImageNet
Metrics: top-1 accuracy testing loss
Baselines: Magnitude pruning WoodFisher Combinatorial Brain Surgeon LR sparse linear regression pruning
Builds on: WoodFisher Combinatorial Brain Surgeon Optimal Brain Surgeon Sinkhorn-Knopp algorithm
Closest papers by meaning
- Sparsest Models Elude Pruning: An Exposé of Pruning’s Current Capabilities (ICML 2024 · similarity 0.85)
- A Gradual, Semi-Discrete Approach to Generative Network Training via Explicit Wasserstein Minimization (ICML 2019 · similarity 0.85)
- Sparsistency for inverse optimal transport (ICLR 2024 · similarity 0.85)
- Pruning via Sparsity-indexed ODE: a Continuous Sparsity Viewpoint (ICML 2023 · similarity 0.84)
- SAFE: Finding Sparse and Flat Minima to Improve Pruning (ICML 2025 · similarity 0.84)
- Towards Meta-Pruning via Optimal Transport (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.