🤖 machine learning

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

FLARE++ is a low-rank attention architecture that enhances the efficiency of PDE surrogates on irregular domains by introducing dynamic, input-conditioned token routing via an extra encoding step, thereby achieving competitive accuracy improvements over fixed-query baselines while maintaining linear complexity and supporting scalable multi-GPU implementation.

Original authors: Vedant Puri, Yongjie Jessica Zhang, Levent Burak Kara

Published 2026-08-13

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Vedant Puri, Yongjie Jessica Zhang, Levent Burak Kara

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a computer to predict how things move, flow, or stretch—like wind swirling around a race car, water rushing through a pipe, or stress building up in a bridge. In the world of science, these are called "partial differential equations" (PDEs). To solve them, computers break the object down into millions of tiny dots, or "tokens," and ask each dot to talk to every other dot to figure out the final picture. This is like a massive classroom where every student must whisper a secret to every other student to solve a puzzle. While this "full attention" method is incredibly accurate, it's also exhausting. If you have a million dots, the number of whispers needed grows so fast that it becomes impossible to run on normal computers. Scientists have been looking for a shortcut: a way to let the dots talk efficiently without needing a million conversations. They found a trick called "low-rank attention," where dots don't talk to everyone directly. Instead, they send their messages to a small group of "summarizers" (latent tokens), who then pass the condensed news back to the dots. It's like having a class representative who listens to the whole class and then tells everyone the main points.

This paper introduces a new version of that shortcut called FLARE++. The original shortcut (FLARE) had a small flaw: the "summarizers" were fixed. They were like pre-written templates that didn't change, no matter what the problem was. Whether you were modeling a bridge or a storm, the same summarizers tried to do the job. The authors realized this was limiting. They built FLARE++, a system where the summarizers are created on the fly, specifically for the problem at hand. Instead of using a static template, FLARE++ looks at the current situation, creates a custom set of summarizers, and then uses those to organize the information. The result is a system that is just as fast as the old one but significantly smarter. In tests on standard engineering problems, this new dynamic approach reduced errors by an average of 24% compared to the fixed version, and it even improved performance on general language tasks, proving that this "custom summarizer" trick works even outside of physics simulations.

The Problem: The "Whispering" Classroom

Imagine you are in a giant stadium filled with 100,000 people (the "tokens"). You need to know the temperature at every single seat. In the old-school method (Full Self-Attention), every person has to ask every other person, "What's your temperature?" and then combine all those answers. This is incredibly accurate, but it's a logistical nightmare. If you double the number of people, the number of conversations quadruples. It's like trying to organize a party where everyone must shake hands with everyone else; eventually, you run out of time and space.

To fix this, scientists invented a "middleman" system. Instead of everyone talking to everyone, the crowd picks a small group of 100 "representatives" (latent tokens). Everyone tells their temperature to the nearest representative. The representatives mix the information and then tell the crowd the result. This is much faster. However, the original version of this system (FLARE) had a rigid rule: the representatives were always the same 100 people, chosen before the game started. They were like a fixed team of scouts who had to interpret every single situation, from a gentle breeze to a hurricane, using the exact same strategy. Sometimes, a fixed team just can't adapt well enough to a weird or complex situation.

The Solution: The "Chameleon" Summarizers

The authors of this paper asked a simple question: What if the representatives could change depending on who is in the room?

Enter FLARE++. Instead of using a pre-set team of 100 representatives, FLARE++ looks at the crowd first. It runs a quick, smart scan of the current situation and synthesizes a brand-new set of 100 representatives specifically designed for that moment. It's like a chameleon changing its colors to match its environment, or a chef tasting the soup before deciding which spices to add.

Here is how it works in practice:

The Scan: The system takes the input (the crowd of dots) and runs a quick calculation to create a custom set of "routing queries." These are the instructions for the new, temporary representatives.
The Gathering: The crowd sends their data to these custom representatives.
The Redistributing: The representatives mix the data and send it back to the crowd.

The magic is that this whole process still happens very fast. The paper shows that even though FLARE++ does a little extra work to create the custom representatives, it doesn't slow things down enough to matter. In fact, because the representatives are better suited to the specific problem, the system makes fewer mistakes.

What They Found

The team tested FLARE++ on five different engineering challenges, ranging from how a car wing handles air (Airfoil) to how water flows through a porous rock (Darcy). They compared it against the old fixed-FLARE system and other popular methods.

Better Accuracy: On average, FLARE++ reduced the error rate by 24% compared to the fixed-FLARE system. On some specific tests, like the "Elasticity" problem (simulating how materials stretch), it was nearly 45% more accurate.
Depth vs. Breadth: The researchers found that FLARE++ is so good at picking the right representatives that it can achieve the same accuracy as the old system using half the number of layers (or "depth"). It's like getting the same grade on a test by studying smarter, rather than studying longer.
General Skills: They also tested it on a general language puzzle called "Long Range Arena." Even though this wasn't a physics problem, FLARE++ still improved the score by 2.3 points on average, showing that the "custom summarizer" trick is a powerful tool for many types of data, not just physics.

The Cost and the Catch

Is there a downside? The paper is honest about the trade-offs. Creating those custom representatives does take a little bit more time—about 1.3 to 1.5 times longer per step than the fixed version. However, because the system is so much more accurate, you don't need to run it as long or as deep to get a good result. The authors measured this on powerful graphics cards (NVIDIA H100) and found that the extra time is worth the jump in quality.

They also built a special version that can run across multiple computers at once. They split the crowd of dots across different machines, and the "custom representatives" were created without ever needing to gather all the dots onto a single machine. This means the system can handle massive problems (millions of dots) without running out of memory, keeping the efficiency high even as the problem gets bigger.

