FLARE: Motor de Roteamento de Atenção Rápida de Baixa Ordem

FLARE: Fast Low-rank Attention Routing Engine

August 18, 2025

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

cs.AI

Resumo

A complexidade quadrática da auto-atenção limita sua aplicabilidade e escalabilidade em malhas não estruturadas de grande porte. Apresentamos o Fast Low-rank Attention Routing Engine (FLARE), um mecanismo de auto-atenção com complexidade linear que direciona a atenção por meio de sequências latentes de comprimento fixo. Cada cabeça de atenção realiza comunicação global entre N tokens projetando a sequência de entrada em uma sequência latente de comprimento fixo de M ll N tokens usando tokens de consulta aprendíveis. Ao direcionar a atenção por meio de uma sequência de gargalo, o FLARE aprende uma forma de atenção de baixo posto que pode ser aplicada com custo O(NM). O FLARE não apenas escala para tamanhos de problemas sem precedentes, mas também oferece precisão superior em comparação com os melhores substitutos neurais de EDP em diversos benchmarks. Também disponibilizamos um novo conjunto de dados de manufatura aditiva para estimular pesquisas adicionais. Nosso código está disponível em 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.