Sparse, 8-bit quantized S5 linear RNNs match dense model audio denoising accuracy with 2x less compute and 36% less memory, and run 42x faster with 149x lower energy on Loihi 2 than a dense FP32 model on Jetson Orin Nano.
JaxPruner: A concise library for sparsity research
1 Pith paper cite this work, alongside 1 external citations. Polarity classification is still indexing.
1
Pith paper citing it
1
external citations · Pith
abstract
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.
fields
cs.LG 1years
2025 1verdicts
CONDITIONAL 1representative citing papers
citing papers explorer
-
Accelerating Linear Recurrent Neural Networks for the Edge with Unstructured Sparsity
Sparse, 8-bit quantized S5 linear RNNs match dense model audio denoising accuracy with 2x less compute and 36% less memory, and run 42x faster with 149x lower energy on Loihi 2 than a dense FP32 model on Jetson Orin Nano.