dm-haiku: オブジェクト指向のモデル定義を可能にするJAX用のニューラルネットワークライブラリ

何を解決するか

JAXは強力な数値計算ライブラリですが、その変換（jitやgradなど）を利用するには関数が純粋（pure）である必要があります。Haikuは、開発者が使い慣れたオブジェクト指向プログラミングモデルを使用してニューラルネットワークを定義できるようにすることで、この問題を解決します。定義された「不純な」オブジェクト指向の定義を、JAXが処理可能な純粋な関数へと自動的に変換します。

仕組み

Haikuは、オブジェクト指向設計と関数的な純粋性の間のギャップを埋めるために、主に2つのツールを提供します。 \n* hk.Module: ネットワークのレイヤーやコンポーネントを定義するために使用されるPythonオブジェクトです。これらのモジュールはパラメータやメソッドへの参照を保持し、ユーザーが標準的なニューラルネットワークライブラリのようなコードを書けるようにします。

• hk.transform: hk.Moduleを使用する関数を、init（初期パラメータ値を収集する）とapply（計算のためにそれらのパラメータを関数に注入する）という一対の純粋な関数に変換する関数変換です。

バッチ正規化（batch normalization）のように内部的な可変状態を必要とするモデルの場合、Haikuはパラメータと状態を個別に管理するhk.transform_with_stateを提供します。

対象ユーザー

ニューラルネットワーク構築のためのオブジェクト指向APIの生産性を求めつつ、JAXの関数変換やハードウェアアクセラレーションの全能力を必要とする研究者や開発者。

ハイライト

• DeepMind Scale: DeepMindの研究者によって、大規模な画像、言語、および強化学習の実験でテストされています。
• ライブラリであり、フレームワークではない: カスタムのオプティマイザやチェックポイント形式を強制することなく、パラメータと状態の管理に焦点を当てた軽量なライブラリとして設計されています。
• hk.next_rng_key(): 決定論的なキーのシーケンスを提供することで、JAX内での乱数生成を簡略化します。
• JAX Compatibility: 複数のアクセラレータにわたる分散トレーニングのためのjax.pmapと完全に互換性があります。