JAXを使用して研究を加速化する

'Accelerating research with JAX.'

DeepMindのエンジニアは、ツールの構築、アルゴリズムの拡大、仮想および物理的な世界の作成などにより、人工知能（AI）システムのトレーニングとテストを加速させるための研究を行っています。この仕事の一環として、私たちは常に新しい機械学習ライブラリやフレームワークを評価しています。

最近、Google Researchチームが開発した機械学習フレームワークであるJAXが、ますます多くのプロジェクトに適していることがわかりました。JAXは私たちのエンジニアリングの哲学に合致し、昨年から私たちの研究コミュニティで広く採用されています。ここでは、JAXとの作業の経験を共有し、AI研究に役立つ理由を概説し、研究者をサポートするために構築しているエコシステムの概要を説明します。

なぜJAXを選ぶのか？

JAXは、特に機械学習の研究に適した高性能な数値計算のためのPythonライブラリです。数値関数のAPIは科学計算で使用される関数のコレクションであるNumPyに基づいています。PythonとNumPyは広く使用され、馴染み深いため、JAXはシンプルで柔軟で採用しやすいです。

• MuZero ルールなしでGo、チェス、将棋、アタリをマスターする
• AI研究によるスポーツ分析の進歩
• アルファフォールドの力を世界の手に

NumPyのAPIに加えて、JAXには機械学習の研究をサポートするための拡張可能な関数変換システムが含まれており、以下のような機能を提供しています：

• 微分：勾配ベースの最適化は機械学習にとって基本的です。JAXは、grad、hessian、jacfwd、jacrevなどの関数変換を介して、任意の数値関数の順方向および逆方向の自動微分をネイティブにサポートしています。
• ベクトル化：機械学習の研究では、しばしば大量のデータに対して単一の関数を適用します。例えば、バッチ全体の損失を計算したり、差分プライバシー学習のために個別の例ごとの勾配を評価したりします。JAXは、この形式のプログラミングを簡素化するvmap変換を介した自動ベクトル化を提供します。例えば、新しいアルゴリズムを実装する際にバッチ処理について考慮する必要がありません。JAXはまた、関連するpmap変換を介した大規模なデータ並列処理もサポートし、単一のアクセラレータのメモリに収まりきらないデータをエレガントに分散します。
• JITコンパイル：XLAは、JAXプログラムをGPUおよびCloud TPUアクセラレータ上でジャストインタイム（JIT）コンパイルおよび実行するために使用されます。JITコンパイルは、JAXのNumPy互換のAPIと組み合わせることで、ハイパフォーマンスコンピューティングの以前の経験がない研究者でも簡単に1つまたは複数のアクセラレータにスケールすることを可能にします。

JAXを使用することで、新しいアルゴリズムやアーキテクチャの迅速な実験が可能になり、最近の多くの論文で使用されています。詳細については、NeurIPS仮想会議で開催されるJAX Roundtableに参加してください（日時：2021年12月9日午後7時GMT）。

DeepMindにおけるJAX

最先端のAI研究を支援するには、迅速なプロトタイピングと素早い反復を、本番システムに関連付けられるスケールで実験を展開する能力とのバランスを取る必要があります。この種のプロジェクトを特に難しいものにしているのは、研究の状況が急速に変化し予測が難しいことです。いつでも、新しい研究のブレイクスルーが、そして定期的に変化し、全チームの軌道と要件を変える可能性があります。この変化し続ける状況の中で、私たちのエンジニアリングチームの重要な責任は、1つの研究プロジェクトのために学んだ教訓と書かれたコードが次のプロジェクトで効果的に再利用されることを確認することです。

成功したアプローチの1つは、モジュール化です：各研究プロジェクトで開発された最も重要で重要な構成要素を、よくテストされ効率的なコンポーネントに抽出します。これにより、研究者は研究に集中することができるだけでなく、コードの再利用、バグ修正、アルゴリズムの性能向上による利点も得ることができます。また、各ライブラリに明確に定義されたスコープがあること、および相互運用可能で独立していることを確認することも重要です。研究者にとって最大の柔軟性を提供し、常に適切なツールを選択するための最大の柔軟性を提供するために、インクリメンタルな参加が必要です。

JAXエコシステムの開発に関して考慮された他の要素には、既存のTensorFlowライブラリ（SonnetやTRFLなど）の設計と一貫性を保つことが含まれます。また、可能な限り基礎となる数学に一致するようにコンポーネントを構築し、自己説明的であり、論文からコードへの移行を最小限に抑えることも重視しています。最後に、私たちはライブラリをオープンソース化し、研究成果の共有を容易にし、広いコミュニティにJAXエコシステムの探索を促すために選択しました。

