TensorFlow、危険です! Google自身が放棄している

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

約 166,000 個のスターを獲得し、ディープラーニングの台頭を目の当たりにした TensorFlow は、現在、危うい立場に立たされています。

そして今回、その衝撃は古いライバルである PyTorch からではなく、その新人であるJAXからもたらされた。

AI界隈での最近の白熱した議論の中で、fast.aiの創設者であるジェレミー・ハワード氏も次のように述べている。

JAX が徐々に TensorFlow に取って代わっていることはよく知られています。それは今起こっています（少なくとも Google 内では）。

LeCun 氏は、ディープラーニング フレームワーク間の熾烈な競争が新たな段階に入ったとさえ考えています。

LeCun 氏は、当時は Google の TensorFlow の方が Torch よりも確かに人気があったと語った。しかし、Meta の PyTorch の登場により、その人気は TensorFlow を上回りました。

現在、Google Brain、DeepMind をはじめ、多くの外部プロジェクトが JAX を使い始めています。

代表的な例としては、最近話題のDALL・E Miniが挙げられます。著者はTPUを最大限に活用するために、プログラミングにJAXを使用しました。使用後にため息をついた人もいました。

これは PyTorch よりもはるかに高速です。

Business Insiderによると、JAXは今後数年以内に機械学習技術を使用するすべてのGoogle製品をカバーする予定だという。

この観点から見ると、現在社内で積極的に行われている JAX の推進は、むしろ Google がこのフレームワークに対して開始した「自力救済」キャンペーンに近いと言えます。

JAX はどこから来たのですか?

JAX に関しては、Google は実際に十分な準備ができています。

2018年に、Google Brain の 3 人からなる小規模チームによって構築されました。

研究結果は、「高レベルトレースによる機械学習プログラムのコンパイル」と題された論文に掲載されました。

Jax は高性能数値計算用の Python ライブラリであり、ディープラーニングはその機能の 1 つにすぎません。

創業以来、その人気は高まっています。

最大の特徴はスピードです。

例を見てみましょう。

たとえば、NumPy を使用して行列の最初の 3 つの累乗の合計を計算する場合、計算には約 478 ミリ秒かかります。

JAX を使用すると 5.54 ミリ秒しかかからず、これは NumPy よりも 86 倍高速です。

なぜそんなに速いのですか?理由は多数ありますが、例えば次のようなものがあります。

1. NumPy アクセラレータ。 NumPy の重要性は説明の必要はありません。科学計算や機械学習に Python を使用する場合、NumPy なしでは作業できませんが、NumPy は GPU などのハードウェア アクセラレーションをネイティブにサポートしていません。

JAX のコンピューティング関数 API はすべて NumPy に基づいているため、モデルを GPU や TPU 上で簡単に実行できます。この点は多くの人の注目を集めました。

2.XLA。 XLA (Accelerated Linear Algebra)は、高速化された線形代数、最適化コンパイラです。 JAX は XLA 上に構築されており、JAX の計算速度の上限が大幅に向上します。

3. JIT。研究者は XLA を使用して関数をジャストインタイム コンパイル(JIT)バージョンに変換できます。これは、計算関数に単純な関数修飾子を追加することで計算速度を数桁向上させることに相当します。

さらに、JAX は Autograd と完全に互換性があり、grad、hessian、jacfwd、jacrev などの関数変換による自動微分化をサポートします。逆モードと順モードの微分化をサポートしており、2 つを任意の順序で合成できます。

もちろん、JAX にも欠点はいくつかあります。

例えば：

1. JAX はアクセラレータとして知られていますが、CPU コンピューティングにおけるあらゆる操作に対して完全に最適化されているわけではありません。

2. JAX はまだ新しいため、TensorFlow のような完全な基本エコシステムが形成されていません。そのため、Google によって完成品の形ではまだリリースされていません。

3. デバッグにかかる​​時間とコストは不確実であり、「副作用」も完全には明らかではありません。

4. Windows システムには対応しておらず、上記の仮想環境でのみ実行可能です。

5. データローダーがないため、TensorFlow または PyTorch を使用する必要があります。

…

それにもかかわらず、シンプルで柔軟性があり、使いやすい JAX が DeepMind で最初に人気を博しました。 HaikuやRLaxなど、2020年にリリースされたいくつかのディープラーニングライブラリは、これをベースに開発されました。

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

