PyTorchに負けました！ GoogleはTensorFlowを放棄し、JAXに賭ける

私はネットユーザーの一部が言ったことにとても満足しています。

「この子は本当にダメだ、もう一匹産もう。」

Google は実際にこれを実行しました。

7年間育てられてきたTensorFlowは、ついにMetaのPyTorchに敗北した、というところまで来ました。

何かがおかしいと気づいた Google は、すぐに別のフレームワーク「JAX」を要求しました。これは、新しい機械学習フレームワークです。

最近人気の DALL·E Mini については誰もが知っています。そのモデルは JAX に基づいてプログラムされており、Google TPU の利点を最大限に活用しています。

TensorFlow の黄昏と PyTorch の台頭

2015年にGoogleが開発した機械学習フレームワーク「TensorFlow」がリリースされました。

当時、TensorFlow は Google Brain の小さなプロジェクトに過ぎませんでした。

TensorFlow がリリースされてすぐにこれほど人気が​​出るとは誰も予想していませんでした。

UberやAirbnbなどの大企業や、NASAなどの国家機関もこれを利用しています。そして、それらはすべて、最も複雑なプロジェクトで使用されます。

2020 年 11 月現在、TensorFlow は 1 億 6000 万回ダウンロードされています。

しかし、Google は多くのユーザーの気持ちをあまり気にしていないようです。

奇妙なインターフェースと頻繁な更新により、TensorFlow はますますユーザーフレンドリーではなくなり、操作がますます難しくなります。

Google 社内でも、このフレームワークは衰退しつつあると感じています。

実際、Google は頻繁にアップデートする以外に選択肢がありません。結局のところ、これが機械学習の分野における急速な反復に追いつく唯一の方法なのです。

その結果、プロジェクトに参加する人が増え、チーム全体の集中力が徐々に失われていきました。

当初 TensorFlow が選ばれるツールとなった主な特徴は、膨大な数の要因に埋もれてしまい、もはや真剣に受け止められなくなっています。

この現象は、Insider によって「猫とネズミのゲーム」と表現されました。会社は猫のようなもので、継続的な反復で発生する新しい要求はネズミのようなものです。猫は常に警戒し、いつでもネズミに襲い掛かる準備ができていなければなりません。

このジレンマは、市場に最初に参入する企業にとっては避けられないものです。

たとえば、検索エンジンに関して言えば、Google が最初ではありませんでした。そのため、Google は先行企業 (AltaVista、Yahoo など) の失敗から学び、それを自社の開発に応用することができました。

残念ながら、TensorFlow に関しては、Google が行き詰まっています。

まさに上記の理由により、もともと Google で懸命に働いていた開発者たちは、次第に古巣への信頼を失っていったのです。

かつては広く普及していた TensorFlow は徐々に衰退し、Meta の新星である PyTorch に敗北しました。

2017 年に、PyTorch のベータ版がオープンソース化されました。

2018年、FacebookのAI研究ラボはPyTorchのフルバージョンをリリースしました。

PyTorch と TensorFlow はどちらも Python ベースで開発されていますが、Meta はオープンソース コミュニティの維持に重点を置き、多くのリソースを投資していることも注目に値します。

さらに、Meta は Google の問題に気づき、同じ過ちを繰り返すべきではないと考えました。彼らは少数の機能に焦点を合わせており、それらの機能を非常にうまく実行しています。

Meta は Google の足跡をたどらなかった。このフレームワークは、最初は Facebook で開発され、徐々に業界のベンチマークとなってきました。

機械学習のスタートアップ企業の研究エンジニアは、「私たちは基本的にPyTorchを使用しています。そのコミュニティとオープンソースは最高です。質問に対する答えがあるだけでなく、示されている例も非常に実用的です」と語った。

この状況に直面して、Google の開発者、ハードウェアの専門家、クラウド プロバイダー、Google の機械学習に関係するすべての人がインタビューで同じことを言いました。彼らは、TensorFlow が開発者の心を失ってしまったと考えています。

一連の公然かつ秘密裏の闘争の後、Meta はついに勝利を収めました。

一部の専門家は、Google が将来も機械学習をリードし続ける可能性は徐々に失われつつあると指摘している。

PyTorch は徐々に、一般の開発者や研究者が選ぶツールになってきました。

Stack Overflow が提供するインタラクション データから判断すると、開発者フォーラムでは PyTorch に関する質問が増えている一方で、TensorFlow に関する質問は近年停滞しているようです。

記事の冒頭で触れた Uber のような企業も PyTorch を導入しています。

実際、PyTorch のその後のすべてのアップデートは、TensorFlow に対する侮辱のように思えます。

Google における機械学習の未来 — JAX

TensorFlow と PyTorch が激しく競争する中、Google 内の「小さなダークホース研究チーム」が、TPU をより便利に利用できる新しいフレームワークの開発に取り組み始めました。

2018 年、Roy Frostig、Matthew James Johnson、Chris Leary によって執筆された「高レベル トレーシングによる機械学習プログラムのコンパイル」という論文によって、JAX プロジェクトが注目を浴びました。

左から右にこの三人の偉大な神々がいます

その後、PyTorch のオリジナル作者の 1 人である Adam Paszke も、2020 年初頭に JAX チームにフルタイムで参加しました。

