JaxPruner: A concise library for sparsity research

要約

タイトル：JaxPruner: スパーシティ（希薄性）研究のための簡潔なライブラリ

要約：本論文では、オープンソースのJAX（高性能なPythonライブラリ）に基づくプルーニングとスパーストレーニングのライブラリであるJaxPrunerを紹介する。JaxPrunerは、最小限のメモリ使用量とレイテンシーオーバーヘッドで、人気のあるプルーニングとスパーストレーニングアルゴリズムを簡潔かつ実装することで、スパースニューラルネットワークの研究を加速することを目的としている。JaxPrunerで実装されたアルゴリズムは、共通のAPIを使用し、Optaxという人気のある最適化ライブラリとシームレスに連携するため、既存のJAXベースのライブラリと簡単に統合できる。Scenic、t5x、Dopamine、FedJAXの4つのコードベースで例を提供し、人気のあるベンチマークのベースライン実験を提供することで、統合の容易さを示す。

要点：

– 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