FLARE: 高速低ランク注意ルーティングエンジン

要旨

自己注意機構の二次計算量は、大規模な非構造化メッシュにおける適用性と拡張性を制限しています。本論文では、Fast Low-rank Attention Routing Engine (FLARE)を提案します。これは、固定長の潜在シーケンスを通じて注意をルーティングする線形計算量の自己注意機構です。各注意ヘッドは、学習可能なクエリトークンを使用して入力シーケンスをM ≪ Nトークンの固定長潜在シーケンスに投影し、Nトークン間のグローバルな通信を行います。ボトルネックシーケンスを通じて注意をルーティングすることで、FLAREはO(NM)のコストで適用可能な低ランク形式の注意を学習します。FLAREは前例のない問題規模にスケールするだけでなく、多様なベンチマークにおいて最先端のニューラルPDEサロゲートと比較しても優れた精度を実現します。また、さらなる研究を促進するために、新しい積層造形データセットを公開します。コードはhttps://github.com/vpuri3/FLARE.pyで利用可能です。

English

The quadratic complexity of self-attention limits its applicability and scalability on large unstructured meshes. We introduce Fast Low-rank Attention Routing Engine (FLARE), a linear complexity self-attention mechanism that routes attention through fixed-length latent sequences. Each attention head performs global communication among N tokens by projecting the input sequence onto a fixed length latent sequence of M ll N tokens using learnable query tokens. By routing attention through a bottleneck sequence, FLARE learns a low-rank form of attention that can be applied at O(NM) cost. FLARE not only scales to unprecedented problem sizes, but also delivers superior accuracy compared to state-of-the-art neural PDE surrogates across diverse benchmarks. We also release a new additive manufacturing dataset to spur further research. Our code is available at https://github.com/vpuri3/FLARE.py.

FLARE: 高速低ランク注意ルーティングエンジン

FLARE: Fast Low-rank Attention Routing Engine

要旨

Support