- 長いコンテキストのTransformerで self-attentionのコスト がボトルネックになる状況に対し、SPECTREはFFTベースのトークンミキサーにより、レイヤーあたりの複雑度を O(L²) から O(L log L) に削減
- 各attention headは、高速なreal FFT、コンテンツ適応型スペクトルゲート、inverse FFT の組み合わせに置き換えられ、既存のTransformer構造は維持される
- 自己回帰生成では Prefix-FFT cache により、各ステップでのFFT再計算の負担を軽減し、選択式waveletモジュールでローカル特徴の損失を補える
- Llama-3.2-1BバックボーンでSDPA、FlashAttention-2、SPECTREを比較し、NVIDIA A100-80GBで512〜128kトークンのスループットとレイテンシを測定
- SPECTREはPG-19とImageNet-1kでベースライン性能と同等以上を示し、6%未満のパラメータ追加で一般的なGPUにおける長文コンテキスト処理を目指す
Self-attentionの二次コストをFFTで削減する方法
- 長いコンテキストのTransformerは、multi-turn dialogue、書籍長の要約、高解像度ビジョンのような 数万トークン を扱うタスクで必要とされる
- 従来のself-attentionは O(n²d) のコストのため、コンテキストが長くなるほど推論レイテンシとメモリ使用量が増大する
- SPECTREはself-attentionレイヤーを 周波数領域トークンミキサー に置き換えるdrop-inな代替方式である
- トークンを orthonormal Fourier basis に射影する
- コンテンツ適応型の対角ゲートと選択式low-rankゲートを適用する
- inverse transform によって再びトークン空間へ戻す
- 周辺のネットワークアーキテクチャを変更せずに、レイヤーあたりの複雑度を O(n log n) に下げることが中核である
トークンミキサーの構成と生成対応
- SPECTREのattention head置き換え構成は、高速なreal FFT、spectral gate、inverse FFT である
- spectral gating は n/2 + 1 個の周波数係数上で動作し、計算量とメモリ使用を抑えつつ表現力を維持するよう設計されている
- Prefix-FFT cache は標準的なKV-cacheに近い役割を果たし、ストリーミングデコードを支援する
- 自己回帰生成で各 time step ごとにFFTを再計算しなければならない従来のspectral mixerの弱点を軽減する
- 固定メモリ予算内で効率的な生成を可能にする構造である
- 選択式 Wavelet Refinement Module は、純粋なspectral方式で失われうるローカルな詳細を補完し、計算オーバーヘッドは小さい
既存Transformerへの適用方法
- SPECTREはmulti-head attentionレイヤーを直接置き換えられるため、追加のアーキテクチャ改変を必要としない
- 既存の事前学習モデルはSPECTREレイヤーでfine-tuningできる
- 更新対象は新たに導入されたパラメータである
- 追加パラメータは全重みの 6%未満 である
- specialized optimization や非標準アーキテクチャを必要とする手法と異なり、周辺のTransformer構造を維持する
Llama-3.2-1Bベースの実験
- 同一のLlama-3.2-1Bバックボーンに3種類のattention kernelを適用して比較した
- standard softmax-dot-product attention(SDPA)
- FlashAttention-2
- SPECTRE mixer
- 測定環境は NVIDIA A100-80GB、シーケンス長は L ∈ {512, 1k, 4k, 8k, 32k, 128k}
- 指標は tokens-per-second のスループットと single-batch latency
- スループットは高いほどよい
- latency は低いほどよい
- SPECTREはバックボーンの精度を維持しながら、ほぼ O(n log n) に近い実行時間を示した
- 32kトークンまでは実行時間がほぼフラットに保たれる
- abstract基準では、128k-token context でFlashAttention-2より最大 7×高速
- 本文のcontribution一覧基準では、32kトークンでFlashAttention-2より最大 7×高速な推論 を示す
ベンチマーク結果と実用範囲
- SPECTREは PG-19 言語モデリングと ImageNet-1k 分類で、ベースライン性能と同等以上の結果を示した
- 長いコンテキスト処理でself-attentionの二次コストを回避しつつ、グローバルなコンテキスト混合を維持する
- sparse pattern、kernel approximation、low-rank structure ベースのattention高速化手法は、厳密性の犠牲、非標準最適化、ストリーミング生成非対応といった限界を持つことがある
- SPECTREは、FFTがcircular convolutionを対角化し、グローバル混合をelement-wise productへ変換する周波数領域アプローチを用いる
- 追加パラメータを6%未満に抑え、specialized hardwareなしで 汎用GPU 上の hundred-kilotoken context 処理を目指す
1件のコメント
Hacker News のコメント
基本的には畳み込み定理を活用する方式。元の空間では高コストな畳み込みが、双対空間では単純な乗算になり、その逆も成り立つ。
データに畳み込み演算があるなら、共役領域に変換して乗算に置き換えればよい。
別の言い方をすれば、データにとって自然な領域で作業せよ、ということ。
https://en.wikipedia.org/wiki/Convolution_theorem
最適化するには各ステップを最適化し、可能な限り最も効率のよい空間で多く作業すればよい。
単に計算が簡単だというのとは別の話ではないのか？
そのうえ計算に複素数を使う必要があり、数値的にも安定性が低い。私の知る限り、FFT は一般的な畳み込みではメリットがない。
自己アテンションやこの論文の用途では、n ははるかに大きいかもしれない。論文は読んでいない。それでも複素数の問題は残る。
Google は 2022 年に FNet: Mixing Tokens with Fourier Transforms でこのアイデアを導入した。
その後、ほとんどの状況で TPU の行列乗算性能のほうが FFT より速いことが分かった。
https://arxiv.org/abs/2105.03824
「全体として、FNet、Performer、スパース Transformer のようなアプローチは、固定または近似的なトークン混合によって計算負荷を減らせることを示しているが、われわれの適応型スペクトルフィルタリング戦略は、FFT の効率性と、学習可能で入力依存のスペクトルフィルタを独自に組み合わせている。これは複雑なシーケンスモデリングタスクに重要な、スケーラビリティと適応性の強力な組み合わせを提供する。」
その後に比較セクションもある。
ところで DSP には FFT を支援する専用ハードウェアがあるのだろうか？ 純粋に気になって聞いている。使ったことはないが、なんとなく役に立ちそうに思える。
「TPU はフーリエ変換ではあまりに非効率なため、研究者らは 4096 未満のシーケンスでは FFT アルゴリズムを使わず、事前計算した DFT 行列を使う二次スケーリングのフーリエ変換実装を選んだ。」
「Nvidia Quadro P6000 GPU では、FNet アーキテクチャにおいてフーリエ変換が推論時間の最大 30% を占めた。」
この会社は 2021 年に、Google が TPU に同社の光チップを使えば推論時間を 40% 削減できると主張していた。FFTNet がさらに多くを担うようになれば、もっと削減できるかもしれない。
https://scribe.rip/optalysys/attention-fourier-transforms-a-...
フーリエ変換は「トークン」次元に沿って適用される。しかし多くの応用では、この次元に意味がない。そのため Transformer は順列不変データを扱うのに良い選択肢になる。
あまり知られていない有限群上のフーリエ変換を使った追加実験を見てみたい。これは順列不変でありながら、標準的なフーリエ変換と多くの性質を共有している。
また、これが LLM の次の大きな流れになるなら、vLLM や llama.cpp のような推論エンジンがどれほど簡単に統合できるのかも気になる。
https://en.wikipedia.org/wiki/Fourier_transform_on_finite_gr...
llama は入力内の位置に応じてベクトルに回転を適用していると理解している。
数学は完全に頭の上を通り過ぎていて、数式まわりの説明もかろうじて理解できる程度です。誰か簡単な言葉で、これがどうやってアテンション機構と同等なのか説明してくれませんか？
ここで言う周波数とは何で、トークン間の位置関係はどうエンコードするのでしょうか？
信号解析や画像では最も直感的です: https://homepages.inf.ed.ac.uk/rbf/HIPR2/fourier.htm
周波数空間は本質的に複素数で表される「複素」空間です。周波数には、問題を大域的に見ることができるという利点があります。
このメカニズムはアテンション機構と同等ではなく、明確なトレードオフがあります。ただし、アテンションが捉える重要な関係のかなりの部分を捉えられる可能性はあります。
modReLUについては今すぐよい直感はありませんが、周波数を修正しつつ逆フーリエ変換を保つため、重要なものに見えます。
ここでの「周波数」は、おそらくかなり抽象的なものである可能性が高いです。FFTは、明確な周波数解釈がない方法でもよく使われます。畳み込み定理のような便利な数学的性質のために使われる場合が多いです。
本当にうまく動くならかなり驚きですし、とてもエレガントです。
つまり、出力の各ベクトルは、その入力ベクトルが他のすべての入力ベクトルの何らかの関数によって変換されたものに依存します。
https://medium.com/optalysys/attention-fourier-transforms-a-...
概念的には、これが少し単純化された畳み込みとどう似ているかを見ることができます: https://openreview.net/pdf?id=8l5GjEqGiRG
畳み込みは、何らかの形で大域的な状態を考慮したいときによく使われます。
このフレームワークに因果マスキングを入れるには、n個の異なるFFTを行う必要がありそうですが、位置埋め込みへの言及もありません。
なので比較対象の自己アテンション実装は非因果のNoPEのように見えますし、そうだとするとベースラインを意図的に弱く設定した例なので、あまり印象的ではないかもしれません。
結果が最先端に近かったなら、著者はおそらく言及していたはずです。
関連文献のようです: https://arxiv.org/abs/2111.13587
Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, Bryan Catanzaro
ここで周波数領域で見ることがなぜ役に立つのか、直感があるのか気になります。
直流成分は理解できますが、入力データが、ほかの周波数に意味が生じるほど十分に周期的だとは期待していません。
数年前にすでにO(n log n)の全文脈混合を示していたHyena Operatorの先行研究が言及されていないようです。
https://arxiv.org/abs/2302.10866
https://arxiv.org/abs/2111.00396
ビッグオー記法はある程度感覚はつかめるものの、コンピューターサイエンスや電気工学に関するたいていの内容と同じく、これも頭の上を通り過ぎていく感じです。
数学が本当に苦手な身としては、こういう内容を理解したり、少なくとも学んで工学の学位や資格まで取れる人たちがうらやましいです。
FFTについて知っているのは、信号を変換し、ある種の信号処理に使われ、昔は核爆発検知の中核だったと聞いたことがある、という程度です。
基本的な考え方はこうです。ほとんどすべての有用な信号は、異なる周波数と位相を持つサイン波の和として表せます。たとえば電気信号や音波は、x軸が時間である1次元信号です。見た目には扱いにくい複雑なうねうねした線かもしれません。
フーリエ変換を使うと、時間ベースの信号に含まれる個々の周波数を分離できます。そのうえで、特定の周波数を望む形に修正できます。たとえば信号にランダムな尖ったノイズが多い場合、それは高い周波数として現れます。整えるにはフーリエ変換を行い、あるしきい値より高い周波数のデータを捨ててから、残ったデータに逆フーリエ変換を適用し、元の信号のより滑らかなバージョンに戻せばよいです。これはローパスフィルターと呼ばれ、元の信号の移動平均を取るのとかなり似ています。
面白いのは、これをかなり直感的に高次元へ拡張できる点です。x軸とy軸がどちらも空間である2次元信号は画像です。JPEG圧縮はこの概念に基づいています。画像をより小さく保存するために高周波信号を取り除き、その代償として細かなディテールを失ったり、捨てすぎるとリング状のアーティファクトが生じたりします。ここに時間という3つ目の次元を加えると動画になり、さらに拡張し続けられます。
これらはすべて視覚的に理解しやすいので、数学をすべて深く知らなくてもよい直感を得られます。可視化とインタラクティブな例が多い良いページです: https://www.jezzamon.com/fourier/index.html
3Blue1Brownの動画も説明が上手です: https://youtu.be/spUNpyF58BY?si=dz0z-s8NftW3Htun
FFTが離散版であるフーリエ変換は、その1次元の時間領域信号を、周波数ごとの大きさと位相成分に分解します。
周波数は基本的には音の高さです。純粋なサイン波、つまり純音は、昔、深夜にテレビ放送が終了するときに聞こえた音に似ています。この場合、ほとんどは0で、その音の周波数の位置に1つの「スパイク」が生じます。信号の振幅が大きいほど、スパイクの大きさも大きくなります。音高、つまり周波数が上がったり下がったりすると、このスパイクの位置が横軸に沿って上下に動きます。
位相は基本的には信号の時間オフセットです。何らかの形で遅延したトーンは別の位相として現れます。ただしこれは絶対的な測定ではなく相対的な測定です。単位がラジアン、つまり角度なので、円を一周するとまた「リセット」されます。そのため、信号が1秒遅れたのか2秒遅れたのかといったことは分かりません。
つまり、1つの信号、すなわち時間に対する振幅から、実際には周波数に対する大きさと位相という2つの情報が得られます。
虚数や複素変数を理解しているなら、この2つの信号は実のところ、複素関数であるFFT出力の大きさと偏角にすぎません。
テレメトリーの時代に、クラウドテレメトリーへFFTを適用して周期的な異常や準安定なシステムを、事故が起きた後ではなく前に見つけないのは、大きな機会を逃しているように思います。
残念ながら、これは私が気づける程度のものではありますが、実装できる技術レベルにはなく、すでに予定もぎっしりです。
「SLAはサービスのデプロイ後23〜25分後に最も破られやすい。ふむ、なぜだろう……あ、まずい」
冗談はさておき、これが本当に収益につながり得る領域は、トラフィックの周期を予測してサーバーインスタンスを増減させ、コストを削減することです。
こうした作業は個人の時間でやっても会社は絶対に承認してくれないでしょうが、既製品としてパッケージ化すれば会社がすぐ買う類のものです。