FLARE : Moteur de Routage d'Attention Rapide à Faible Rang

FLARE: Fast Low-rank Attention Routing Engine

August 18, 2025

papers.authors: Vedant Puri, Aditya Joglekar, Kevin Ferguson, Yu-hsuan Chen, Yongjie Jessica Zhang, Levent Burak Kara

cs.AI

papers.abstract

La complexité quadratique de l'auto-attention limite son applicabilité et son extensibilité sur de grands maillages non structurés. Nous introduisons Fast Low-rank Attention Routing Engine (FLARE), un mécanisme d'auto-attention à complexité linéaire qui dirige l'attention à travers des séquences latentes de longueur fixe. Chaque tête d'attention effectue une communication globale parmi N tokens en projetant la séquence d'entrée sur une séquence latente de longueur fixe de M ≪ N tokens à l'aide de tokens de requête apprenables. En dirigeant l'attention à travers une séquence goulot, FLARE apprend une forme d'attention de faible rang qui peut être appliquée à un coût de O(NM). FLARE non seulement s'adapte à des tailles de problèmes sans précédent, mais offre également une précision supérieure par rapport aux substituts de PDE neuronaux de pointe sur divers benchmarks. Nous publions également un nouveau jeu de données sur la fabrication additive pour stimuler des recherches ultérieures. Notre code est disponible à l'adresse 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.