JAX は、機械学習における最も複雑な問題の 1 つであるマルチコア プロセッサのスケジュール設定に対して、より直接的なアプローチを提供します。

アプリケーションに応じて、JAX は 1 つのチップを単独で動作させるのではなく、複数のチップを自動的に小さなグループに結合します。

これの利点は、できるだけ多くの TPU に瞬時に応答できるため、「錬金術の世界」が活性化されることです。

最終的に、肥大化した TensorFlow と比較して、JAX は Google 内の大きな問題、つまり TPU にすばやくアクセスする方法を解決しました。

以下は、JAX を構成する Autograd と XLA の簡単な紹介です。

Autograd は主に勾配ベースの最適化に使用され、Python コードと Numpy コードを自動的に区別できます。

ループ、再帰、クロージャを含む Python のサブセットを処理するために使用でき、導関数の導関数を微分するためにも使用できます。

さらに、Autograd は勾配の逆伝播をサポートしています。つまり、配列値パラメータに関するスカラー値関数の勾配、順方向微分、およびこれら 2 つの任意の組み合わせを効率的に取得できます。

XLA (Accelerated Linear Algebra) は、ソース コードを変更せずに TensorFlow モデルを高速化できます。

プログラムの実行中、すべての操作は実行者によって個別に実行されます。各操作には、エグゼキュータがディスパッチする事前コンパイルされた GPU カーネル実装があります。

例えば：

 モデル関数の定義( x , y , z ):
 tf.reduce_sum ( x + y * z ) を返す

XLA なしで実行すると、この部分は乗算用、加算用、減算用の 3 つのカーネルを起動します。

XLA は、加算、乗算、減算を単一の GPU コアに「融合」することで最適化を実現できます。

この融合操作では、メモリによって生成された中間値が y*z メモリ x+y*z に書き込まれるのではなく、これらの中間計算の結果が完全に GPU 上に保持されたまま、ユーザーに直接「ストリーミング」されます。

実際には、XLA は約 7 倍のパフォーマンス向上と約 5 倍のバッチ サイズの改善を実現できます。

さらに、XLA と Autograd は任意の方法で組み合わせることができ、pmap メソッドを使用して複数の GPU または TPU コアを同時に使用するようにプログラムすることもできます。

JAX を Autograd および Numpy と組み合わせると、CPU、GPU、TPU 向けのプログラミングが簡単で高性能な機械学習システムが実現します。

明らかに、Google は今回の教訓を学んだ。同社は独自のオープンソース エコシステムを全面展開するだけでなく、オープンソース エコシステムの構築促進にも特に積極的に取り組んでいる。

2020年、DeepMindは正式にJAXに参加し、Googleも個人的に関与することを発表しました。それ以来、さまざまなオープンソースライブラリが次々と登場しています。

賈陽青氏は、「公開と秘密の闘争」全体を見て、TensorFlowを批判する過程で、AIシステムはPython的な科学研究だけが必要なのだと信じていたと述べた。

しかし、一方では、純粋な Python では効率的なソフトウェアとハ​​ードウェアの共同設計を実現できず、他方では、上位レベルの分散システムでは依然として効率的な抽象化が必要です。

JAX はより良いバランスを模索しており、Google が自社の実用主義を覆す意欲は学ぶ価値がある。

causact R パッケージと関連するベイズ分析の教科書の著者は、Google が TF からよりクリーンなソリューションである JAX に移行したことを嬉しく思っていると述べました。

Googleの挑戦

新参者として、Jax は先行する 2 つの技術である PyTorch と TensorFlow の利点から学ぶことができますが、後発であることで不利な点が生じることもあります。

まず第一に、JAX はまだ「若すぎる」です。実験的なフレームワークであるため、成熟した Google 製品の標準に到達するにはほど遠いです。

さまざまな隠れたバグに加えて、JAX はいくつかの問題に関して依然として他のフレームワークに依存しています。

たとえば、データの読み込みと前処理では、ほとんどのセットアップを処理するために TensorFlow または PyTorch を使用する必要があります。

明らかに、これは理想的な「ワンストップ」フレームワークからは程遠いものです。

第二に、JAX は TPU に対して高度に最適化されていますが、GPU と CPU ではさらに劣ります。

一方、2018年から2021年にかけてのGoogleの組織的、戦略的な混乱により、GPUサポートの研究開発への資金が不足し、関連する問題への対処の優先順位が低くなりました。

同時に、自社のTPUがAIアクセラレーションのシェアをより大きく獲得できるようにすることに集中しすぎたためか、GPUのサポート強化などの詳細については言うまでもなく、NVIDIAとの連携は当然ながら非常に乏しいものとなった。

一方、Google 自身の社内研究は間違いなく TPU に重点を置いており、そのせいで Google は GPU の使用に関する適切なフィードバック ループを失ってしまいました。

さらに、デバッグ時間が長くなること、Windows との非互換性、副作用を追跡できないリスクなどにより、Jax の使用限界と使いやすさが高まります。

現在、PyTorch はほぼ 6 年が経過していますが、TensorFlow が過去に示していたような衰退は見られません。

ジャックスが追いつきたいなら、まだ長い道のりが残っているようだ。