FLARE++：带动态注意力路由的低秩注意力机制

中文摘要 AI 辅助

全自注意力是用于不规则域PDE代理模型的强大token混合器，但其二次复杂度限制了在高分辨率问题上的应用。Fast Low-rank Attention Routing Engine（FLARE）等高效隐注意力模型通过将所有N个token路由至M<<N个学习到的隐查询来规避该开销，但这些查询是参数：训练完成后，相同的学习到的查询模板适用于所有输入。FLARE++消除了这一限制，它是一种带动态token路由的低秩注意力架构。FLARE++复用FLARE自身的编码器构建其路由查询：学习到的隐种子驱动一次额外的编码调用，将N个输入token聚合为M个输入条件查询，这些查询随后决定如何压缩和重新分配相同的token。这保留了FLARE显式低秩分解与线性O(NM)复杂度，且仅通过标准缩放点积注意力（SDPA）调用即可表达完整的路由操作。我们还提供了多GPU上下文并行实现，可在设备间分片输入token，无需在单个设备上聚合完整token序列。FLARE++在一组标准PDE代理基准测试中表现具有竞争力，较固定查询的FLARE平均提升24%，且在Long Range Arena上平均准确率提升2.3个点。

英文摘要

Full self-attention is a strong token mixer for PDE surrogates on irregular domains, but its quadratic cost limits its use on large problems. Latent-space attention methods such as PerceiverIO, Transolver, and FLARE (Fast Low-rank Attention Routing Engine) avoid that cost by routing attention among $N$ tokens through $M\ll N$ learned latents. They compress by dot-product matching of the input tokens against $M$ learned query tokens or projection weights: once trained, the same learned templates serve every input. We remove this restriction with FLARE++, a low-rank attention architecture with input-conditioned routing queries. FLARE++ uses FLARE's own encoder to map the $N$ input tokens to $M$ query tokens, which correct the learned queries. The adapted queries then determine how that same input is compressed and redistributed. This preserves FLARE's explicit low-rank factorization and linear $\mathcal O(NM)$ complexity, and expresses the complete routing operation with standard scaled dot-product attention (SDPA) calls alone. We also provide a multi-GPU context-parallel implementation that shards input tokens across devices without ever gathering the full token sequence on one of them. FLARE++ reduces FLARE's error by $25\%$ on average across five standard PDE benchmarks, achieving the lowest errors among the efficient models compared. The gains persist on industrial-scale DrivAerML aerodynamics and on Long Range Arena, where average accuracy rises by $3.5\%$ over fixed-query FLARE. Code is available at https://github.com/vpuri3/FLARE.py.

发表机构

• Carnegie Mellon University（卡内基梅隆大学）

机构由 AI 辅助整理，请以论文原文为准。