FLARE: Motor de Enrutamiento de Atención de Bajo Rango Rápido

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

Resumen

La complejidad cuadrática de la auto-atención limita su aplicabilidad y escalabilidad en mallas grandes no estructuradas. Presentamos Fast Low-rank Attention Routing Engine (FLARE), un mecanismo de auto-atención de complejidad lineal que dirige la atención a través de secuencias latentes de longitud fija. Cada cabeza de atención realiza comunicación global entre N tokens proyectando la secuencia de entrada en una secuencia latente de longitud fija de M ll N tokens utilizando tokens de consulta aprendibles. Al dirigir la atención a través de una secuencia de cuello de botella, FLARE aprende una forma de atención de bajo rango que puede aplicarse con un costo de O(NM). FLARE no solo escala a tamaños de problemas sin precedentes, sino que también ofrece una precisión superior en comparación con los sustitutos de EDP neuronales más avanzados en diversos puntos de referencia. También publicamos un nuevo conjunto de datos de fabricación aditiva para fomentar más investigaciones. Nuestro código está disponible en 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.