現在、JAX のオープンソース プロジェクトは GitHub で 18.4k のスターを獲得しており、これは TensorFlow よりもはるかに高い数字です。

この期間中、TensorFlow に取って代わる可能性が高いとの声が多くあったことは注目に値します。

一方では、これは JAX の強さによるものであり、他方では、主に TensorFlow 自体の多くの理由に関連しています。

Google はなぜ JAX に切り替えたのでしょうか?

2015年に誕生したTensorFlowは、かつては非常に人気がありました。発売後、Torch、Theano、Caffeなどの多くの「トレンドセッター」をすぐに追い越し、最も人気のある機械学習フレームワークになりました。

しかし、2017 年に、まったく新しい PyTorch が復活しました。

MetaがTorchをベースに構築した機械学習ライブラリです。使いやすく、理解しやすいため、多くの研究者の間で急速に支持を集め、TensorFlowを上回る傾向さえあります。

対照的に、TensorFlow は頻繁なアップデートとインターフェースの反復により肥大化し、徐々に開発者の信頼を失っていきました。

（Stack Overflowでの質問の割合から見ると、PyTorchは年々増加しているのに対し、TensorFlowは停滞している）

競争の中で、TensorFlow の欠点が徐々に露呈してきました。不安定な API、複雑な実装、高い学習コストなどの問題はアップデートではあまり解決されておらず、むしろ構造が複雑化しています。

対照的に、TensorFlow は「運用効率」などの利点を活用し続けていません。

学術界では、PyTorch が TensorFlow よりも徐々に人気が高まっています。

特にACLやICLRなどの主要なトップカンファレンスでは、近年PyTorchで実装されたアルゴリズムフレームワークが80%以上を占めるようになっています。一方でTensorFlowの使用率は依然として低下傾向にあります。

まさにこの理由から、Google は黙っていられず、JAX を使用して機械学習フレームワークの「優位性」を取り戻そうとしました。

JAX は名目上は「ディープラーニング専用に構築された汎用フレームワーク」ではありませんが、リリース以来、Google のリソースは JAX に傾倒してきました。

一方、Google Brain と DeepMind は、JAX 上に徐々にライブラリを構築しています。

Google Brain の Trax、Flax、Jax-md をはじめ、DeepMind のニューラル ネットワーク ライブラリ Haiku や強化学習ライブラリ RLax はすべて JAX に基づいて構築されています。

Googleによると:

JAX エコシステムが発展するにつれて、その設計が Sonnet や TRFL などの既存の TensorFlow ライブラリと (可能な限り) 一貫性を保つようにすることも検討されます。

一方、JAX をベースにしたプロジェクトも増えてきており、最近話題の DALL·E ミニプロジェクトもその 1 つです。

JAX は Google TPU の利点をより有効に活用できるため、PyTorch よりもはるかに優れた実行パフォーマンスを実現し、以前は TensorFlow で構築されていた産業プロジェクトも JAX に切り替えています。

ネットユーザーの中には、JAX が最近とても人気がある理由について冗談を言う人もいました。TensorFlow ユーザーはこのフレームワークを本当に我慢できないのかもしれません。

では、JAX は TensorFlow に取って代わり、PyTorch と競合する新たな勢力となる見込みがあるのでしょうか?

どのフレームワークを好みますか?

全体的に、多くの人が依然として PyTorch を強く支持しています。

彼らは、Google が毎年新しいフレームワークをリリースするペースを気に入らないようです。

「JAX は魅力的ですが、人々が PyTorch を放棄して JAX を使用するほど革新的ではありません。」

しかし、JAX に対して楽観的な見方をする人もたくさんいます。

PyTorch は完璧だと言う人もいますが、JAX もその差を縮めています。

中には、JAX は PyTorch より 10 倍優れていると称賛し、「Meta がさらに努力を続けなければ、Google が勝つだろう」と言った人もいました。 （手動犬頭）

しかし、誰が勝つか負けるかを気にせず、長期的なビジョンを持っている人は常に存在します。

最高というものは存在しません。より良いものだけがあるのです。最も重要なことは、より多くのプレーヤーと優れたアイデアが参加し、オープンソースが真に優れたイノベーションと同義になることです。

プロジェクトアドレス:
https://github.com/google/jax