The Bottom Line

FLARE++ doesn't reinvent the wheel; it just makes the spokes adjustable. By letting the system decide how to summarize information based on what it's looking at, rather than forcing a one-size-fits-all approach, it solves complex physics problems faster and more accurately. The paper suggests that this dynamic routing is a significant step forward, proving that in the world of AI and physics, being flexible is often better than being fixed.

Technical Summary: FLARE++: Low-Rank Attention with Dynamic Attention Routing

Problem Statement
Full self-attention mechanisms have become dominant for Partial Differential Equation (PDE) surrogates on irregular domains because they allow every discretization point (token) to communicate with every other. However, this generality incurs a quadratic computational cost ( $O(N^2)$ ) in the number of tokens $N$ , rendering it infeasible for the high-resolution meshes typical of engineering problems. While efficient latent-attention models like the Fast Low-rank Attention Routing Engine (FLARE) reduce this cost to linear $O(NM)$ by routing $N$ tokens through $M \ll N$ learned latent queries, they suffer from a rigidity: the routing queries are fixed parameters. Once trained, these query templates remain static for every input, regardless of the specific geometry or boundary conditions. Consequently, the "compression template" (how information is gathered and redistributed) is decoupled from the specific field being processed.

Methodology: FLARE++
FLARE++ addresses this limitation by introducing dynamic token routing. Instead of using fixed learned parameters as routing queries, FLARE++ synthesizes these queries dynamically from the input tokens themselves.

The architecture preserves FLARE's explicit low-rank factorization and its encode–decode structure, which consists of two Scaled Dot-Product Attention (SDPA) calls:

Dynamic Query Synthesis: FLARE++ reuses its own encoder mechanism. Learned latent seeds ( $\tilde{Q}_h$ ) act as queries for an additional SDPA call that gathers the $N$ input tokens into $M$ input-conditioned vectors. These $M$ vectors then serve as the routing queries ( $Q_h(X)$ ) for the main attention block.
Encode–Decode Routing: The synthesized queries $Q_h(X)$ are used in the standard FLARE encode–decode pair. The encoder gathers the $N$ input tokens into $M$ latent values, and the decoder scatters these values back to the $N$ output positions.

This approach ensures that the routing weights are determined by the current input field, allowing the compression template to adapt to the specific geometry and boundary conditions of each sample. Crucially, the method does not introduce a separate latent-processing stage (unlike Transolver) nor does it add depth or width to the residual stream. The routing matrix remains explicitly rank- $M$ for any fixed input, and the complexity remains linear $O(NM)$.

Key Contributions

Synthesized Compression Template: FLARE++ constructs the routing queries from the input tokens using standard SDPA, ensuring the input determines how it is compressed and redistributed. This preserves FLARE's explicit rank- $M$ operator and fused-SDPA implementation without adding residual depth.
Multi-GPU Context Parallelism: The authors provide an exact token-sharded implementation. Input tokens are distributed across devices, and the encoder gathers them into replicated latent tokens. Communication is limited to latent outputs and softmax statistics, avoiding the need to gather the full $N$ -token sequence on a single device. This maintains parallel efficiency near unity in both time and memory.
Comprehensive Evaluation: The paper evaluates FLARE++ against fixed-query FLARE, the Transolver family, and full self-attention across standard PDE benchmarks and the Long Range Arena (LRA).

Results

PDE Surrogate Benchmarks: On five standard benchmarks (Elasticity, Darcy, Airfoil, Pipe, and DrivAerML-40K), FLARE++ achieves the lowest relative $L_2$ error among all tested models. It reduces the relative $L_2$ error of fixed-query FLARE by an average of 24% and reduces the error of Transolver-3 by an average of 31%.
Efficiency vs. Accuracy Trade-off: While a single FLARE++ block costs 1.3–1.5 $\times$ the step time of a FLARE block due to the extra projection and SDPA call, FLARE++ reaches a given accuracy with significantly fewer residual blocks. For instance, on the Elasticity benchmark, FLARE++ matches the accuracy of FLARE at half the residual depth.
Latent Budget Behavior: On the Elasticity benchmark (fundamentally low-rank), FLARE++ continues to improve as the latent budget $M$ increases, whereas fixed-query FLARE saturates or even degrades. On the Darcy benchmark (rank-limited), both models benefit from larger $M$ , but FLARE++ extracts more value from the budget.
Long Range Arena: Outside of PDEs, FLARE++ improves the average accuracy of the FLARE architecture by 2.3 points on the LRA benchmark, demonstrating that the dynamic routing mechanism is not tuned exclusively to PDE discretizations.
Scaling: The token-sharded implementation achieves parallel efficiency of 0.92–0.98 in time and 0.95–0.98 in memory across up to four NVIDIA H100 GPUs, allowing training on meshes with up to $10^6$ points.

Significance
The paper claims that FLARE++ demonstrates that the "compression template" in low-rank attention is a critical component that should be input-dependent. By synthesizing routing queries from the input field, the model adapts its low-rank bottleneck to the specific sample, improving accuracy without increasing the network's depth or width. The work isolates the benefit of dynamic routing from other architectural changes, showing that the construction of the template itself—rather than just the latent processing—is what drives performance gains. Furthermore, the successful multi-GPU implementation proves that these dynamic, low-rank operators can scale to industrial-resolution meshes without the memory bottlenecks of full self-attention.

The Problem: The "Whispering" Classroom

The Solution: The "Chameleon" Summarizers

What They Found

The Cost and the Catch

The Bottom Line

More like this