FLARE: Быстрый механизм маршрутизации внимания с низким рангом

FLARE: Fast Low-rank Attention Routing Engine

August 18, 2025

Авторы: Vedant Puri, Aditya Joglekar, Kevin Ferguson, Yu-hsuan Chen, Yongjie Jessica Zhang, Levent Burak Kara

cs.AI

Аннотация

Квадратичная сложность механизма самовнимания ограничивает его применимость и масштабируемость на больших неструктурированных сетках. Мы представляем Fast Low-rank Attention Routing Engine (FLARE) — механизм самовнимания с линейной сложностью, который направляет внимание через фиксированные по длине латентные последовательности. Каждый блок внимания выполняет глобальную коммуникацию между N токенами, проецируя входную последовательность на фиксированную латентную последовательность длины M, где M ≪ N, с использованием обучаемых токенов запроса. Направляя внимание через последовательность-«бутылочное горлышко», FLARE обучает низкоранговую форму внимания, которая может быть применена с вычислительной сложностью O(NM). FLARE не только масштабируется до беспрецедентных размеров задач, но и демонстрирует превосходную точность по сравнению с современными нейронными суррогатами дифференциальных уравнений в частных производных на различных тестовых наборах. Мы также публикуем новый набор данных по аддитивному производству для стимулирования дальнейших исследований. Наш код доступен по адресу 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.