チューリヒ工科大学の研究チームが開発した「UltraFastBERT」は、言語モデルの効率を大幅に向上させる技術だ。この技術は、BERTモデルの一種であり、従来のフィードフォワード層を高速フィードフォワードネットワークに置き換えることで、推論時の計算効率を飛躍的に高めている。その結果、最適化されたベースラインのフィードフォワード実装に比べて、驚異的な78倍の速度向上を実現した。
言語モデルの規模が拡大し、数十億のパラメータを持つようになると、多様なタスクでのパフォーマンスが向上する一方で、計算コストの増大が実世界のアプリケーションにおける大きな障壁となっている。この課題に対処するため、研究者たちは大規模言語モデル（LLM）の効率を高める方法を模索している。
これらの大規模言語モデルの大部分のパラメータは、フィードフォワード層に存在するが、推論時にすべてのニューロンがアクティブである必要はない。この点に着目し、計算効率を最適化する機会がある。
「Exponentially Faster Language Modelling」と題された新しい論文で、チューリヒ工科大学の研究チームはUltraFastBERTを紹介している。このモデルは、4095個のニューロンを持ち、推論時にはわずか12個（0.03%）のみを使用する。標準的な下流タスクでのファインチューニング後、UltraFastBERTはBERTのカウンターパートと同等のパフォーマンスを発揮する。
UltraFastBERTのアーキテクチャは、cramming BERTから着想を得ているが、中間層のフィードフォワードネットワークを高速フィードフォワードネットワークに置き換えることで差別化している。これらのネットワークに適用された簡素化の変更には、一様な活性化関数、すべてのノードの出力重み、出力バイアスの除去、リーフサイズを1に固定すること、および複数のFFFツリーを並行して使用することが含まれる。
トレーニングフェーズでは、cramming BERTの最終トレーニング手順に従い、評価フェーズでは、GLUEベンチマークのさまざまなタスクでUltraFastBERTモデルをファインチューニングする。注目すべきは、1日間のA6000 GPUでのトレーニングのみで、UltraFastBERTのバリエーションが元のBERTbaseモデルのGLUE下流予測パフォーマンスの少なくとも96.0%を維持していることである。UltraFastBERT-1×11-longは、使用するニューロンがわずか0.3%であるにもかかわらず、元のBERT-baseモデルと同等のパフォーマンスを発揮する。
さらに、このチームは、最適化されたベースラインのフィードフォワード実装に比べて78倍の速度向上を示す高レベルのCPUコードを提供しており、PyTorch実装では、同等のバッチ処理されたフィードフォワード推論に比べて40倍の速度向上を実現している。
研究者らによると、このアプローチは、特に非常に大規模な言語モデルにおいて、非常に大きなスピードアップの可能性を秘めているという。例えば、OpenAIの言語モデルGPT-3では、理論上、1回の推論に必要なニューロンの数を従来のわずか0.03％まで減らすことができる。
理論上の可能性を実際に活用するには、ソフトウェアレベルでの最適化が必要である。科学者たちによれば、条件付き行列乗算を効率的に実装することで、処理を341倍高速化することが可能になるという。
しかし、必要な知識は自由に手に入るものではない：
密行列乗算は、コンピュータ史上最も最適化された数学演算である。これを可能な限り高速に実行するために、メモリ、チップ、命令セット、ソフトウェア・ルーチンの設計に多大な努力が払われてきた。これらの進歩の多くは、その複雑さのためであれ、競争上の優位性のためであれ、秘密にされ、強力だが制限の多いプログラミング・インターフェースを通してのみエンドユーザーに公開されてきた。そのため、新しいハードウェアを必要としないにもかかわらず、CMMを実装するためには高水準の線形代数ルーチンを組み合わせることを余儀なくされ、その結果、高速性が低下している。
その第一歩として、研究者たちは、最適化されたベースラインのフィードフォワード実装に比べて78倍のスピードアップを達成する実用的なCPUレベルのコードを書いた。
研究者は画像モデルの分野でもブレークスルーを果たし、必要な計算時間を大幅に短縮した。Stability AIのSDXL Turboは、ほぼ同じ品質を維持しながら、画像生成のステップ数を50から1に減らしている。
論文
参考文献
研究の要旨
言語モデルは、個々の推論にニューロンの指数関数的な割合しか使用する必要がない。その証明として、類似のBERTモデルと同等の性能を発揮しながら、推論中に0.3%のニューロンを使用するBERT変種であるUltraFastBERTを紹介する。UltraFastBERTは、各層の推論に4095個のニューロンのうち12個だけを選択的に使用する。これは、フィードフォワードネットワークを高速フィードフォワードネットワーク（FFF）に置き換えることで達成される。現在、条件付きニューラル実行の高速化の可能性を完全に引き出す真に効率的な実装は存在しないが、最適化されたベースラインフィードフォワード実装に対して78倍の高速化を達成する高レベルCPUコードと、同等のバッチフィードフォワード推論に対して40倍の高速化を達成するPyTorch実装を提供する。学習コード、ベンチマーク設定、モデル重みを公開します。
コメントを残す