Julia は誕生した瞬間から科学計算の分野に目を向け、それ以来ずっと密かに Python と競争してきました。

ニューラルネットワークフレームワークに関して言えば、Python には PyTorch と TensorFlow があり、これらはディープラーニング開発でほぼ推奨されるフレームワークであり、Meta や Google から技術的および財政的支援を受け、急速な発展につながっています。

JuliaにもFlux.jlフレームワークがありますが、Juliaコミュニティは常に言語自体の高いパフォーマンスに頼って生産性を高めてきました。そのため、Flux.jlのコードサイズはPythonフレームワークと比較して特に「軽量」と言えます。例えば、PyTorchとTensorFlowは独立した言語とコンパイラ（Torchscript、XLAなど）を完全に含んでいますが、Flux.jlはJulia言語のみで記述されています。

もちろん、世の中にタダ飯はありません。別の視点から見ると、機械学習の分野でシンプルかつ汎用的で高性能なフレームワークを開発することはほぼ不可能です。私たちにできるのは、常にトレードオフを繰り返すことだけです。

たとえば、特定の問題に対して、スパースで小さなモデルが必要な場合、最高のパフォーマンスを実現する最善の方法は、一般的なフレームワークを使用するのではなく、それを書き直すことです。

最近、Julia コミュニティは SimpleChains.jl という新しいフレームワークをオープンソース化しました。これは、小規模なモデル シナリオでは PyTorch よりも少なくとも 5 倍高速です。

コードリンク: https://github.com/PumasAI/SimpleChains.jl

開発者は、このフレームワークはすべての人に役立つわけではないが、必要とする人にとっては非常に役立つと述べています。

あるネットユーザーはこれに強く同意し、「タスクによって必要なツールは異なる」と述べました。TensorFlowとPyTorchはメモリを大量に消費し、インプレース演算が不足しているため、小規模なモデルでは時間の無駄になるからです。数年前、Netflixで働いていた彼は、Vectorflowと呼ばれるD言語フレームワークを設計・開発しました。このフレームワークは現在、GitHubで1200個のスターを獲得しています。

機械学習モデルの仮定

SimpleChains.jlは、Pumas-AIとJulia ComputingがRoche社およびメリーランド大学ボルチモア校と共同で開発したライブラリです。その主な目的は、小規模ニューラルネットワークに可能な限り最高のパフォーマンスを提供することです。

SimpleChains.jl は当初、科学的機械学習 (SciML) のソリューションとして医療データ分析に使用されていました。小さなニューラル ネットワーク (およびフーリエ シーケンスやチェビシェフ多項式展開などの他の近似値) を既知の半生理学的モデルと組み合わせて、これまで知られていなかったメカニズムや予後因子を発見することができます。

ブラックホールのダイナミクスから耐震建築物の開発まで、SciML メソッドの有効性は多くの分野で実証されており、(生物)物理学的方程式の柔軟な発見/誘導を可能にしています。

アプリケーション シナリオが大きく変化する場合は、モデルのパフォーマンスを向上させるために、特殊なニューラル ネットワークの使用が必要になることがあります。

具体的には、機械学習モデルの研究では、ニューラルネットワークが十分に大規模であるため、行列乗算（畳み込みなど）にかかる時間コストO(n^3)が実行時間の大部分を占めるという仮定が一般的に用いられています。これは、機械学習ライブラリのほとんどのメカニズムの背後にある4つの指針の1つでもあります。

1. 行列乗算の計算量は3乗であるのに対し、メモリ割り当ての規模は線形です。したがって、ベクトル演算に非割り当てメモリを使用することは、優先度が高くありません。

2. 現在のAIアクセラレーションの取り組みは、主にGPUカーネルアクセラレーションに焦点を当てており、命令実行を可能な限り高速化しています。これらの大規模な行列演算はGPU上で最も高速であり、大規模モデルの主要なボトルネックでもあるため、パフォーマンスベンチマークは基本的にこれらの特定のカーネルの速度のみを測定しています。

3. 自動微分バックプロパゲーションを実行する場合、メモリ割り当てが大きなカーネル呼び出しによって隠されるため、値をメモリにコピーする操作はほとんど認識されません。

4. ユーザーは任意のテープに書き込んでバックプロパゲーションを生成することができます。これによりフォワードプロパゲーション中に辞書を構築するコストが増加しますが、カーネル呼び出しの規模を縮小することでコストを抑制できます。

しかし、これらの仮定はすべて、現実世界のケースでも本当に当てはまるのでしょうか?

それがうまくいかない場合、より高いコンピューティング パフォーマンスを実現するために、これらの領域の改善に重点を置くことは可能でしょうか?

小規模ニューラルネットワークのボトルネックは何ですか?

初心者の場合は、まず Julia コードを使用してメモリ割り当て時間や GPU 計算時間などをテストし、仮定 1 と 2 をテストできます。

ご覧のとおり、100x100*100x100 などの大規模な行列乗算演算を実行する場合、メモリ割り当てによって発生するオーバーヘッドは基本的に無視できます。

しかし、低速域ではかなり大幅なパフォーマンス向上が見られることも分かります。これらの利点は、Julia LoopVectorization.jl のみを使用することで実現されます。これは、標準的な BLAS ツールではこの領域で追加のスレッドオーバーヘッドが発生することが多いためです（繰り返しますが、この領域では最適化は実行されません）。

これまでGPUを詳細に調べずに利用してきた方にとって、この事実は驚くかもしれません。GPUは多数のコアを搭載した低速チップとして設計されているため、大規模な行列乗算などの高度な並列処理にしか効果を発揮しません。この点から、仮定2は大規模なネットワーク処理と見なすことができます。

