スタンフォード大学の研究者が『FlashFFTConv』を導入：長いシーケンスのFFT畳み込みを最適化するための新しい人工知能システム

「スタンフォード大学の研究者が『FlashFFTConv』を導入：長いシーケンスのFFT畳み込みを効率化する新たなAIシステム」

効率的な推論は、機械学習において長いシーケンスを取り扱う上での主要な困難です。最近では、畳み込みがシーケンスモデリングにおいて重要な基本操作となり、言語モデリング、時系列解析、コンピュータビジョン、DNAモデリングなどでの最先端のパフォーマンスをサポートしています。これらの印象的な結果や、安定性の向上、シーケンスの長さが増すにつれてスケーラビリティが向上するなどの追加の利点を考慮しても、畳み込みシーケンスモデルは依然としてTransformersよりも遅いです。

その主な原因は、信頼性のないハードウェアサポートです。シーケンスモデリングにおける畳み込みは、通常、視覚的なアプリケーションで使用される短いフィルタとは異なり、入力シーケンスと同じ長さのフィルタを使用します。高速フーリエ変換（FFT）畳み込みアルゴリズムは、入力uと畳み込みカーネルkの畳み込みを入力と出力の周波数にマッピングすることで計算します。

FFT畳み込みアルゴリズムは漸近的に効率的ですが、現在のアクセラレータ上では壁時計時間が短いです。しかし、システムの技術的進歩により、Transformersは現在のアクセラレータの限界まで到達し、FlashAttention-v2を使用する場合のエンドツーエンドFLOP使用率が72%以上になっています。

• 「浙江大学の研究者がUrbanGIRAFFEを提案し、難しい都市のシーンに対する制御可能な3D認識画像の生成に取り組む」
• MITの研究者たちは「MechGPT」を導入しました：メカニクスと材料モデリングにおいてスケール、学問領域、およびモダリティをつなぐ言語ベースのパイオニア
• 『NVIDIAの研究者たちが、現行のCTCモデルと互換性のあるGPU加速の重み付き有限状態トランスデューサ（WFST）ビームサーチデコーダを導入』

長いコンテキストの機能を提供するために、スタンフォード大学の新しい研究では、現代のアクセラレータ上でFFT畳み込みメソッドを最適化する方法を調査しています。研究者らは、システムの進歩によってFlashAttentionのようなモデルや新しいアテンションアルゴリズムが生まれたように、FFT畳み込みの最適化も新しいアルゴリズムを生み出し、畳み込みシーケンスモデルの品質を向上させると考えています。

FFT畳み込みは短いシーケンスに対して簡単に最適化できます。フィルタのFFTを再利用することが一般的な実践であり、再利用する前にフィルタのFFTを事前計算することが可能となります。したがって、FFT畳み込みはバッチとフィルタに跨る並列処理が可能であり、カーネルフュージョンにより中間の畳み込み出力をSRAMやレジスタにキャッシュすることができます。

1. しかし、チームはシーケンスの長さが増すにつれて2つの主要なボトルネックが発生することを指摘しています。現在のアクセラレータでは、FFT畳み込みは特殊な行列行列乗算ユニットを最適化的に利用していません。
2. 第二に、シーケンスがSRAMに収まりきらないほど長くなると、カーネルフュージョンが失敗し、コストのかかるI/O操作が必要になります。因果関係のためのパディング操作や、実数値の入出力から複素数値のFFT中間生成物への変換も、これらのI/Oコストをさらに増加させる可能性があります。

それに対応して、研究者はFlashFFTConvという革新的なアルゴリズムを提案しています。このアルゴリズムは、FFTを長いシーケンスに最適化するためにモナーク分解を使用します。モナーク分解は、FFTをp個の行列乗算操作の系列として書き直すための方法であり、pの値が大きくなるほど、より小さな行列のためFLOPコストが減少しますが、中間結果を伝達するためにより多くのI/Oが必要となります。したがって、トレードオフが存在します。

この研究では、FLOPコストとI/Oコストをシーケンスの長さに基づいた単純なコストモデルを使用して、GPU上でpの最適化方法を示しています。FlashFFTConvは、より長いシーケンス長でのカーネルフュージョンを容易にし、SRAM内に保持する必要のあるシーケンスの量を減らすことができます。したがって、FlashFFTConvは256文字から400万文字までのシーケンスを容易に処理できます。実数値のFFTアルゴリズムを使用し、入力がゼロパディングされている場合には一部の行列乗算操作をスキップすることで、FlashFFTConvはFFT操作の長さを半分以上短縮することができます。さらに、FFT畳み込みの行列ビューは2つのアーキテクチャの修正を実装するためのシンプルなインタフェースを提供し、畳み込みカーネルが入力シーケンスよりも短い長さで学習する部分畳み込みや、周波数空間でカーネルの一部をゼロにする周波数疎畳み込みなどが簡単に実装できます。両方のアプローチは、Transformersにおける疎な/近似的なアテンションの畳み込みバージョンと考えることができます。

研究者は、FlashFFTConvがFFT畳み込みを加速し、より優れた品質、より効率的な長いシーケンスモデルを実現することを実証しています。

• FlashFFTConvは効率の向上により、畳み込みシーケンスモデルの品質を改善します。同じ計算予算の場合、FlashFFTConvはHyena-GPT-sが2.3ポイント改善したperplexityを実現し、M2-BERT-baseが最大3.3高い平均GLUEスコアを達成することができます。これは、モデルのパラメータを倍増させた場合のパフォーマンス向上と同等です。
• FlashFFTConvはPyTorchと比較して、畳み込みの効率性を最大7.93、メモリの節約を最大5.60向上させます。この効率性は、シーケンス長で4桁以上にわたって保持されます。FlashFFTConvは、FlashAttention-v2に比べて、シーケンス長2K以上において壁時計時間が速くなり、エンドツーエンドのFLOP使用率が最大で62.3%となります。
• FlashFFTConvは長いシーケンスのモデルも実現可能です。FlashFFTConvは、長大なarenaベンチマークのPath-512ジョブ（シーケンス長256K）を完了する唯一のモデルを生み出しました。また、FlashFFTConvは、単一ヌクレオチドの分解能で最長の人間の遺伝子（最大2.3M塩基対）を埋め込む最初のモデルでもあり、部分畳み込みを介してHyenaDNAを4Mのシーケンス長に拡張することができます。

チームは、FlashFFTConvが畳み込みシーケンスモデルのより広範な使用を可能にし、学んだ教訓がよりリソース効率の良いコンピュータアーキテクチャにつながることを期待しています。

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