google-deepmind/dm-haiku

解決する課題

HaikuはJAX用のニューラルネットワークライブラリであり、開発者が（SonnetやTensorFlowで見られるような）馴染みのあるオブジェクト指向プログラミングモデルを使用しながら、JAXの純粋関数変換への完全なアクセスを維持することを可能にします。初期化のために膨大なボイラープレートを書く必要なく、モデルのパラメータや内部状態の管理を簡素化します。

仕組み

Haikuは、オブジェクト指向設計と関数的な純粋性の間のギャップを埋めるために、主に2つのツールを提供します：

• hk.Module: ニューラルネットワークのレイヤーやコンポーネントを定義するために使用されるPythonオブジェクト。これらのモジュールはパラメータやメソッドへの参照を保持しますが、定義中は関数的に「不純」として扱われます。
• hk.transform: これらの不純なモジュールベースの関数を、init（初期パラメータ値を収集する）とapply（計算のためにそれらのパラメータを関数に注入する）という一対の純粋関数に変換する関数変換です。

可変状態（Batch Normalizationの移動平均など）を必要とするモデルのために、Haikuはパラメータと状態の両方を個別に管理するhk.transform_with_stateを提供しています。

対象者

ニューラルネットワーク構築にオブジェクト指向APIの生産性を求めつつ、JAXのパフォーマンスと変換機能（jax.jit、jax.grad、jax.pmapなど）を必要とする研究者や開発者。

ハイライト

• JAX統合: JAXの自動微分およびGPU/TPUサポートと完全に互換性があります。
• SonnetライクなAPI: Sonnet 2 APIに極めて近い設計となっており、TensorFlow/Sonnetからの移行が容易です。
• 簡素化されたRNG管理: 変換された関数内で決定論的に乱数生成を扱うためのhk.next_rng_key()を提供します。
• スケーラブル: DeepMindの研究者によって、画像、言語、強化学習タスクにおいて大規模にテストされています。
• ライブラリであり、フレームワークではない: パラメータと状態の管理に厳密に焦点を当て、オプティマイザやチェックポインティングは他のライブラリに委ねています。