DeepMindがニューラルネットワークと強化学習ライブラリをリリース、ネットユーザー：JAXの開発を促進

この記事はAI新メディアQuantum Bit（公開アカウントID：QbitAI）より許可を得て転載しています。転載の際は出典元にご連絡ください。

DeepMind は本日、JAX をベースにした 2 つのライブラリ、Haiku と RLax をリリースしました。

JAX は Google によって提案され、TensorFlow 用の簡略化されたライブラリです。これは、線形代数用のコンパイラである XLA と、ネイティブ Python コードと Numpy コードを自動的に区別するライブラリである Autograd を組み合わせて、高性能機械学習の研究に使用します。

今回リリースされた 2 つのライブラリは、それぞれニューラル ネットワークと強化学習を対象としており、JAX の使用を大幅に簡素化します。

Haiku は JAX をベースにしたニューラル ネットワーク ライブラリであり、ユーザーは使い慣れたオブジェクト指向プログラミング モデルを使用し、JAX の純粋な関数変換に完全にアクセスできます。

RLax は、強化学習エージェントを実装するための便利なビルディング ブロックを提供する、JAX 上のライブラリです。

興味深いことに、Reddit ユーザーは Haiku ライブラリの名前が「ax」で終わらないことに驚きました。

もちろん、ネットユーザーの中には、この2つのライブラリを肯定する意見も出ている。

それがJAXを推進したことは間違いありません。

それでは、Haiku と RLex の本当の姿を見てみましょう。

俳句

Haiku は、JAX のニューラル ネットワーク ライブラリであり、ユーザーは使い慣れたオブジェクト指向プログラミング モデルを使用しながら、JAX の純粋に機能的な変換に完全にアクセスできます。

モジュール抽象化 hk.Module と単純な関数変換 hk.transform という 2 つのコア ツールを提供します。

hk.Module は、独自のパラメータ、他のモジュール、およびユーザー入力に関数を適用するためのメソッドへの参照を含む Python オブジェクトです。

hk.transform を使用すると、JAX の純粋に機能的な変換に完全にアクセスできます。

JAX には多くのニューラル ネットワーク ライブラリがありますが、Haiku の特別な点は何でしょうか?ポイントは5つあります。

1. HaikuはDeepMindの研究者によって大規模にテストされている

DeepMind は、Haiku と JAX で多くの実験を比較的簡単に再現しました。これらには、画像および言語処理、生成モデル、強化学習における大規模な結果が含まれます。

2. Haikuはライブラリであり、フレームワークではない

これは、モデル パラメータやその他のモデル状態の管理など、いくつかの特定の事項を簡素化するように設計されています。他のライブラリと組み合わせて記述でき、JAX の他の部分と連携して動作します。

3. Haikuは新しいスタートアップではない

これは、DeepMind のほぼ普遍的に採用されているニューラル ネットワーク ライブラリである Sonnet のプログラミング モデルと API に基づいて構築されています。 JAX 関数変換へのアクセスを維持しながら、状態管理のための Sonnet のモジュールベースのプログラミング モデルを保持します。

4. 俳句への移行は簡単

慎重な設計により、TensorFlow と Sonnet から JAX と Haiku への移行は比較的容易です。新しい関数 (hk.transform など) とは別に、Haiku は Sonnet 2 の API となる予定です。

5. Haiku は JAX を簡素化します

乱数を処理するためのシンプルなモデルを提供します。変換された関数では、hk.next_rng_key() は一意の rng キーを返します。

では、Haiku はどのようにインストールするのでしょうか?

Haiku は純粋な Python で書かれていますが、JAX を介して C++ コードに依存しています。

まず、以下のリンクの手順に従って、関連するアクセラレータ サポートを備えた JAX をインストールします。

https://github.com/google/jax#インストール

その後、インストールを完了するには、簡単な pip コマンドを 1 つだけ実行する必要があります。

 $ pip git+https: //github.com/deepmind/haikuをインストールします 

次に、ニューラル ネットワークと損失関数の例を示します。

俳句をhkとしてインポート
 
 jax.numpyをjnpとしてインポートする
 
 def softmax_cross_entropy(ロジット、ラベル): 
 
  one_hot = hk.one_hot(ラベル、logits.shape[- 1 ]) 
 
 -jnp.sum(jax.nn.log_softmax(logits) * one_hot、axis=- 1 )を返します。 
 
 def loss_fn(画像、ラベル): 
 
  モデル = hk.Sequential([ 
 
      hk.Linear( 1000 )、 
 
      jax.nn.relu、 
 
      hk.Linear( 100 )、 
 
      jax.nn.relu、 
 
      hk.Linear( 10 )、 
 
  ]) 
 
  logits = モデル(画像) 
 
 jnp.mean(softmax_cross_entropy(logits, labels))を返します。 
 
 loss_obj = hk.transform(loss_fn)

RLax

RLax は、強化学習エージェントを実装するための便利なビルディング ブロックを提供する、JAX 上のライブラリです。

提供される操作と関数は完全なアルゴリズムではなく、強化学習のための特定の数学的操作の実装です。

RLax のインストールも非常に簡単で、pip コマンドで実行できます。

 pip で git+git をインストールします: //github.com/deepmind/rlax.git  

JAX jax.jit 関数を使用すると、すべての RLax コードを異なるハードウェア上でコンパイルできます。

RLax に関して注意する必要があるのは、その命名規則です。

多くの関数は、出力を計算するために、連続する時間ステップにわたるポリシー、アクション、報酬、および値を考慮します。この場合、接尾辞 _t と tm1 は通常、各入力がどのステップで生成されたかを示すために使用されます。次に例を示します。

q_tm1: 遷移のソース状態における操作の値。

a_tm1: ソース状態で選択された操作。

r_t: 目標状態で収集された結果の報酬。

q_t: ターゲット状態における操作値。

Haiku と RLax は両方とも GitHub でオープンソース化されています。興味のある読者は「ポータル」のリンクからアクセスできます。

ポータル

俳句：

https://github.com/deepmind/haiku

RLax:

https://github.com/deepmind/rlax