FLARE: Fast Low-rank Attention Routing Engine

FLARE is a novel token mixing layer that bypasses the quadratic cost of self-attention by leveraging low-rankness. FLARE is built on the argument that projecting input sequences onto shorter latent sequences, and then unprojecting to the original sequence length, is equivalent to constructing a low-rank form of attention with rank at most equal to the number of latent tokens (see figure below).

Furthermore, we argue that multiple simultaneous low-rank projections could collectively capture
a full attention pattern. Our design allocates a distinct slice of the latent tokens to
each head resulting in distinct projection matrices for each head. This allows each head to learn
independent attention relationships, opening up a key direction of scaling and exploration, wherein
each head may specialize in distinct routing patterns.

FLARE is built entirely from standard fused attention primitives, ensuring high GPU utilization and ease of integration into existing transformer architectures. By replacing full self-attention with low-rank projections and reconstructions, FLARE achieves linear complexity in the number of points (see plot below). As such, FLARE enables end-to-end training on unstructured meshes with one million points on a single GPU – the largest scale demonstrated for transformer-based PDE surrogates.