mlsanityData to 7 Oct 2026

JaxPruner: A concise library for sparsity research

arXiv 2304.14082 ↗First posted 27 Apr 2023cs.LG cs.SEcited by 13

Not on the map. Why?

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.

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

Builds on

49 of 76 references.

References · 49

  1. 2017Attention Is All You Needcited by 59,932
  2. 2015Deep Residual Learning for Image Recognitioncited by 51,809
  3. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  4. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
  5. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,491
  6. 2014ImageNet Large Scale Visual Recognition Challengecited by 8,441
  7. 2018Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actorcited by 4,635
  8. 2022Flamingo: a Visual Language Model for Few-Shot Learningcited by 4,337
  9. 2018The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networkscited by 2,250
  10. 2013Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computationcited by 2,031
  11. 2016Pruning Filters for Efficient ConvNetscited by 1,129
  12. 2020Adaptive Federated Optimizationcited by 1,036
37 more
  1. 2015Automatic differentiation in machine learning: a surveycited by 846
  2. 2017Rainbow: Combining Improvements in Deep Reinforcement Learningcited by 774
  3. 2017To prune, or not to prune: exploring the efficacy of pruning for model compressioncited by 596
  4. 2020What is the State of Neural Network Pruning?cited by 578
  5. 2021Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networkscited by 521
  6. 2019Rigging the Lottery: Making All Tickets Winnerscited by 491
  7. 2019The State of Sparsity in Deep Neural Networkscited by 471
  8. 2017Distributional Reinforcement Learning with Quantile Regressioncited by 464
  9. 2018Implicit Quantile Networks for Distributional Reinforcement Learningcited by 362
  10. 2017Scalable Training of Artificial Neural Networks with Adaptive Sparse Connectivity inspired by Network Sciencecited by 325
  11. 2018Efficient Neural Audio Synthesiscited by 218
  12. 2021Accelerating Sparse Deep Neural Networkscited by 213
  13. 2017EMNIST: an extension of MNIST to handwritten letterscited by 212
  14. 2019Sparse Networks from Scratch: Faster Training without Losing Performancecited by 210
  15. 2022MegaBlocks: Efficient Sparse Training with Mixture-of-Expertscited by 204
  16. 2021Efficient and Modular Implicit Differentiationcited by 199
  17. 2022Scaling Up Models and Data with textttt5x and textttseqiocited by 187
  18. 2021Learning N:M Fine-grained Structured Sparse Neural Networks From Scratchcited by 177
  19. 2023The Dormant Neuron Phenomenon in Deep Reinforcement Learningcited by 171
  20. 2020Sparse GPU Kernels for Deep Learningcited by 138
  21. 2018Dopamine: A Research Framework for Deep Reinforcement Learningcited by 121
  22. 2021Accelerated Sparse Neural Training: A Provable and Efficient Method to Find N:M Transposable Maskscited by 103
  23. 2022Better plain ViT baselines for ImageNet-1kcited by 97
  24. 2019Fast Sparse ConvNetscited by 86
  25. 2022The State of Sparse Training in Deep Reinforcement Learningcited by 60
  26. 2021SCENIC: A JAX Library for Computer Vision Research and Beyondcited by 56
  27. 2016Pruning Convolutional Neural Networks for Resource Efficient Inferencecited by 47
  28. 2017Block-Sparse Recurrent Neural Networkscited by 38
  29. 2017Exploring the Regularity of Sparse Structure in Convolutional Neural Networkscited by 37
  30. 2021FedJAX: Federated learning simulation with JAXcited by 32
  31. 2022EvoJAX: Hardware-Accelerated Neuroevolutioncited by 29
  32. 2023Ten Lessons We Have Learned in the New "Sparseland": A Short Handbook for Sparse Neural Network Researcherscited by 25
  33. 2018Format Abstraction for Sparse Tensor Algebra Compilerscited by 16
  34. 2019Neural Network Distiller: A Python Package For DNN Compression Researchcited by 12
  35. 2022Compiler Support for Sparse Tensor Computations in MLIRcited by 8
  36. 2019GraphBLAST: A High-Performance Linear Algebra-based Graph Framework on the GPUcited by 6
  37. 2015Graphs, Matrices, and the GraphBLAS: Seven Good Reasonscited by 2

Built on by

Cited by year

Citing papers per year: 2023: 2; 2024: 6; 2025: 5.

13 papers.

Sort:

  1. 2023Scaling Laws for Sparsely-Connected Foundation Modelscited by 41
  2. 2024In value-based deep reinforcement learning, a pruned network is a good networkcited by 38
  3. 2024Neuroplastic Expansion in Deep Reinforcement Learningcited by 17
  4. 2024Dynamic Sparse Training versus Dense Training: The Unexpected Winner in Image Corruption Robustnesscited by 11
  5. 2025Eau De Q-Network: Adaptive Distillation of Neural Networks in Deep Reinforcement Learningcited by 6
  6. 2024Maxwell's Demon at Work: Efficient Pruning by Leveraging Saturation of Neuronscited by 5
  7. 2025Mind the GAP! The Challenges of Scale in Pixel-based Deep Reinforcement Learningcited by 4
  8. 2025Deep RL Needs Deep Behavior Analysis: Exploring Implicit Planning by Model-Free Agents in Open-Ended Environmentscited by 4
  9. 2024Navigating Extremes: Dynamic Sparsity in Large Output Spacescited by 3
  10. 2025Accelerating Linear Recurrent Neural Networks for the Edge with Unstructured Sparsitycited by 3
  11. 2025Matrix-Free Least Squares Solvers: Values, Gradients, and What to Do With Themcited by 1
  12. 2023ELSA: Partial Weight Freezing for Overhead-Free Sparse Network Deploymentcited by 0
1 more
  1. 2024Sparse-ProxSkip: Accelerated Sparse-to-Sparse Training in Federated Learningcited by 0