✨ 要約🔬 技術概要
あなたは、超スマートなロボットに世界を理解させる方法を教えようとしていると想像してください。現在、最高のロボットは「アテンション(注意)」と呼ばれるツールを使って、テキストのページ全体をスキャンして最も重要な単語を見つけ出す学生のように、あらゆるものを一度に見渡しています。これは短い物語には非常にうまく機能しますが、もしロボットに図書館一館分、あるいは3D映画や複雑な気象マップを与えたら、「アテンション」というツールは圧倒されてしまいます。なぜなら、すべてのデータの一つひとつを他のすべてのデータと比較しなければならず、それには膨大な時間とエネルギーがかかるため、ロボットはクラッシュするか、生き残るためにデータを小さな鍵穴から覗き見る(「パッチ化」と呼ばれるプロセス)しかなくなってしまうのです。
これを解決するために、科学者たちはより高速なツールを構築しようとしてきました。一つの人気のあるアイデアは、テキストを一行ずつ一語ずつ読み進めるようにデータを読み取る「リカレント(再帰的)」モデルを使用することです。しかし、これは3D彫刻を一つの角度からだけ見て、残りの部分を推測しようとするようなものであり、オブジェクトの自然な形状を破壊してしまいます。もう一つのアイデアは、画像の上を滑るように動く虫眼鏡を使ってパターンを見つけ出す「コンボリューション(畳み込み)」です。これらは高速で3Dの形状を維持できますが、通常は、そこに何があるかに注意を払うことなく、あらゆる部分に全く同じパターンを適用する硬直したスタンプのように機能します。コンピュータサイエンスにおける大きな疑問は、「膨大な3Dデータを扱うのに十分速く、データの自然な形状を維持し、かつ、見たものに応じて『虫眼鏡』を変化させることができるほど賢いツールを作れるか?」ということです。
この論文は、HyenaND と呼ばれる新しいツールを紹介しており、それは「イエス」と答えています。HyenaNDを、3Dオブジェクト(医療スキャンや気象シミュレーションなど)の上を、決して1次元の線に平坦化することなく滑っていく、魔法の形を変える虫眼鏡だと考えてください。過去の硬直したスタンプとは異なり、このツールはまずオブジェクト全体を見渡し、どのようなパターンを見つける必要があるかを決定してから、瞬時にレンズを調整して適合させることができます。これは、「フーリエ変換」(画像を音波に変換して高速に処理する方法)を用いた巧妙な数学的トリックを用いることで、データが大きくなってもコストが爆発的に増えるのではなく、ごくわずかな増加に抑えつつ、非常に高い速度を維持します。
著者らは、HyenaNDが複雑で多次元的なデータを扱う上でゲームチェンジャーであることを明らかにしました。テストにおいて、彼らは、標準的な「アテンション」ツールがメモリ制限のためにクラッシュしてしまうようなタスク(3D医療画像の解析や流体力学のシミュレーションなど)を、HyenaNDがいかに処理できるかを示しました。また、ロボットが3Dグリッド内の特定の色の記憶を保持するという「コピー」ゲームでテストしたところ、HyenaNDは古い手法よりも数千倍正確でした。彼らはさらに、この数学的トリックが実際のコンピュータチップ上で電光石火の速さで動作するように、nSubQ と呼ばれる特別なソフトウェアエンジンを構築し、理論上の速度を実際の時間の節約へと変えました。
論文は、HyenaNDは単独でも非常に強力ですが、古い「アテンション」ツールとハイブリッドチームを組むことでさらに効果を発揮すると示唆しています。DNA配列、コンピュータビジョン(ImageNet)、および医療画像を用いた実験において、これらのハイブリッドチームは、純粋なアテンションモデルや、データを直線的に読み取ろうとする他の高速モデルの両方に勝利しました。著者らは、スピードと知能のどちらかを選ぶ必要はもうないのだと主張しています。私たちは、データの3D幾何学を尊重し、大規模なサイズにも対応でき、かつ細部にもしっかりと注意を払うことができるツールを持つことができるのです。彼らは、このアプローチが大きなボトルネックを取り除き、将来のAIが、小さな鍵穴から覗き見ることなく、ついに高解像度のネイティブな3Dの世界を「見る」ことを可能にすると結論付けています。
技術要約: HyenaND および nSubQ
問題提起 ディープラーニングモデルは、高解像度画像、3D医療ボリューム、物理場などの大規模かつ多次元(ND)データをますます処理するようになっています。標準的なTransformerアーキテクチャは、本質的に多次元的ではあるものの、シーケンス長に対して二次的な計算複雑度(O ( L 2 ) O(L^2) O ( L 2 ) )に苦しんでおり、これが高解像度または長コンテキストのタスクへの適用を妨げています。一般的な回避策であるパッチ化(patchification)は、性能を犠牲にして計算可能性を確保するものですが、微細な構造を回復するためにパッチサイズを縮小するとトークン数が増加し、再び二次的なボトルネックが生じます。
既存の劣二次(subquadratic)な代替案は、NDデータへの適用において重大な限界に直面しています:
回帰モデル(例:Mamba): これらは本質的に1次元的かつ因果的(causal)です。NDデータを処理するためには、入力をアドホックな1Dスキャン順序へと「ラスタライズ」する必要がありますが、これはデータの空間構造や幾形的帰納バイアスを破壊してしまいます。さらに、それらの入力依存性は通常トークンレベルであり、近似的に線形時不変(LTI)で、かつ全量(非自己回帰設定)で利用可能な信号に対しては最適ではありません。
標準的な畳み込み: これらはND構造を保持し、高速なFFTベースのパスを提供しますが、従来の長距離畳み込みは入力依存性に欠け、アテンションメカニズムと比較して表現力が制限されます。
手法: HyenaND 著者らは、ラスタライズを行うことなく、多次元データのネイティブな幾何学に対して直接作用するように設計された、劣二次的、グローバル、かつ入力依存的なオペレーターであるHyenaND を導入します。
アーキテクチャの再構成: 本論文は、「状態空間モデル(SSM)」の景観を2つの明確な分岐に分離しています:
一般化された回帰(Generalized Recurrent): 1次元の非LTIデータ(例:テキスト)に適しています。
一般化された畳み込み(Generalized Convolutional): 近似的にLTIである多次元データに適しています。HyenaNDはこの分岐に属します。
ネイティブND畳み込み: HyenaNDは、N N N 次元グリッド(L 1 × ⋯ × L N L_1 \times \dots \times L_N L 1 × ⋯ × L N )上で直接畳み込みを実行します。これは、N N N 次元高速フーリエ変換(FFT)を介して評価される、グローバルで暗黙的にパラメータ化されたカーネル K K K を利用しており、O ( L log L ) O(L \log L) O ( L log L ) の複雑さを達成しています。
レジスタとFiLMによる入力依存カーネル: 標準的な畳み込みにおける入力依存性の欠如を克服するため、HyenaNDはカーネルを個々のトークンではなく、入力サンプル全体に条件付けします。
レジスタ(Registers): 学習可能なレジスタ・トークン R R R が入力グリッドの前に付加されます。これらは、グローバルな情報を統合するために共同で処理されます。
FiLM条件付け: ソフトマックス重み付き平均を介して、レジスタの活性化から制御変数 z ( x ) z(x) z ( x ) が導出されます。この変数は、座標ベースのSIREN(Sinusoidal Representation Networks)MLP f θ f_\theta f θ をFeature-wise Linear Modulation(FiLM)を用いて変調します。
カーネル合成: 最終的なカーネルは K ( x ) = w ( c ) ⊙ f θ ( c ; z ( x ) ) K(x) = w(c) \odot f_\theta(c; z(x)) K ( x ) = w ( c ) ⊙ f θ ( c ; z ( x )) です。ここで w w w は学習可能なガウス窓です。これにより、単一のグローバルFFTパスを維持しながら、カーネルが特定の入力インスタンスに適応することが可能になります。
カーネルのパラメータ化: カーネルは、SIRENを用いた連続的な座標グリッド上で定義されるため、解像度に依存しません。著者らは、カーネルが異なるチャネルにわたって低周波および高周波の両方のコンテンツを解像できるように、SIRENの周波数スケーリング係数 ω 0 \omega_0 ω 0 に対する「ブロック対角マルチ周波数初期化」を導入しています。
境界の処理: 実装は、周期的な物理システムのための循環畳み込み(circular convolutions)と、自然な画像やボリュームのための非循環畳み込み(non-circular convolutions)の両方をサポートしており、FFTの前に入力を2倍のグリッドサイズにゼロパディングすることで、人工的な境界リークを回避しています。
実装: nSubQ 漸近的な複雑度と実測のウォールクロック性能の間のギャップを埋めるため、著者らはIO認識型のCUDAライブラリであるnSubQ を開発しました。
融合カーネル(Fused Kernels): nSubQは、フォワードFFT、スペクトル変調(入力依存カーネルの適用)、逆FFT、および出力セグメンテーションを単一のCUDAカーネルに融合します。
メモリ効率: 中間表現をオンチップの共有メモリ内に保持することで、高帯域幅メモリ(HBM)へのトラフィックを最小限に抑えます。
パフォーマンス: この融合により、理論的な O ( L log L ) O(L \log L) O ( L log L ) の優位性が、特にアテンションメカニズムがメモリ帯域幅のボトルネックに直面する長シーケンスにおいて、実用的なスピードアップへと変わります。
主な貢献
概念的な再構成: 劣二次的なオペレーターを、回帰型(1D)と一般化された畳み込み型(ND)のファミリーへと明確に分離し、後者が多次元のLTIデータにとって自然なプリミティブであることを論じています。
HyenaND オペレーター: ラスタライズを回避し、LTIの帰納バイアスを保持する、ネイティブなND、サンプルレベルの入力依存的、劣二次的オペレーターです。
nSubQ ライブラリ: ND FFT 畳み込みのウォールクロック・スピードアップを実現する、高性能なCUDA実装です。
実証的検証: 1D(ゲノミクス)、2D(ImageNet)、3D(医療画像)、およびPDEモデリングにわたる包括的な評価。
結果
制御された空間的回想(Controlled Spatial-Recall): HyenaNDは、その並進等変性(translation-equivariant)の事前分布により、3D「単純コピー」タスクにおいて、アテンションおよび双方向Mambaよりも大幅に優れた性能(MSEで3〜4桁低い)を示しました。アテンションが失敗(OOM)するか、攻撃的なパッチ化を必要とする一方で、HyenaNDはネイティブな3D解像度を維持しています。
長コンテキスト・ゲノミクス (1D): OpenGenome2において、ハイブリッドHyenaND-attentionモデル(特に2つのMHAレイヤーを持つH2)は、フルTransformerおよび純粋なHyena構成を上回り、最も低いパープレキシティを達成しました。HyenaNDは、10Mトークンのコンテキストにわたって安定したスループットのスケーリングを示しました。
コンピュータビジョン (ImageNet-1K): 純粋なHyenaND(ViT-5-Smallバックボーン)は、アテンションのベースラインに匹敵する81.5%のTop-1精度を達成しました。ハイブリッド構成(例:( H A ) × 6 (H A) \times 6 ( H A ) × 6 )は、純粋なアテンションおよび回帰ベースのMambaハイブリッドの両方を上回り、大幅に少ないFLOPs(小さなパッチサイズで最大82%の節約)で82.1%の精度を達成しました。
PDEモデリング (The Well): HyenaNDは、5つのすべてのデータセット(流体力学、MHDなど)において、フル解像度のConvNeXt U-Netおよびアテンションのベースラインを上回り、最も低い分散正規化RMSE(VRMSE)を達成しました。アテンションに対する性能向上は、シーケンス長とともに増大しました。
3D医療セグメンテーション (PanTS): ハイブリッドHyenaND-attentionモデルは、純粋なアテンションおよび純粋なHyeraのバリアントを凌駕する、最高のMean Diceスコア(0.7559)を達成しました。HyenaNDのバリアントは、アテンションのベースラインと比較して、ピークGPUメモリ割り当てを約10.8%削減しました。
意義と主張 本論文は、HyenaNDとnSubQが、ファウンデーションモデルがパッチ化に依存することを強いている根本的な制約に対処すると主張しています。パッチ化は微細な構造情報を損なわせます。ネイティブなND、入力依存的、かつ劣二次的なオペレータ を提供することで、著者らはモデルがデータの固有の幾何学とネイティブな解像度でデータを取り込むことを可能にします。
著者らは、HyenaNDをアテンションの普遍的な置き換えとしてではなく、補完的なスペシャリストとして位置付けています。純粋なHyenaNDのスタックは強力なアテンションのベースラインに匹敵しますが、ハイブリッド構成(HyenaNDとアテンションのインターリーブ)は、純粋なアテンションおよび回帰ベースのハイブリッドの両方を一貫して上回ります。これは、畳み込みの分岐(グローバルなLTI構造用)とアテンションの分岐(微細な選択性用)が、タッグを組んで使用されるのが最適であることを示唆しています。本研究は、アーキテクチャ設計の転換を促進するために、HyenaNDアーキテクチャとnSubQライブラリの両方を公開しています。
毎週最高の machine learning 論文をお届け。
スタンフォード、ケンブリッジ、フランス科学アカデミーの研究者に信頼されています。
受信トレイを確認して登録を完了してください。
問題が発生しました。もう一度お試しください。
スパムなし、いつでも解除可能。
週刊ダイジェスト — 最新の研究をわかりやすく。 登録 ×