当社のエコシステム

俳句 ‍

組み合わせ可能な関数変換のJAXプログラミングモデルは、状態を持つオブジェクト（例：訓練可能なパラメータを持つニューラルネットワーク）を扱うことが複雑になる場合があります。Haikuは、JAXの純粋な関数型パラダイムのパワーとシンプリシティを活用しながら、ユーザーが馴染みのあるオブジェクト指向プログラミングモデルを使用できるニューラルネットワークライブラリです。

Haikuは、DeepMindとGoogleの数百人の研究者によって積極的に使用されており、Coax、DeepChem、NumPyroなどのいくつかの外部プロジェクトでも採用されています。Haikuは、TensorFlowのニューラルネットワークのモジュールベースのプログラミングモデルであるSonnetのAPIに基づいて構築されており、SonnetからHaikuへの移植を可能な限り簡単にすることを目指しています。

GitHubで詳細を確認してください

Optax

勾配ベースの最適化は、機械学習において基本的なものです。Optaxは、勾配変換のライブラリと、合成演算子（例：chain）を提供し、たった1行のコードで多くの標準的な最適化手法（例：RMSPropまたはAdam）を実装することができます。

Optaxの合成的な性質は、基本的な要素をカスタム最適化手法で再組み合わせすることを自然にサポートします。さらに、確率的勾配推定や二次最適化のためのユーティリティも提供しています。

多くのOptaxユーザーはHaikuを採用していますが、インクリメンタルバイインの哲学に従い、JAXツリー構造でパラメータを表現する任意のライブラリ（例：Elegy、Flax、Stax）がサポートされています。JAXライブラリの豊富なエコシステムについては、こちらをご覧ください。

GitHubで詳細を確認してください

RLax

私たちの最も成功したプロジェクトの多くは、深層学習と強化学習（RL）の交差点にあります。RLaxは、RLエージェントを構築するための有用なビルディングブロックを提供するライブラリです。

RLaxのコンポーネントは、TD学習、ポリシーグラディエント、アクターコントリビューション、MAP、プロキシマルポリシーオプティマイゼーション、非線形値変換、一般的な価値関数、およびいくつかの探索手法など、さまざまなアルゴリズムやアイデアをカバーしています。

いくつかの初心者向けの例のエージェントも提供されていますが、RLaxは完全なRLエージェントシステムの構築と展開を目的としたフレームワークではありません。RLaxコンポーネントを利用したフル機能のエージェントフレームワークの例としては、Acmeなどがあります。

GitHubで詳細を確認してください

Chex

テストはソフトウェアの信頼性にとって重要であり、研究コードも例外ではありません。研究実験から科学的な結論を導くためには、コードの正確性に自信を持つことが必要です。Chexは、ライブラリの作者が一般的なビルディングブロックの正確性と堅牢性を検証するために使用するテストユーティリティのコレクションであり、エンドユーザーが実験コードをチェックするためにも使用されます。

Chexは、JAXに対応した単体テスト、JAXデータ型のプロパティのアサーション、モックとフェイク、およびマルチデバイスのテスト環境など、さまざまなユーティリティを提供しています。Chexは、DeepMindのJAXエコシステムやCoax、MineRLなどの外部プロジェクトでも使用されています。

GitHubで詳細を確認してください

Jraph

グラフニューラルネットワーク（GNN）は、多くの有望な応用分野で研究が行われています。たとえば、Googleマップでの交通予測や物理シミュレーションの最近の研究などをご覧ください。Jraph（発音は「ジラフ」）は、JAXでGNNを扱うための軽量なライブラリです。

Jraphは、グラフの標準化されたデータ構造、グラフの操作に使用するユーティリティのセット、そして容易にフォーク可能で拡張可能なグラフニューラルネットワークモデルの「動物園」を提供します。その他の主な特徴には、ハードウェアアクセラレータを効率的に活用するGraphTupleのバッチ処理、パディングとマスキングを介した可変形状グラフのJITコンパイルサポート、および入力パーティション上で定義された損失があります。Optaxやその他のライブラリと同様に、Jraphはユーザーのニューラルネットワークライブラリの選択に制約を設けません。

ライブラリの使用方法については、豊富な例コレクションから詳しく学ぶことができます。

GitHubで詳細を確認

JAXエコシステムは常に進化しており、ML研究コミュニティには当社のライブラリとJAXの潜在能力を探求することをお勧めします。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!