# JaxPruner: A concise library for sparsity research

arXiv: https://arxiv.org/abs/2304.14082 · HTML: https://mlsanity.com/paper/2304.14082/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 49 of its 76 references that are in the map.
- First posted on arXiv 27 Apr 2023.

## Builds on (most cited first, 30 of 49)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Flamingo: a Visual Language Model for Few-Shot Learning](https://mlsanity.com/paper/2204.14198/) (2022)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Pruning Filters for Efficient ConvNets](https://mlsanity.com/paper/1608.08710/) (2016)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [To prune, or not to prune: exploring the efficacy of pruning for model compression](https://mlsanity.com/paper/1710.01878/) (2017)
- [What is the State of Neural Network Pruning?](https://mlsanity.com/paper/2003.03033/) (2020)
- [Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks](https://mlsanity.com/paper/2102.00554/) (2021)
- [Rigging the Lottery: Making All Tickets Winners](https://mlsanity.com/paper/1911.11134/) (2019)
- [The State of Sparsity in Deep Neural Networks](https://mlsanity.com/paper/1902.09574/) (2019)
- [Distributional Reinforcement Learning with Quantile Regression](https://mlsanity.com/paper/1710.10044/) (2017)
- [Implicit Quantile Networks for Distributional Reinforcement Learning](https://mlsanity.com/paper/1806.06923/) (2018)
- [Scalable Training of Artificial Neural Networks with Adaptive Sparse Connectivity inspired by Network Science](https://mlsanity.com/paper/1707.04780/) (2017)
- [Accelerating Sparse Deep Neural Networks](https://mlsanity.com/paper/2104.08378/) (2021)
- [Efficient Neural Audio Synthesis](https://mlsanity.com/paper/1802.08435/) (2018)
- [EMNIST: an extension of MNIST to handwritten letters](https://mlsanity.com/paper/1702.05373/) (2017)
- [Sparse Networks from Scratch: Faster Training without Losing Performance](https://mlsanity.com/paper/1907.04840/) (2019)
- [MegaBlocks: Efficient Sparse Training with Mixture-of-Experts](https://mlsanity.com/paper/2211.15841/) (2022)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Scaling Up Models and Data with textttt5x and textttseqio](https://mlsanity.com/paper/2203.17189/) (2022)
- [Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch](https://mlsanity.com/paper/2102.04010/) (2021)

## Built on by (most cited first, 13 of 13)

- [Scaling Laws for Sparsely-Connected Foundation Models](https://mlsanity.com/paper/2309.08520/) (2023)
- [In value-based deep reinforcement learning, a pruned network is a good network](https://mlsanity.com/paper/2402.12479/) (2024)
- [Neuroplastic Expansion in Deep Reinforcement Learning](https://mlsanity.com/paper/2410.07994/) (2024)
- [Dynamic Sparse Training versus Dense Training: The Unexpected Winner in Image Corruption Robustness](https://mlsanity.com/paper/2410.03030/) (2024)
- [Maxwell's Demon at Work: Efficient Pruning by Leveraging Saturation of Neurons](https://mlsanity.com/paper/2403.07688/) (2024)
- [Eau De Q-Network: Adaptive Distillation of Neural Networks in Deep Reinforcement Learning](https://mlsanity.com/paper/2503.01437/) (2025)
- [Mind the GAP! The Challenges of Scale in Pixel-based Deep Reinforcement Learning](https://mlsanity.com/paper/2505.17749/) (2025)
- [Deep RL Needs Deep Behavior Analysis: Exploring Implicit Planning by Model-Free Agents in Open-Ended Environments](https://mlsanity.com/paper/2506.06981/) (2025)
- [Navigating Extremes: Dynamic Sparsity in Large Output Spaces](https://mlsanity.com/paper/2411.03171/) (2024)
- [Accelerating Linear Recurrent Neural Networks for the Edge with Unstructured Sparsity](https://mlsanity.com/paper/2502.01330/) (2025)
- [Matrix-Free Least Squares Solvers: Values, Gradients, and What to Do With Them](https://mlsanity.com/paper/2510.19634/) (2025)
- [ELSA: Partial Weight Freezing for Overhead-Free Sparse Network Deployment](https://mlsanity.com/paper/2312.06872/) (2023)
- [Sparse-ProxSkip: Accelerated Sparse-to-Sparse Training in Federated Learning](https://mlsanity.com/paper/2405.20623/) (2024)

## Abstract (arXiv)

This paper introduces JaxPruner, an open-source JAX-based pruning and sparse training library for machine learning research. JaxPruner aims to accelerate research on sparse neural networks by providing concise implementations of popular pruning and sparse training algorithms with minimal memory and latency overhead. Algorithms implemented in JaxPruner use a common API and work seamlessly with the popular optimization library Optax, which, in turn, enables easy integration with existing JAX based libraries. We demonstrate this ease of integration by providing examples in four different codebases: Scenic, t5x, Dopamine and FedJAX and provide baseline experiments on popular benchmarks.