ただし、小規模ネットワークの場合、並列コンピューティングが不足しているため、GPU コアを使用した場合のパフォーマンスは、適切に設計された CPU コアほど良好ではない可能性があります。

行列演算は、バッチ処理が利用可能な場合にのみ実行されます (A*B の行列 B の各列は個別のバッチです)。

常微分方程式の隣接ベクトルのヤコビ積の計算など、多くの科学的機械学習シナリオでは、この演算は行列とベクトルの乗算です。これらの演算の時間計算量はわずかO(n^2)ですが、この場合、メモリのオーバーヘッドが増大する可能性があります。

ニューラル ネットワークの基本的な演算はシグマであるため、O(n) の時間計算量を伴う演算もあり、これによりメモリのオーバーヘッドがさらに深刻になります。

仮説3と4については、バックプロパゲーションの実装を詳しく検討する必要があります。機械学習ライブラリによって自動微分化の手法も異なります。勾配値を即座にバックプロパゲーションするライブラリもあれば、勾配を保存する必要があり、追加のメモリオーバーヘッドが発生するライブラリもあります。

特定のアプリケーションでは、勾配が即座に伝播することが分かっている場合、瞬時に計算できます。一般的な実装と比較して、問題を解くために必要なのはキャッシュされたベクトルのみで、値はその場で代入できます。これにより、自動微分化に伴う余分なオーバーヘッドがすべて排除されます。

これらのアイデアに基づき、研究者たちはSimpleChains.jlをオープンソース化しました。これは、この種の最適化問題を効果的に解決できます。CPU上で小規模なモデルを迅速にフィッティングし、最適化することができます。初期のニューラルネットワークのプロトタイプモデル設計は、主に以下のことを目的としていました。

1. より良いパフォーマンスを実現し、理想的には CPU のピーク FLOP に到達すること。

2. 開発の初期段階では小規模モデルに重点を置き、大規模モデルに対するカーネル最適化操作（キャッシュタイリングなど）の一部を中止します。

3. ベクトルのパラメータと勾配がファーストクラスである API があり、さまざまなオプティマイザーやソルバー (BFGS など) での作業が容易になります。

4. 「純粋な Julia」で記述されているため、開発と最適化が容易です。LoopVectorization.jl を多用しながらも、SimpleChains.jl は BLAS または NN ライブラリに依存しません。

開発者の長期的な目標は、このループコンパイラ最適化アプローチを拡張し、プルバックを自動生成することです。しかし、このコンパイラ中心のアプローチは実装の容易さから既に利用されています。勾配を手作業で記述する必要はありますが、手動で最適化する必要はありません。

SimpleChains.jl は実際にはどの程度うまく機能しますか?

研究者たちは、2×2行列を用いた実験を、AVX512命令セットを搭載したIntel i9-10980XEプロセッサ上で実行しました。10,000エポックの実行に要した時間は0.41秒でしたが、pyTorchでは15秒でした。つまり、この小型ニューラルネットワークでは、約35倍の高速化が実現したことになります。

実験を AVX2 命令を備えた AMD EPYC 7513 マシンに移したところ、Julia 実装では 0.72 秒かかりましたが、PyTorch 実装では 70 秒かかり、その差は 100 倍になりました。

研究者らはAMD Ryzen 9 5950XでもJaxコードをテストしました。Juliaでは1.3秒かかりましたが、Jaxでは14秒かかり、約10倍の短縮となりました。

Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz プラットフォームに切り替えると、Jax では 9 秒かかりましたが、Julia では 0.4 秒かかり、約 22 倍の改善が見られました。

やや性能の劣るプロセッサ（6 コア CPU）に切り替えると、Jax では 19 秒かかりますが、Julia では 9 秒かかり、速度は 2 倍しか向上しません。

もう少し大きくて実際に使えるニューラル ネットワークでも、トレーニング速度にこれほど大きな違いが出るでしょうか?

研究者たちはLeNet5を用いてMNISTをテストしました。この例は、従来の機械学習のユースケースではバッチ処理で行列乗算が使用される可能性があるため、非常に控えめな速度推定値に過ぎません。しかし、この場合でも、比較的小規模なネットワークサイズにより、大幅なパフォーマンスの向上が見られます。

バッチサイズ2048で10エポックのトレーニングを行ったところ、A100マシンではトレーニング時間は17.66秒、17.62秒で、精度はそれぞれ94.91%、96.92%でした。V100マシンでは、トレーニング時間は16.29秒、15.94秒で、精度はそれぞれ95.6%、97.5%でした。

しかし、この問題はGPUにとってはまだ過剰です。バッチサイズが2048でも処理速度は非常に速く、CPUからGPUへの転送に主に時間がかかります。

さらに 2 つの実験が AMD EPYC 7513 と Intel i9 10980XE で実施され、その結果、GPU よりも高速で正確であることが示されました。

SimpleChains.jl に切り替えると、AMD プラットフォームでは所要時間が 3 秒で精度は 98.3% でした。Intel プラットフォームでは所要時間はわずか 1 秒で精度は 98.2% でした。ラップトップの Intel プラットフォームでも所要時間はわずか 5.3 秒で精度は 97% でした。

現在、大規模な機械学習フレームワークは、99.9% のユーザーに最高のパフォーマンスを提供するのに優れていますが、小規模なモデルを使用する残りの 0.1% のユーザーには適していません。

これは、構成可能性と柔軟性の利点です。つまり、機械学習フレームワークを簡単に構築できる言語であると同時に、それらの代替フレームワークに最適化された代替フレームワークを構築できる言語でもあります。