JaxPruner: A concise library for sparsity research

要約

タイトル：JaxPruner：スパーシティー研究のための簡潔なライブラリー

要約：

– JaxPrunerは、オープンソースのJAXベースのプルーニングおよびスパースなトレーニングライブラリであり、スパースニューラルネットワークの研究を加速することを目的としています。
– JaxPrunerは、人気のあるプルーニングおよびスパーストレーニングアルゴリズムの簡潔な実装を提供し、最小限のメモリとレイテンシのオーバーヘッドで動作します。
– JaxPrunerで実装されたアルゴリズムは共通のAPIを使用し、Optaxという人気のある最適化ライブラリとシームレスに連携して動作します。これにより、既存のJAXベースのライブラリとの簡単な統合が可能になります。
– Scenic、t5x、Dopamine、FedJAXの4つのコードベースでのエクスペリメント例を提供し、人気のあるベンチマークのベースライン実験を提供します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, OpenAI