FLARE++: Low-rank attention with dynamic attention routing

2026-08-11

AI Summary

This work addresses the challenge of scaling self-attention for surrogate modeling of partial differential equations (PDEs) on irregular domains, where full self-attention incurs prohibitive quadratic complexity and existing low-rank methods rely on fixed query templates lacking input adaptivity. To overcome this limitation, we propose FLARE++, which introduces, for the first time, a dynamic, input-dependent query mechanism into low-rank attention routing. By reusing encoder features to generate routing queries conditioned on the input, FLARE++ enables adaptive token compression and redistribution while preserving explicit low-rank decomposition and compatibility with standard scaled dot-product attention (SDPA). The method achieves linear O(NM) complexity and supports multi-GPU context parallelism. Experiments demonstrate that FLARE++ outperforms the fixed-query FLARE by 24% on average across standard PDE surrogate benchmarks and improves average accuracy by 2.3 percentage points on Long Range Arena tasks.

Abstract

Full self-attention is a strong token mixer for PDE surrogates on irregular domains, but its quadratic cost limits its use on high-resolution problems. Efficient latent-attention models such as the Fast Low-rank Attention Routing Engine (FLARE) avoid that cost by routing all N tokens through M << N learned latent queries, but those queries are parameters: once trained, the same learned query templates serve every input. We remove this restriction with FLARE++, a low-rank attention architecture with dynamic token routing. FLARE++ reuses FLARE's own encoder to build its routing queries: learned latent seeds drive one extra encode call that gathers the N input tokens into M input-conditioned queries, and those queries then determine how the same tokens are compressed and redistributed. This preserves FLARE's explicit low-rank factorization and linear 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++ is competitive across a set of standard PDE surrogate benchmarks, improving on fixed-query FLARE by 24% on average, and it gains 2.3 points of average accuracy on Long Range Arena.

Problem

• low-rank attention
• dynamic routing
• PDE surrogates
• irregular domains
• computational complexity

Innovation

• low-rank attention
• dynamic routing
• latent queries
• scaled dot-product attention
• context-parallel