FLARE: Fast Low-rank Attention Routing Engine

## Abstract

Abstract 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 N tokens by projecting the input sequence onto a fixed length latent sequence of M ≪ N M\ll N tokens using learnable query tokens.
Routing information through this bottleneck constrains the interaction to have rank at most M M : the resulting operation is equivalent to an attention matrix of rank at most M M , realized at 𝒪 ​ ( N ​ M ) \mathcal{O}(NM) cost.
This formulation allows users to flexibly adjust the effective rank of attention –and thus the balance between accuracy and efficiency– by tuning M M to their resource and performance requirements.
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 introduce a new additive manufacturing dataset to spur further research.
Our code is available at https://github.com/vpuri3/FLARE.py .

## 1 Introduction

High-fidelity simulations of physical systems are often too costly for multi-query applications, such as design optimization or uncertainty quantification.
Machine learning offers a promising alternative via surrogate models that learn system dynamics from data, enabling fast approximations that accelerate experimentation and decision making.

Among machine learning (ML) architectures, transformers (Vaswani et al., 2017) have shown exceptional scalability and generalization capabilities in domains such as natural language processing (Devlin et al., 2019) and computer vision (Dosovitskiy et al., 2020).
This has led to increased interest in adapting transformers for spatially distributed data such as point clouds and meshes in physical simulations.
However, applying transformers directly to large-scale unstructured meshes introduces severe computational bottlenecks.
The core challenge lies in the self-attention mechanism in which, for a sequence of $N$ tokens, each token attends to every other token in the sequence, resulting in $\mathcal{O}(N^{2})$ time and memory complexity.
Although this global communication is key to the expressive power of transformers, enabling them to outperform graph neural networks (GNNs) that are inherently local (Yun et al., 2019), its computational footprint hinders its scalability.

In the context of learning partial differential equation (PDE) surrogates, i.e., regressing continuous fields over point clouds/meshes, each point in the 3D point cloud is treated as a token, with associated features that encode geometric and physical quantities such as coordinates, normals, and material properties.

Several works have proposed approximations to the full self-attention mechanism to mitigate the cost.
Perceiver and PerceiverIO (Jaegle et al., 2021b; a) introduced cross-attention between the input sequence and a length $M$ learnable latent sequence.
In this mechanism, the $M\times N$ attention weights are used to project the input sequence to the latent sequence, followed by self-attention blocks in the latent space.
This decouples the length of the input sequence from the computational complexity by avoiding dense self-attention on the inputs.
However, the latent bottleneck in PerceiverIO can limit accuracy as the model may discard fine-grained features if the number of latent tokens is too low.

In learning PDE surrogates, Transolver (Wu et al., 2024) and Latent Neural Operator (LNO) (Wang & Wang, 2024a) adopt a similar projection-unprojection scheme to map variable-length point clouds to fixed-length latent representations.
Both models obtain projection weights by expanding the feature dimension, and apply self-attention on the resulting latent sequence.
The differences between the three approaches are highlighted in [Table 1](https://arxiv.org/html/2508.12594v2#S1.T1).
Notably, Perceiver and LNO only perform a single encoding and decoding step, whereas Transolver reapplies projection and unprojection in every transformer block.
This enables Transolver to form deep models, where attention connections in downstream layers benefit from the rich embeddings and global context aggregated by previous layers.
However, the attention projector layers proposed in Transolver and LNO cannot utilize existing GPU kernels for scaled dot-product attention (SDPA), a well-optimized algorithm in modern deep learning (Dao et al., 2022).
This limits their scalability to large meshes.

**Table 1: Comparison of FLARE with different latent attention schemes in PerceiverIO (Jaegle et al., 2021a), LNO (Wang & Wang, 2024a), and Transolver (Wu et al., 2024).**
| Criteria | PerceiverIO | Transolver | LNO | FLARE (ours) |
| --- | --- | --- | --- | --- |
| Sequential encoding and decoding blocks | ✗ | ✓ | ✗ | ✓ |
| Parallel encodings & decodings | ✓ | ✓* | ✗ | ✓ |
| Weight sharing b/w encoding and decoding | ✗ | ✓ | ✓ | ✓ |
| Latent space self-attention | ✓ | ✓ | ✓ | ✗ |
| Compatible with fused attention kernels | ✓ | ✗ | ✗ | ✓ |

We propose Fast Low-rank Attention Routing Engine (FLARE),
a simple yet effective mechanism designed to break the scalability barrier in PDE surrogate learning.
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 1](https://arxiv.org/html/2508.12594v2#S1.F1)).

Furthermore, we argue that multiple simultaneous low-rank projections could collectively capture a full attention pattern.
Unlike Transolver which shares projection weights across heads, or LNO which applies only a single projection, 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.

To support this intuition, we draw on known properties of rank-deficient matrices, to develop a linear-time eigenanalysis procedure to study the learned communication matrices.
Our spectral analysis reveals that attention heads operate in low-rank subspaces, validating our core design choice.
Furthermore, we observe diversity in the eigen spectrum across heads, reinforcing the motivation for independent head-wise projections.
These findings highlight the efficiency and expressiveness of our low-rank formulation.
We summarize our main contributions below.

Figure: Figure 1: Schematic of a FLARE block. In FLARE, each head projects the input sequence with $N$ tokens to a fixed-length sequence of $M$ tokens via the cross-attention matrix $W_{\text{encode}}=\mathrm{softmax}(Q\cdot K^{T})$, and then projects back to the original length via the cross-attention matrix $W_{\text{decode}}=\mathrm{softmax}(K\cdot Q^{T})$. The overall operation is equivalent to token mixing on the input sequence with the rank-deficient matrix $\left(W_{\text{decode}}\cdot W_{\text{encode}}\right)$.
Refer to caption: x1.png

- •
Linear complexity token mixing.
FLARE is self-attention mechanism designed to learn on long sequences such as point clouds.
By replacing full self-attention with low-rank projections and reconstructions, FLARE achieves linear complexity in the number of points.
- •
Superior accuracy.
Across multiple PDE benchmarks, FLARE achieves superior predictive accuracy compared to leading neural surrogate models, despite operating with fewer parameters, and at much lower computational complexity.
- •
Scalability.
FLARE is built entirely from standard fused attention primitives, ensuring high GPU utilization and ease of integration into existing transformer architectures.
As such, FLARE enables end-to-end training on unstructured meshes with one million points (see [Figure 2](https://arxiv.org/html/2508.12594v2#S2.F2)) without distributed computing (Luo et al., 2025) or memory offloading – the largest scale demonstrated for transformer-based PDE surrogates.
- •
Benchmark dataset of additive manufacturing simulations.
We release a large-scale, high-resolution dataset for residual displacement prediction in metal additive manufacturing to encourage further research in scalable PDE surrogate modeling.
The dataset is available at
[https://huggingface.co/datasets/vedantpuri/LPBF_FLARE](https://huggingface.co/datasets/vedantpuri/LPBF_FLARE).

## 2 Related work

#### Neural PDE surrogates.

Learning neural surrogates for PDEs is an exciting and rapidly growing direction in scientific machine learning.
Neural operators (Li et al., 2020; Lu et al., 2021; Kovachki et al., 2023) have been proposed for learning mappings between infinite-dimensional input and output function spaces, enabling mesh-independent generalization.
Extensions to neural operators that incorporate geometric priors and graph-based representations further improve performance, especially for problems defined on complex unstructured meshes (Li et al., 2023b; c).

Graph networks have also been widely explored as PDE surrogates (Pfaff et al., 2020; Elrefaie et al., 2024; Ferguson et al., 2025), leveraging their inherent ability to model local neighborhood interactions on meshes.
Transformer-based architectures have emerged more recently as powerful PDE surrogates (Alkin et al., 2024; Cao, 2021; Li et al., 2022; 2023a; Hao et al., 2023), allowing global context aggregation and dynamic prediction of complex phenomena (Qian et al., 2025).
Recent works (Alkin et al., 2024; Wu et al., 2024; Wang & Wang, 2024a) leverage latent space attentions for PDE modeling, achieving high accuracy with reduced computational cost.
Building upon these advances, our FLARE method employs a linear complexity attention mechanism with a learnable latent bottleneck to further improve accuracy and efficiency on PDE surrogate modeling tasks.

#### Efficient attention mechanisms.

Several methods have been proposed to address the quadratic complexity of the attention mechanism.
In (Wang et al., 2020), attention matrices are shown to be low-rank.
Their proposed model, called *Linformer*, employs learned linear mappings to project the input key and value sequences to a shorter sequence length, resulting in a linear complexity form of attention. A drawback of this approach is that it cannot directly handle variable input sequence lengths.
Other low-rank approaches include Reformer (Kitaev et al., 2020), which uses locality-sensitive hashing for efficient approximate attention, Nyströmformer (Xiong et al., 2021), which approximates self-attention using the Nyström method, and Scatterbrain (Chen et al., 2021), which combines sparse and low-rank methods for further gains.
Another efficient attention method, called
multi-head latent attention (Liu et al., 2024), jointly compresses key and value tokens to reduce the memory footprint of language models during inference.
However, this does not help with the quadratic bottleneck in self-attention.
Low-Rank Adaptation (LoRA) (Hu et al., 2022) also leverages low-rank projections, though primarily for efficient fine-tuning rather than architectural design or modeling capacity, as we explore in this work.

Figure: Figure 2: Time and memory requirements of different attention schemes. On an input sequence of one million tokens, FLARE (red) is over $200\times$ faster than vanilla attention, while consuming marginally more memory. All models are implemented with flash attention (Dao et al., 2022), and the memory upper bound on a single H100 80GB GPU is depicted with a dashed line. Note that the curves for FLARE are somewhat overlapping.
Refer to caption: x2.png

## 3 Method

### 3.1 Preliminary: Multi-Head Self-Attention

Let $X\in\mathbb{R}^{N\times C}$ denote the input sequence of $N$ tokens with $C$ features each.
The query, key, and value matrices $Q,K,V\in\mathbb{R}^{N\times C}$ are obtained by applying learned linear projections to $X$, we have

$$ $Q=X\cdot W^{q},\hskip 10.00002ptK=X\cdot W^{k},\hskip 10.00002ptV=X\cdot W^{v},$ (1) $$

where $W^{q},W^{k},W^{v}\in\mathbb{R}^{C\times C}$.
The $Q,\,K,\,V$ matrices are then split along the feature dimension and passed to $H$ heads, each with dimension $D=C/H$, enabling parallel computation of attention:

$$ $\displaystyle\left[Q_{1},\ldots,Q_{H}\right]=Q,\hskip 10.00002pt\left[K_{1},\ldots,K_{H}\right]=K,\hskip 10.00002pt\left[V_{1},\ldots,V_{H}\right]=V.$ (2) $$

The scaled dot-product attention (SDPA) operation, introduced by (Vaswani et al., 2017), computes the output as

$$ $Y_{h}=\text{SDPA}(Q_{h},K_{h},V_{h},s)=\mathrm{softmax}\left(\frac{Q_{h}\cdot K_{h}^{T}}{s}\right)\cdot V_{h},$ (3) $$

where $Q_{h},\,K_{h},\,V_{h}\in\mathbb{R}^{N\times D}$ are query, key, and value matrices belonging to head $h$, and
$s$ is typically $\sqrt{D}$.
Note that $\mathrm{softmax}$ is taken along the row-dimension.
The outputs $Y_{h}$ from all heads are then concatenated along the feature dimension to form the final output

$$ $Y=\left[Y_{1},\ldots,Y_{H}\right].$ (4) $$

This concatenation, followed by a linear layer, enables the model to integrate information across attention heads efficiently.

The greatest cost in multi-head self-attention is the call to SDPA which is $\mathcal{O}(N^{2})$ in time and memory complexity.
This is because $Q_{h}\cdot K_{h}^{T}\in\mathbb{R}^{N\times N}$ requires $\mathcal{O}(N^{2})$ storage and softmax, matrix-vector product with $V_{h}$ takes $\mathcal{O}(N^{2})$ operations.
Fortunately, GPU optimized multi-head implementations of SDPA are available in PyTorch (Paszke, 2019).

Figure: Figure 3: PyTorch code for multi-head token mixing operation in FLARE. See [Figure 6](https://arxiv.org/html/2508.12594v2#A2.F6) for an implementation without the fused attention kernel.

### 3.2 FLARE: Fast Low-rank Attention Routing Engine

FLARE is a linear-complexity token mixing layer that learns low-rank global communication structures via attention projections.
The FLARE mechanism introduces a set of $M\ll N$ learnable latent tokens that serve as a bottleneck for information exchange.
The process consists of two stages:

- 1.
Encoding.
The input sequence is projected onto the latent tokens via cross-attention, compressing global information.
- 2.
Decoding.
The latent tokens are then projected back onto the input sequence, distributing the aggregated information.

Formally, we define a learnable query matrix $Q\in\mathbb{R}^{M\times C}$, where each row corresponds to a latent token.
The key and value matrices, $K,V\in\mathbb{R}^{N\times C}$, are obtained by applying deep residual multi-layer perceptrons (MLPs) detailed in [Appendix A](https://arxiv.org/html/2508.12594v2#A1) to the input $X$.
Compared to just a linear layer, these allow the model to learn higher-order feature interactions and deeper nonlinear transformations. Refer to [Appendix F](https://arxiv.org/html/2508.12594v2#A6) for ablation studies.

The matrices $Q$, $K$, and $V$ are first split along the feature dimension into $H$ heads, each of dimension $D=C/H$.
Then, for encoding, each head performs SDPA with a scaling factor $s=1$:

$$ $Z_{h}=\mathrm{SDPA}(Q_{h},K_{h},V_{h},s=1).$ (5) $$

Here, $Q_{h}\in\mathbb{R}^{M\times D}$, $K_{h},V_{h}\in\mathbb{R}^{N\times D}$ are query, key, and value matrices belonging to head $h$ and $Z_{h}\in\mathbb{R}^{M\times D}$ is the latent sequence for head $h$.
For decoding and propagating information back to the input tokens, we perform a second SDPA operation, swapping the roles of queries and keys and using the latent sequence as values:

$$ $Y_{h}=\mathrm{SDPA}(K_{h},Q_{h},Z_{h},s=1)$ (6) $$

where $Y_{h}\in\mathbb{R}^{N\times D}$ is the output for each head.
Similar to multi-head self attention,
the outputs from all heads are concatenated along the feature dimension and passed through a final linear projection to mix information across heads.
As the query matrix has only $M$ tokens, the cost of SDPA calls in [Eq. 5](https://arxiv.org/html/2508.12594v2#S3.E5) and [Eq. 6](https://arxiv.org/html/2508.12594v2#S3.E6) is $\mathcal{O}(NM)$.
PyTorch code for the multi-head implementation is presented in [Figure 3](https://arxiv.org/html/2508.12594v2#S3.F3).
Note that we use a scaling factor $s=1$ instead of $\sqrt{D}$ in typical transformers (Vaswani et al., 2017)) in SDPA.
This modification is explained in [Appendix F](https://arxiv.org/html/2508.12594v2#A6).

#### Low-rank communication.

The two-step attention process can be written as

$$ $Y_{h}=\left(W_{\text{decode},h}\cdot W_{\text{encode},h}\right)\cdot V_{h}$ (7) $$

where

$$ $\displaystyle W_{\text{encode},h}$ $\displaystyle=\mathrm{softmax}(Q_{h}\cdot K_{h}^{T})\in\mathbb{R}^{M\times N},\,\text{and}$ (8) $\displaystyle W_{\text{decode},h}$ $\displaystyle=\mathrm{softmax}(K_{h}\cdot Q_{h}^{T})\in\mathbb{R}^{N\times M}.$ $$

Note that $\mathrm{softmax}$ is taken along the row-dimension.
We define

$$ $\displaystyle W_{h}=W_{\text{decode},h}\cdot W_{\text{encode},h}\in\mathbb{R}^{N\times N}$ (9) $$

as the dense global communication matrix with rank at most $M$.
This low-rank structure, illustrated in [Figure 1](https://arxiv.org/html/2508.12594v2#S1.F1), enables efficient all-to-all communication without explicitly forming $W_{h}$; instead, $W_{\text{encode},h}$ and $W_{\text{decode},h}$ are applied sequentially, resulting in an overall cost of $\mathcal{O}(MN)$ per head.

#### Discussion on the design principles of FLARE.

We defer to [Appendix E](https://arxiv.org/html/2508.12594v2#A5) a deeper analysis of how FLARE’s latent tokens enable gather–scatter communication (acting as selective pooling hubs and broadcasters), why the symmetry between encoding and decoding operators promotes stable information flow, and how fixing latent queries highlights a tradeoff between query dynamics and the need for expressive key/value projections.
Together, these perspectives clarify the structural principles that underlie FLARE’s efficiency and expressivity.

#### FLARE block.

[Figure 1](https://arxiv.org/html/2508.12594v2#S1.F1) (left) illustrates a single FLARE block.
Given input tokens $X\in\mathbb{R}^{N\times C}$, the output of an FLARE block is computed as

$$ $\displaystyle X$ $\displaystyle=X+\text{FLARE}\left(\text{LayerNorm}\left(X\right)\right)$ (10) $\displaystyle X$ $\displaystyle=X+\text{ResMLP}\left(\text{LayerNorm}\left(X\right)\right).$ $$

Here, ResMLP (He et al., 2016) denotes a residual MLP block detailed in [Appendix A](https://arxiv.org/html/2508.12594v2#A1), and LayerNorm denotes layer normalization (Ba et al., 2016) operation.
To summarize, a FLARE block consists of a token mixing operation via FLARE, a pointwise residual MLP, and layer normalization in pre-norm format (Xiong et al., 2020).
Deep residual MLPs within the block enable complex, token-level feature transformations and improve accuracy.

#### Overall design.

The overall architecture is given by $B$ sequential FLARE blocks sandwiched between an input projection and an output projection, which are detailed in [Appendix A](https://arxiv.org/html/2508.12594v2#A1).
Such a design enables the model to efficiently integrate local and global information across multiple layers, making it well suited for large-scale, high-dimensional data such as point clouds.

### 3.3 Spectral analysis

The matrix $W_{h}$ in [Eq. 9](https://arxiv.org/html/2508.12594v2#S3.E9) represents the attention weights between $N$ tokens, where $[W_{h}]_{ij}$ quantifies how much token $j$ communicates to token $i$ within head $h$.
Since $W_{h}$ has rank at most $M$, the model can capture at most $M$ independent global communication patterns per head.
The eigenvalues of $W$ indicate the relative importance or energy of each latent dimension in forming the attention matrix $W$.
We apply an algorithm to obtain the eigen decomposition of $W$ in $\mathcal{O}(M^{3}+M^{2}N)$ time compared to $\mathcal{O}(N^{3})$ for a dense communication matrix.
The algorithm is predicated on computing the eigenspectra of an $M\times M$ matrix $JJ^{T}$
where $J\in\mathbb{R}^{M\times N}$ is chosen such that $J^{T}J$ is similar to $W$.
The algorithm is detailed in [Section B.1](https://arxiv.org/html/2508.12594v2#A2.SS1), and summarized in Algorithm 1.
[Figure 7](https://arxiv.org/html/2508.12594v2#A2.F7) presents the $M$ nonzero eigenvalues of $W_{h}$ for an FLARE model trained on the elasticity benchmark problem with 972 points per input.
The distinct spectra of the heads indicates that each head learns distinct attention patterns.

The eigenvalue analysis detailed in [Section B.2](https://arxiv.org/html/2508.12594v2#A2.SS2) shows that while FLARE provides capacity for rank-$M$ attention, the model learns to use only a small fraction of this in early blocks indicating effective compression.
In deeper blocks, more of the latent capacity is utilized, with diverse spectral profiles across heads, validating our design choice of independent head-wise projections.

## 4 Experiments

**Table 2: (Top) relative $L_{2}$ error ($\times 10^{-3}$) and (bottom) parameter count for different models across PDE benchmark problems. The best results (smallest error) are made bold, and the second best results are underlined. A backslash ($\backslash$) indicates that the model cannot be applied to the benchmark, and tilde ($\sim$) indicates that the model is prohibitively slow on the benchmark.**
| Model | Elasticity | Darcy | Airfoil | Pipe | DrivAerML-40k | LPBF |
| --- | --- | --- | --- | --- | --- | --- |
| Vanilla Transformer<br>(Vaswani et al., 2017) | 5.37<br>660k | 4.38<br>660k | 6.28<br>660k | $\sim$ | $\sim$ | $\sim$ |
| PerceiverIO<br>(Jaegle et al., 2021a) | 23.4<br>1.93m | 21.5<br>1.93m | 162<br>1.93m | 7.14<br>1.93m | 760<br>1.93m | 56.3<br>1.93m |
| GNOT<br>(Hao et al., 2023) | 13.3<br>4.87m | 16.9<br>4.90m | 103<br>4.90m | 5.89<br>4.90m | 115<br>4.87m | 24.3<br>4.87m |
| LNO<br>(Wang & Wang, 2024a) | 9.25<br>1.83m | 7.64<br>762k | 17.8<br>762k | 8.10<br>762k | 146<br>762k | 24.7<br>762k |
| Transolver w/o conv<br>(Wu et al., 2024) | 6.40<br>713k | 18.6<br>713k | 8.24<br>713k | 4.87<br>713k | 70.5<br>713k | 20.4<br>713k |
| Transolver with conv<br>(Wu et al., 2024) | $\backslash$ | 5.94<br>2.8m | 5.50<br>2.8m | 3.90<br>2.8m | $\backslash$ | $\backslash$ |
| FLARE (ours) | 3.38<br>592k | 5.10<br>691k | 4.28<br>691k | 2.85<br>625k | 60.8<br>691k | 18.5<br>625k |

### 4.1 PDE surrogate benchmarks

#### Benchmark problems.

We consider a diverse set of benchmark datasets ([Table 3](https://arxiv.org/html/2508.12594v2#A3.T3)) for regressing PDE solutions on point clouds spanning structured and unstructured grids with up to 50,000 points.
Note that FLARE is mesh-agnostic, and operates solely on the input point cloud.
The 2D elasticity, darcy, airfoil, and pipe benchmarks (Li et al., 2020; 2023b) cover a wide range of physical phenomena, and the 3D DrivAerML benchmark (Ashton et al., 2024) provides automotive aerodynamic simulations.
Additional details of the dataset are presented in [Section C.2](https://arxiv.org/html/2508.12594v2#A3.SS2).
We also introduce a 3D field-prediction benchmark derived from laser powder bed fusion (LPBF) simulations, with diverse 3D-printed parts containing up to 50,000 grid points (see [Appendix G](https://arxiv.org/html/2508.12594v2#A7)).

#### Baselines.

We compare FLARE with state-of-the-art PDE surrogates: generic attention models (vanilla Transformer (Vaswani et al., 2017), PerceiverIO (Jaegle et al., 2021a)); attention-based PDE surrogates (Transolver (Wu et al., 2024), LNO (Wang & Wang, 2024a)); and the neural operator GNOT (Hao et al., 2023). We exclude graph-based models because graph connectivity is unavailable for most problems. Transolver++ (Luo et al., 2025) is a concurrent work; although we tested the recently released official implementation with reported hyperparameters, preliminary results did not match those reported in the paper, so we do not include it in our comparisons. We follow the experimental setup of Transolver as it is the preeminent surrogate model and attempt to match its parameter count. Note that Transolver can be instantiated in two configurations: *without convolution*, where point-to-point communication relies solely on physics attention,
and *with convolution*, where convolution layers are added to inject information from neighboring points when the input grid is structured.
We evaluate these two configurations separately to isolate the impact of convolution versus physics attention.
In our model, we choose not to employ any convolution layers and rely entirely on FLARE for token mixing. Because these PDE problems are relatively small (up to 50,000 points), we train all models in FP32. As shown in [Figure 8](https://arxiv.org/html/2508.12594v2#A3.F8), the vanilla Transformer is drastically slower than FLARE and Transolver on large point clouds; accordingly, we evaluate it only on problems up to $\sim 10{,}000$ points.

#### Discussion.

The results in [Table 2](https://arxiv.org/html/2508.12594v2#S4.T2) clearly demonstrate that the proposed FLARE architecture achieves the lowest relative $L_{2}$ error across all but one benchmark PDE problems, outperforming both LNO and Transolver on every dataset.
Notably, FLARE also achieves these gains with a consistently lower parameter count than Transolver or LNO, highlighting its efficiency in addition to higher accuracy.
These results underline the robustness and versatility of FLARE across diverse problem settings.
We also note that the poor performance of Transolver without convolutions indicates that the inter-point communication via Transolver’s built-in physics-attention mechanism is not enough.
With convolutions, the input projections amass information from neighboring points, which in turn helps the physics attention learn the global structure.

Although the vanilla transformer is extremely effective on small-scale PDE problems, it becomes prohibitively slow on large point clouds due to its quadratic cost as illustrated in [Figure 2](https://arxiv.org/html/2508.12594v2#S2.F2).
On the contrary, PerceiverIO (with only a single encoding and decoding step) performs poorly even with $M=1,024$ latent tokens and $B=8$ latent self-attention blocks.
This validates our hypothesis that
multiple latent self-attention operations can be unnecessary and potentially suboptimal so long as the projections are sufficiently expressive.
This is because information loss during projection is not recoverable via latent self-attention alone.
Instead, performing multiple (head-wise) parallel projections and reconstructions directly between the input and latent sequences preserves expressivity while simplifying the architecture.

### 4.2 Field-prediction on million-point geometries

Figure: Figure 4: We train FLARE on the DrivAerML dataset (Ashton et al., 2024) with one million points per geometry on a single Nvidia H100 80GB GPU. We present (left) the test relative error, (middle) time per epoch (s), and (right) peak memory utilization (GB) as a function of the number of FLARE blocks ($B$) for different number of latent tokens ($M$).
Refer to caption: x3.png

Although the benchmark problems in [Section 4.1](https://arxiv.org/html/2508.12594v2#S4.SS1) represent a wide variety of PDE problems, they are relatively small compared to industrial use cases that demand PDE solutions on complex geometries with millions of grid points (Ashton et al., 2024).
So far, attention-based surrogate models have not been able to scale to million-scale regression problems due to quadratic time and memory complexity, as illustrated in [Figure 2](https://arxiv.org/html/2508.12594v2#S2.F2).
The flash attention (Dao et al., 2022) algorithm has alleviated the memory bottleneck thanks to online softmax computation; however, these methods remain impractical due to their long training times.
Furthermore, SOTA models such as Transolver and LNO cannot be implemented with off-the-shelf fused attention algorithms like flash attention because of the need to explicitly materialize the projection weights.

We demonstrate in [Figure 4](https://arxiv.org/html/2508.12594v2#S4.F4) that FLARE can scale to million-scale geometries by training on the DrivAerML dataset (Ashton et al., 2024) where each mesh is subsampled to contain one million points.
These calculations are performed in mixed precision on a single Nvidia H100 80GB GPU provisioned through Google Cloud Platform.
We note the clear trend in [Figure 4](https://arxiv.org/html/2508.12594v2#S4.F4) (left) that the error consistently decreases as we increase the number of FLARE blocks.
To our knowledge, this is the first attention-based neural surrogate model trained on one million points on a single GPU without memory offloading or distributed computing.

### 4.3 Model analysis and ablations

In this section, we scale FLARE to examine how expressivity trades off against time and memory complexity.
Complementary ablations in [Appendix F](https://arxiv.org/html/2508.12594v2#A6) analyze architectural choices such as the depth of residual MLPs in key/value projections, the design of the feedforward block, and the effect of head dimension on accuracy.

#### Time and memory complexity.

[Figure 2](https://arxiv.org/html/2508.12594v2#S2.F2) illustrates the time and memory complexity of a single forward and backward pass for different attention schemes on long sequences.
The experiment is done in mixed precision (FP16 in forward pass, FP32 in backward pass) using PyTorch’s autocast functionality with $C=128$ features and $H=8$ heads for all models.
The flash attention backend (Dao et al., 2022) is employed for SDPA wherever possible.

Although vanilla self-attention has the lowest memory cost thanks to the flash-attention algorithm, which eliminates the need to materialize the score matrices ($Q_{h}\cdot K_{h}^{T}$), its compute time still scales poorly with the sequence length.
In contrast, the compute time for FLARE exhibits strong scaling with sequence length.
Its memory requirement is marginally greater than vanilla attention due to the presence of deep residual networks for key/value projections, and due to the need to materialize $Z_{h}$, the latent sequence of $M$ tokens.
As these costs are marginal compared to the SDPA operation, the curves for different $M$ values of FLARE are somewhat overlapping.
Finally, the compute time for Physics Attention of Transolver (Wu et al., 2024) exhibits somewhat good scaling.
However, its memory cost and compute time blow up for large slice counts due to the need for materializing the projection matrices.

Figure: Figure 5: Effect of number of blocks ($B$) and number of latent tokens ($M$) on test accuracy on the elasticity (left) and darcy (right) test cases.
Refer to caption: x4.png

#### Number of blocks ( B B ) and latent tokens ( M M ).

[Figure 5](https://arxiv.org/html/2508.12594v2#S4.F5) presents the test relative error of FLARE on the elasticity (left) and darcy (right) benchmark datasets as a function of the number of blocks $(B)$ and the number of latent tokens $(M)$.
[Figure 4](https://arxiv.org/html/2508.12594v2#S4.F4) (left) presents the same for the DrivAerML dataset with one million points per geometry.
In all cases, we note the favorable trend that relative error consistently decreases as we increase the number of blocks.
Similarly, we observe that the relative error generally decreases with $M$, though the trend is not strictly monotonic.
In the elasticity problem, improvements with rank diminish rapidly, indicating that global communication in that problem is fundamentally low-rank.
On the other hand, increasing $M$ monotonically increases performance on the darcy problem, indicating that the problem is *rank-limited*.
This also explains why vanilla transformer with a full-rank attention pattern outperforms rank-deficient FLARE on the darcy problem.
However, the accuracy gain comes at the cost of greater latency as the vanilla transformer is $\sim 5\times$ slower than FLARE on the Darcy problem.
[Figure 4](https://arxiv.org/html/2508.12594v2#S4.F4) indicates that time per epoch (middle) and memory (right) scaling of FLARE with $B$ and $M$.
Here, increasing $M$ leads to increased latency, and that increasing $M$ does not come at the cost of greater memory requirements.

## 5 Conclusion

FLARE is a token mixing layer that bypasses the quadratic cost of self-attention by leveraging low-rankness.
Mechanically, FLARE routes attention through a fixed-size latent sequence via cross-attention projection and unprojection.
FLARE achieves SOTA accuracy on a set of diverse PDE benchmarks, and easily scales to PDE problems with million-scale geometries.

As transformers are the backbone of modern deep learning, we postulate that an efficient attention mechanism has several applications.
We also identify potential areas for improvement:
FLARE’s reliance on deep residual MLPs can introduce sequential bottlenecks and increase latency, suggesting that further speedups are possible by addressing this issue.
Additional enhancements include
(1) incrementally increasing the number of latent tokens during training;
(2) conditioning latent tokens on time for diffusion modeling;
and
(3) designing decoder-only variants for auto-regressive modeling.

## LLM Usage

LLMs assisted with copyediting (grammar/clarity) and with scaffolding boilerplate in our experimental scripts (e.g., debugging, refactors) inside an IDE with LLM support.
LLMs did *not* generate ideas, architectures, analyses, results, or datasets.
All LLM outputs were reviewed, edited, and tested by the authors, who take full responsibility for the content.

## Acknowledgments

This work was supported by the Air Force Research Laboratory under contract FA8650-21-F-5803 and PA Manufacturing Fellows Initiative. Zhang was supported in part by the National Science Foundation under grants CMMI-1953323 and CBET-2332084.
The authors appreciate the support of Camfer, Inc. (Camfer, 2025) and Professor Amir Barati Farimani in providing computing resources.
Part of this research was conducted on the Bridges-2 Supercomputer at the Pittsburgh Supercomputing Center, and parts were conducted using ORCHARD, a high-performance cloud computing cluster made available by Carnegie Mellon University.
The authors thank Jay Pathak of Ansys, Inc. (Ansys, 2025) for insightful discussions and Andrew Porco for his assistance with data processing.

## References

- Alkin et al. (2024)
Benedikt Alkin, Andreas Fürst, Simon Schmid, Lukas Gruber, Markus Holzleitner, and Johannes Brandstetter.
Universal physics transformers.
*arXiv e-prints*, pp. arXiv–2402, 2024.
- Ansys (2025)
Ansys.
Ansys – engineering simulation software.
[https://www.ansys.com/](https://www.ansys.com/), 2025.
Accessed: 2025-08-14.
- Ashton et al. (2024)
Neil Ashton, Charles Mockett, Marian Fuchs, Louis Fliessbach, Hendrik Hetmann, Thilo Knacke, Norbert Schonwald, Vangelis Skaperdas, Grigoris Fotiadis, Astrid Walle, et al.
Drivaerml: High-fidelity computational fluid dynamics dataset for road-car external aerodynamics.
*arXiv preprint arXiv:2408.11969*, 2024.
- Autodesk (2025)
Autodesk.
*NetFabb Simulation Utility and Local Simulation*.
Autodesk Inc., San Francisco, California, United States, 2025.
- Ba et al. (2016)
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.
Layer normalization.
*arXiv preprint arXiv:1607.06450*, 2016.
- Camfer (2025)
Camfer.
camfer – ai‑powered mechanical engineering.
[https://www.camfer.dev/](https://www.camfer.dev/), 2025.
Accessed: 2025‑08‑07.
- Cao (2021)
Shuhao Cao.
Choose a transformer: Fourier or galerkin.
*Advances in Neural Information Processing Systems*, 34:24924–24940, 2021.
- Chen et al. (2021)
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré.
Scatterbrain: Unifying sparse and low-rank attention.
*Advances in Neural Information Processing Systems*, 34:17413–17426, 2021.
- Dao et al. (2022)
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.
Flashattention: Fast and memory-efficient exact attention with io-awareness.
*Advances in Neural Information Processing Systems*, 35:16344–16359, 2022.
- Denlinger et al. (2014)
Erik R. Denlinger, Jeff Irwin, and Pan Michaleris.
Thermomechanical modeling of additive manufacturing large parts.
*Journal of Manufacturing Science and Engineering*, 136(6):061007, 10 2014.
ISSN 1087-1357.
doi: 10.1115/1.4028669.
URL [https://doi.org/10.1115/1.4028669](https://doi.org/10.1115/1.4028669).
- Devlin et al. (2019)
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.
BERT: Pre-training of deep bidirectional transformers for language understanding.
In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), *Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)*, pp. 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
doi: 10.18653/v1/N19-1423.
URL [https://aclanthology.org/N19-1423/](https://aclanthology.org/N19-1423/).
- Dosovitskiy et al. (2020)
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, G Heigold, S Gelly, et al.
An image is worth 16x16 words: Transformers for image recognition at scale.
In *International Conference on Learning Representations*, 2020.
- Elrefaie et al. (2024)
Mohamed Elrefaie, Angela Dai, and Faez Ahmed.
Drivaernet: A parametric car dataset for data-driven aerodynamic design and graph-based drag prediction.
In *International Design Engineering Technical Conferences and Computers and Information in Engineering Conference*, volume 88360, pp. V03AT03A019. American Society of Mechanical Engineers, 2024.
- Ferguson et al. (2025)
Kevin Ferguson, Yu-hsuan Chen, Yiming Chen, Andrew Gillman, James Hardin, and Levent Burak Kara.
Topology-agnostic graph U-Nets for scalar field prediction on unstructured meshes.
*Journal of Mechanical Design*, 147(4):041701, 2025.
- Hao et al. (2023)
Zhongkai Hao, Zhengyi Wang, Hang Su, Chengyang Ying, Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song, and Jun Zhu.
GNOT: A general neural operator transformer for operator learning.
In *International Conference on Machine Learning*, pp. 12556–12569. PMLR, 2023.
- He et al. (2016)
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.
Deep residual learning for image recognition.
In *Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition*, pp. 770–778, 2016.
- Hendrycks & Gimpel (2016)
Dan Hendrycks and Kevin Gimpel.
Gaussian error linear units (GELUs).
*arXiv preprint arXiv:1606.08415*, 2016.
- Hu et al. (2022)
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.
LoRA: Low-rank adaptation of large language models.
In *International Conference on Learning Representations*, 2022.
URL [https://openreview.net/forum?id=nZeVKeeFYf9](https://openreview.net/forum?id=nZeVKeeFYf9).
- Jaegle et al. (2021a)
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al.
Perceiver IO: A general architecture for structured inputs & outputs.
*arXiv preprint arXiv:2107.14795*, 2021a.
- Jaegle et al. (2021b)
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira.
Perceiver: General perception with iterative attention.
In *International Conference on Machine Learning*, pp. 4651–4664. PMLR, 2021b.
- Kitaev et al. (2020)
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya.
Reformer: The efficient transformer.
*arXiv preprint arXiv:2001.04451*, 2020.
- Kovachki et al. (2023)
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.
Neural operator: Learning maps between function spaces with applications to PDEs.
*Journal of Machine Learning Research*, 24(89):1–97, 2023.
- Lambourne et al. (2021)
Joseph G. Lambourne, Karl D.D. Willis, Pradeep Kumar Jayaraman, Aditya Sanghi, Peter Meltzer, and Hooman Shayani.
Brepnet: A topological message passing system for solid models.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)*, pp. 12773–12782, June 2021.
- Li et al. (2022)
Zijie Li, Kazem Meidani, and Amir Barati Farimani.
Transformer for partial differential equations’ operator learning.
*arXiv preprint arXiv:2205.13671*, 2022.
- Li et al. (2023a)
Zijie Li, Dule Shu, and Amir Barati Farimani.
Scalable transformer for PDE surrogate modeling.
*Advances in Neural Information Processing Systems*, 36:28010–28039, 2023a.
- Li et al. (2020)
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.
Fourier neural operator for parametric partial differential equations.
*arXiv preprint arXiv:2010.08895*, 2020.
- Li et al. (2021)
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.
Fourier neural operator for parametric partial differential equations, 2021.
URL [http://arxiv.org/abs/2010.08895](http://arxiv.org/abs/2010.08895).
arXiv:2010.08895.
- Li et al. (2023b)
Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar.
Fourier neural operator with learned deformations for PDEs on general geometries.
*Journal of Machine Learning Research*, 24(388):1–26, 2023b.
- Li et al. (2023c)
Zongyi Li, Nikola Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Otta, Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzadenesheli, et al.
Geometry-informed neural operator for large-scale 3D PDEs.
*Advances in Neural Information Processing Systems*, 36:35836–35854, 2023c.
- Liang et al. (2019)
Xuan Liang, Qian Chen, Lin Cheng, Devlin Hayduke, and Albert C To.
Modified inherent strain method for efficient prediction of residual deformation in direct metal laser sintered components.
*Computational Mechanics*, 64(6):1719–1733, 2019.
- Liu et al. (2024)
Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al.
DeepSeek-V2: A strong, economical, and efficient mixture-of-experts language model.
*arXiv preprint arXiv:2405.04434*, 2024.
- Loshchilov & Hutter (2019)
Ilya Loshchilov and Frank Hutter.
Decoupled weight decay regularization, 2019.
URL [http://arxiv.org/abs/1711.05101](http://arxiv.org/abs/1711.05101).
arXiv:1711.05101.
- Lu et al. (2021)
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis.
Learning nonlinear operators via deeponet based on the universal approximation theorem of operators.
*Nature Machine Intelligence*, 3(3):218–229, 2021.
- Luo et al. (2025)
Huakun Luo, Haixu Wu, Hang Zhou, Lanxiang Xing, Yichen Di, Jianmin Wang, and Mingsheng Long.
Transolver++: An accurate neural solver for pdes on million-scale geometries.
In *Forty-second International Conference on Machine Learning*, 2025.
- Paszke (2019)
A Paszke.
Pytorch: An imperative style, high-performance deep learning library.
*arXiv preprint arXiv:1912.01703*, 2019.
- Pfaff et al. (2020)
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter Battaglia.
Learning mesh-based simulation with graph networks.
In *International Conference on Learning Representations*, 2020.
- Qian et al. (2025)
Kuanren Qian, Genesis Omana Suarez, Toshihiko Nambara, Takahisa Kanekiyo, and Yongjie Jessica Zhang.
High-throughput machine learning framework for predicting neurite deterioration using MetaFormer attention.
*Computer Methods in Applied Mechanics and Engineering*, 442:118003, 2025.
- Reijonen et al. (2024)
Joni Reijonen, Alejandro Revuelta, Sini Metsä-Kortelainen, and Antti Salminen.
Effect of hard and soft re-coater blade on porosity and processability of thin walls and overhangs in laser powder bed fusion additive manufacturing.
*The International Journal of Advanced Manufacturing Technology*, 130(5):2283–2296, 2024.
- Smith & Topin (2019)
Leslie N Smith and Nicholay Topin.
Super-convergence: Very fast training of neural networks using large learning rates.
In *Artificial Intelligence and Machine Learning for Multi-Domain Operations Applications*, volume 11006, pp. 369–386. SPIE, 2019.
- Vaswani et al. (2017)
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin.
Attention is all you need.
*Advances in Neural Information Processing Systems*, 30, 2017.
- Wang et al. (2020)
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma.
Linformer: Self-attention with linear complexity.
*arXiv preprint arXiv:2006.04768*, 2020.
- Wang & Wang (2024a)
Tian Wang and Chuang Wang.
Latent neural operator for solving forward and inverse pde problems.
*arXiv preprint arXiv:2406.03923*, 2024a.
- Wang & Wang (2024b)
Tian Wang and Chuang Wang.
Latent neural operator for solving forward and inverse PDE problems, June 2024b.
URL [http://arxiv.org/abs/2406.03923](http://arxiv.org/abs/2406.03923).
arXiv:2406.03923 [cs, math].
- Wu et al. (2024)
Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, and Mingsheng Long.
Transolver: A fast transformer solver for pdes on general geometries.
*arXiv preprint arXiv:2402.02366*, 2024.
- Xiong et al. (2020)
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Yan Liu.
On layer normalization in the transformer architecture.
In *Proceedings of the 37th International Conference on Machine Learning*, pp. 10524–10533, 2020.
- Xiong et al. (2021)
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh.
Nyströmformer: A nyström-based algorithm for approximating self-attention.
In *Proceedings of the AAAI conference on artificial intelligence*, volume 35, pp. 14138–14148, 2021.
- Yun et al. (2019)
Seongjun Yun, Minbyul Jeong, Raehyun Kim, Jaewoo Kang, and Hyunwoo J Kim.
Graph transformer networks.
*Advances in Neural Information Processing Systems*, 32, 2019.

## Appendix A Architecture details

### A.1 Input/ output projection

#### ResMLP.

We implement a residual MLP block to serve as a flexible non-linear function approximator.
Given input/output dimensions $C_{\text{i}}$ and $C_{\text{o}}$, the layer first applies a linear transformation to a hidden space of size $C_{\text{h}}$, followed by $L$ residual layers, each consisting of a linear layer with GELU activation (Hendrycks & Gimpel, 2016).
These are the only instances of pointwise nonlinear activations in the model.
An optional input residual connection is applied after the first layer when $C_{\text{i}}=C_{\text{h}}$, and an optional output residual connection is applied at the end when
$C_{\text{h}}=C_{\text{o}}$.
The final output is projected to dimension $C_{\text{o}}$ via a linear layer.
This design allows control over depth and expressivity while preserving stability through residual connections.

#### Input projection.

The input projection consists of a ResMLP with $L=2$, $C_{i}$ is the input feature dimension, and $C_{\text{h}}=C_{\text{o}}$ are set to $C$, the feature dimension of the model.

#### Output projection.

The output projection consists of a Layer Norm (Ba et al., 2016) followed by a ResMLP with $C_{\text{i}}=C$, $L=2$, and $C_{\text{o}}$ is the output label dimension.

### A.2 FLARE block

The FLARE block illustrated in [Figure 1](https://arxiv.org/html/2508.12594v2#S1.F1), and detailed in [Section 3](https://arxiv.org/html/2508.12594v2#S3), consists of the pointwise ResMLP layer, and the FLARE token mixer.
For the ResMLP, we set $C_{\text{i}}=C_{\text{h}}=C_{\text{o}}=C$, set the number of layers to $3$, and allow residuals to flow through the entire block.

#### FLARE.

FLARE consists of two ResMLPs for key/value projections, the token operation described in [Figure 3](https://arxiv.org/html/2508.12594v2#S3.F3), and an output projection.
For the key/value projections, we set $C_{\text{i}}=C_{\text{h}}=C_{\text{o}}=C$, $L=3$, and allow residuals to flow through the entire block.
[Figure 6](https://arxiv.org/html/2508.12594v2#A2.F6) presents a mathematically equivalent PyTorch implementation for multi-head token-mixing operation without the fused SDPA kernels.
The primary memory bottleneck in this implementation is materializing the $M\times N$ encoding weights and the $N\times M$ decoding weights.
Its storage requirement is, thus, $\mathcal{O}(MN)$.
Finally, the output projection is set to a single linear layer.

## Appendix B Spectral analysis

Figure: Figure 6: Pseudocode of FLARE if attention kernel is not available. See [Figure 3](https://arxiv.org/html/2508.12594v2#S3.F3) for efficient implementation.

### B.1 Eigenanalysis procedure

We exploit the low-rank structure of the global communication matrix
$W=W_{\text{decode}}\cdot W_{\text{encode}}\in\mathbb{R}^{N\times N}$
to obtain its eigen decomposition in $\mathcal{O}(M^{3}+M^{2}N)$ time, compared to the $\mathcal{O}(N^{3})$ cost for a dense communication matrix.
We find the eigen-decomposition of the FLARE attention matrix without actually forming the $N\times N$ matrix.
We first note that $W_{\text{encode}}$ and $W_{\text{decode}}$ can be written in terms of the exponentiated score matrix
$A=\exp(Q\cdot K^{T})\in\mathbb{R}^{M\times N}$ as

$$ $\displaystyle W_{\text{encode}}=\Lambda_{M}\cdot A,\text{ and }W_{\text{decode}}=\Lambda_{N}\cdot A^{T},$ (11) $$

where
$\Lambda^{M}\in\mathbb{R}^{M\times M}$,
and
$\Lambda^{N}\in\mathbb{R}^{N\times N}$
are diagonal matrices whose entries are

$$ $\displaystyle[\Lambda_{M}]_{m}=\frac{1}{\sum_{n=1}^{N}[A]_{m,n}},\text{ and }[\Lambda_{N}]_{n}=\frac{1}{\sum_{m=1}^{M}[A]_{m,n}}.$ (12) $$

Thus we have

$$ $\displaystyle W=\Lambda_{N}A^{T}\Lambda_{M}A$ (13) $$

as the low-rank attention matrix.
We observe that $W$ is similar to $J^{T}J\in\mathbb{R}^{N\times N}$
where
$J=\Lambda_{M}^{1/2}A\Lambda_{N}^{1/2}\in\mathbb{R}^{M\times N}$.
This is because

$$ $\displaystyle W$ $\displaystyle=\Lambda_{N}A^{T}\Lambda_{M}A=\underbrace{(\Lambda_{N}^{1/2}\Lambda_{N}^{1/2})}_{\Lambda_{N}}A^{T}\underbrace{(\Lambda_{M}^{1/2}\Lambda_{M}^{1/2})}_{\Lambda_{M}}A\underbrace{(\Lambda_{N}^{-1/2}\Lambda_{N}^{1/2})}_{I_{N}}$ (14) $\displaystyle=\Lambda_{N}^{1/2}\underbrace{(\Lambda_{N}^{1/2}A^{T}\Lambda_{M}^{1/2})}_{J^{T}}\underbrace{(\Lambda_{M}^{1/2}A\Lambda_{N}^{1/2})}_{J}\Lambda_{N}^{-1/2}.$ $$

Thus $J^{T}J$ is symmetric, positive semi-definite, with rank at most $M$.
Now suppose a singular value decomposition of $J$ as
$J=U\Sigma V^{T}$ where $U\in\mathbb{R}^{M\times M}$ and $V\in\mathbb{R}^{N\times M}$ are the matrices whose columns are the left and right singular vectors of $J$ respectively, and $\Sigma\in\mathbb{R}^{M\times M}$ is the diagonal matrix of singular values.
Then, we obtain

$$ $\displaystyle J^{T}J=V\Sigma\underbrace{U^{T}U}_{I_{M}}\Sigma V^{T}=V\Sigma^{2}V^{T},$ (15) $$

and

$$ $\displaystyle W=\Lambda_{N}^{1/2}V\Sigma^{2}V^{T}\Lambda_{N}^{-1/2}.$ (16) $$

Post-multiplying both sides by $\Lambda_{N}^{1/2}V$, we have

$$ $\displaystyle W(\Lambda_{N}^{1/2}V)=(\Lambda_{N}^{1/2}V)\Sigma^{2}.$ (17) $$

Therefore, the $M$ nonzero eigenvalues of $W$ are the squares of the singular values of $J$,
and the corresponding eigenvectors are the columns of $\Lambda_{N}^{1/2}V$.
Obtaining the eigenvalues and eigenvectors of $W$ this way requires the singular value decomposition of $J\in\mathbb{R}^{M\times N}$.
We can do better by relating $V$ and $\Sigma$ to the eigen decomposition of $JJ^{T}$.
Consider the matrix $JJ^{T}\in\mathbb{R}^{M\times M}$ with singular value decomposition, we have

$$ $\displaystyle JJ^{T}=U\Sigma\underbrace{V^{T}V}_{I_{M}}\Sigma U^{T}=U\Sigma^{2}U^{T}.$ (18) $$

We note that the nonzero eigenvalues of $W$ are the same as the singular values of $JJ^{T}$.
To obtain the eigenvectors of $W$, we need an expression for $V$ in terms of $U$, $J$, and $\Sigma$.
We do so by noting that

$$ $\displaystyle J^{T}U=(V\Sigma U^{T})U=V\Sigma\implies V=J^{T}U\Sigma^{-1}.$ (19) $$

Therefore, the eigenvectors of $W$ are

$$ $\displaystyle\Lambda_{N}^{1/2}V=\Lambda_{N}^{1/2}J^{T}U\Sigma^{-1}.$ (20) $$

Figure: Algorithm 1 Eigenvalues and Eigenvectors from $Q,\,K$

To find the eigenvalues and eigenvectors of $W$, one only needs to compute the eigen-decomposition of the $M\times M$ matrix $J^{T}J$.
The overall algorithm, summarized in Algorithm 1, takes
$\mathcal{O}(M^{3}+NM^{2})$ time where the $\mathcal{O}(M^{3})$ is for computing the SVD of $JJ^{T}$.

Figure: Figure 7: $M=64$ nonzero eigenvalues of $W_{h}$ ([Eq. 9](https://arxiv.org/html/2508.12594v2#S3.E9)) for $H=8$ heads in different blocks of FLARE with $B=8$ blocks $C=64$ features trained on the elasticity dataset ($972$ points per input). The number of latents ($M$) is depicted with a dashed red line, and FP32 precision with a dashed black line.
Refer to caption: x6.png

### B.2 Qualitative Analysis

The matrix $JJ^{T}$, being $M\times M$, captures the structure of this latent space and provides insights into how these $M$ dimensions contribute to the attention mechanism.
Large eigenvalues correspond to dominant latent dimensions that contribute significantly to attention patterns.
If some eigenvalues are small or zero, those latent dimensions contribute little, suggesting redundancy in the latent space.
The number of nonzero eigenvalues gives the effective rank of $W$, which reflects how many independent patterns the attention mechanism captures.

[Figure 7](https://arxiv.org/html/2508.12594v2#A2.F7) presents the $M=64$ nonzero eigenvalue spectra of an FLARE model trained on the elasticity dataset with $M=64$ latents.
Some observations are as follows.
In all blocks, especially block 1, the eigenvalues drop sharply within the first $10-20$ indices.
This indicates that even though the communication matrices $W_{h}$ could have rank up to $M=64$, most of the energy (information) is captured by a much smaller subset of modes.
This result validates the hypothesis that the global communication pattern is inherently low-rank.

We also observe that the eigenvalue curves in block 5 and block 8 decay more slowly, retaining more moderate-magnitude eigenvalues beyond index 20–30.
This indicates that as depth increases, the attention mechanism seems to leverage more of the latent space — i.e., the effective rank increases.
This shows that deeper layers learn richer global dependencies, and the model may be using more of the projection capacity in later blocks.

Finally, we note that the curves for different heads (colors) have distinct decay patterns, especially in later blocks.
This reinforces the claim that separate projection matrices per head enables specialized routing.
This supports the idea that FLARE benefits from head-wise diversity, rather than using shared projections like in Transolver (Wu et al., 2024).

## Appendix C Benchmarking and Comparison

### C.1 Benchmark metrics

The primary evaluation metric for all benchmarks is the relative $\mathcal{L}_{2}$ error, which quantifies the normalized discrepancy between the predicted solution $\hat{u}$ and the ground truth solution $u$ over all points in the domain. For a given test sample, the relative $L_{2}$ error is defined as:

$$ $\text{Relative }L_{2}=\frac{\|\hat{u}-u\|_{2}}{\|u\|_{2}}$ (21) $$

where $\|\cdot\|_{2}$ denotes the standard Euclidean norm.
For datasets where each sample consists of $N$ points (or grid locations), this expands to:

$$ $\text{Relative }L_{2}=\frac{\left(\sum_{i=1}^{N}(\hat{u}_{i}-u_{i})^{2}\right)^{1/2}}{\left(\sum_{i=1}^{N}u_{i}^{2}\right)^{1/2}}.$ (22) $$

The reported metric is averaged over all test samples.

**Table 3: Summary of PDE benchmarks.**
| Benchmark | #Dim | Grid Type | #Points | #Input/ Output<br>Features | #Train/ Test<br>Cases |
| --- | --- | --- | --- | --- | --- |
| Elasticity | 2D | Unstructured | 972 | 2 / 1 | 1000 / 200 |
| Darcy |  |  | 7,225 |  |  |
| Airfoil | 2D | Structured | 11,271 | 2 / 1 | 1000 / 200 |
| Pipe |  |  | 16,641 |  |  |
| DrivAerML-40k | 3D | Unstructured | 40,000 | 3 / 1 | 387 / 97 |
| LPBF |  |  | 1,000–50,000 |  | 1100 / 290 |

### C.2 Benchmark datasets

We evaluate all models on six benchmark datasets and our proposed AM dataset, each designed to assess generalization, scalability, and robustness to domain irregularity in PDE surrogate modeling. A summary is provided in [Table 3](https://arxiv.org/html/2508.12594v2#A3.T3).

#### Elasticity.

This benchmark estimates the inner stress distribution of elastic materials based on their structure. Each sample consists of a tensor of shape $972\times 2$ representing the 2D coordinates of discretized points, and the output is the stress at each point ($972\times 1$). The dataset contains 1,000 training and 200 test samples with different material structures (Li et al., 2023b).

#### Darcy.

This benchmark models fluid flow through a porous medium.
The porous structure is discretized on a $421\times 421$ regular grid, downsampled to $85\times 85$ for main experiments.
The output is the fluid pressure at each grid point.
There are 1,000 training and 200 test samples with varying medium structures (Li et al., 2021).

#### Airfoil.

This task estimates the Mach number distribution around airfoil shapes. The input geometry is discretized into a structured mesh of shape $221\times 51$, representing deformations of the NACA-0012 profile, and the output is the Mach number at each mesh point. The dataset includes 1,000 training and 200 test samples with unique airfoil designs (Li et al., 2023b).

#### Pipe.

This benchmark predicts horizontal fluid velocity in pipes with varying internal geometries. Each sample is represented as a $129\times 129\times 2$ tensor encoding mesh point positions, with the output being the velocity at each mesh point ($129\times 129\times 1$). The dataset consists of 1,000 training and 200 test samples, with pipe shapes generated by controlling the centerline (Li et al., 2023b).

#### DrivAerML-40k.

The DrivAerML dataset (Ashton et al., 2024) has high-fidelity automotive aerodynamic simulations, featuring 500 parametrically morphed DrivAer vehicles.
CFD simulations are performed on 160 million volumetric mesh grids using hybrid RANS-LES, the highest-fidelity scale-resolving CFD approach routinely deployed by the automotive industry (Ashton et al., 2024).
Each sample in the dataset includes a surface mesh with approximately 8.8 million points and corresponding pressure values.
Since no official dataset split is provided, an 80/20 random split is used, with 40,000 points sampled per case for training and evaluation.

#### LPBF.

We introduce the laser powder bed fusion (LPBF) additive manufacturing dataset which involves field prediction on complex geometries.
In metal additive manufacturing, variations in design geometry can affect the dimensional accuracy of the part and lead to shape distortions.
We perform several LPBF process simulations of a set of geometries to obtain the deformation field over the geometry.
We select a subsample of the dataset with up to 50,000 points per geometry and divide it into 1,100 training and 290 cases.
We train models to learn the $Z$ (vertical) component of the deformation field at each grid point.
Additional details are provided in [Appendix G](https://arxiv.org/html/2508.12594v2#A7).

### C.3 Benchmark models and training details

**Table 4: Standard training configuration on PDE datasets. Identical values are grouped for clarity.**
| Dataset | Batch Size | Weight Decay | Learning Rate | Epochs | Loss |
| --- | --- | --- | --- | --- | --- |
| Elasticity | 2 | $10^{-5}$ | $10^{-3}$ | 500 | Rel. $L_{2}$ |
| Darcy |  |  |  |  | Rel. $L_{2}$ + 0.1 $L_{\text{g}}$ |
| Airfoil | 2 | $10^{-5}$ | $10^{-3}$ | 500 | Rel. $L_{2}$ |
| Pipe |  |  |  |  | Rel. $L_{2}$ |
| DrivAerML | 1 | $10^{-2}$ | $10^{-3}$ | 500 | Rel. $L_{2}$ |
| LPBF |  | $10^{-4}$ |  | 250 |  |

**Table 5: Model configurations for FLARE for PDE datasets. Identical values are grouped for clarity.**
| Dataset | #Heads ($H$) | #Latents ($M$) | #Blocks ($B$) | #Features $(C)$ |
| --- | --- | --- | --- | --- |
| Elasticity | 8 | 64 | 8 | 64 |
| Darcy | 16 | 256 |  |  |
| Airfoil | 8 | 256 | 8 | 64 |
| Pipe | 8 | 128 |  |  |
| DrivAerML-40k | 8 | 256 | 8 | 64 |
| LPBF | 16 | 256 |  |  |

We follow the recommended hyperparameter configuration for LNO, Transolver, and GNOT wherever possible.
For consistency, we set the number of blocks to $B=8$,
and strive to match the parameter counts of Transolver without (w/o) convolution for all models.
As such, the hidden feature dimension for vanilla transformer, is set to $C=80$.
Its head dimension is set to $D=16$ and MLP ratio is 4 which is typical for transformers (Vaswani et al., 2017).
For FLARE, we set the hidden feature dimension to $C=64$, and use $H=8$ or $16$ heads, which result in a head dimension of $D=8$ or $4$ respectively.
The number of latents is chosen from $M\in\{64,\,128,\,256\}$ depending on the problem.
As PerceiverIO was not designed to be a surrogate model,
we generously set its channel dimension to $C=128$, and number of latents to $M=1,024$.
Furthermore, the input and output projections for vanilla transformer, perceiver, and FLARE are held consistent to facilitate an equitable comparison of their point-to-point communication schemes.

We evaluate Transolver, GNOT with the hyperparameter configurations provided in Wu et al. (2024) and LNO on the ones provided in Wang & Wang (2024a)
on the 2D test cases.
For the remaining problems, we choose the best performing parameter set from the 2D cases.

Unless otherwise stated, all models are trained with the AdamW (Loshchilov & Hutter, 2019) optimizer ($\beta_{1}=0.9$, $\beta_{2}=0.999$) to minimize the relative $L_{2}$ error.
We use the OneCycleLR (Smith & Topin, 2019) scheduler with 10% of epochs dedicated to warming-up to a learning rate of $10^{-3}$, followed by cosine decay.
We train on LPBF for 250 epochs, and 500 epochs for all other models.
Note that we use gradient clipping with max_norm$=1.0$ unless otherwise stated.
Unless otherwise stated, the weight decay regularization parameter is set to $10^{-5}$ for the 2D test cases,
$10^{-4}$ for DrivAerML-40k,
and $10^{-4}$ for LPBF.
The batch size is set to 2 for the 2D problems and 1 for the 3D problems unless otherwise stated.
Unless otherwise stated, all models are trained in full precision (FP32).

#### Vanilla transformer.

For the vanilla transformer, we set the hidden feature dimension to $C=80$, and choose $H=5$ heads so that the head dimension $D=16$.
The number of blocks is set to $B=8$, and the MLP ratio for the feedforward block is set to 4.
The vanilla transformer can be prohibitively slow for test cases with over 10,000 test points.

#### PerceiverIO.

For PerceiverIO, we use $B=8$, $C=128$, $H=8$, and set the latent sequence length to $M=1,024$.

#### Transolver.

Transolver (Wu et al. (2024)) introduces Physics-Attention: each mesh point is softly assigned to a few learnable *slices*, shrinking thousands of points to only tens of tokens.
Self-attention runs on this compact set and the tokens are then desliced back to points.
Following the hyperparameter recommendations in the code of Wu et al. (2024),
Transolver is trained with 30% of the steps dedicated to warm-up, and gradient clipping with max_norm $=0.1$.
The recommended batch size for Transolver is 1 for elasticity, and 4 for the remaining 2D problems.
For the 3D problems, we set the batch size to 1.

#### Latent Neural Operator (LNO).

LNO (Wang & Wang (2024b)) moves computation into a small latent space.
An embedder lifts the input field, cross-attention compresses $N$ points into $M$ latent tokens, transformers act solely on these tokens, and a decoder maps them to any query location.
In line with the recommended hyperparameter configuration in Wang & Wang (2024b),
we set $\beta_{2}=0.99$ in the AdamW optimizer,
do warmup for the first 20% of epochs,
and clip gradient norms greater than 1,000.
The number of hidden features are set to $C=192$ for elasticity, and $C=128$ for all other test cases.
The number of residual layers is set to 3 for elasticity and 4 for other test cases.
The number of latent self attention blocks is 8 for pipe and airfoil test cases and 4 for all other cases.
The number of latent modes is set to 256 for all test cases.
The batch size during training is set to 4 for the 2D test cases and 1 for the 3D test cases.
The LNO code also recommends weight decay regularization of $5\cdot 10^{-5}$ for the 2D test cases.

In running our experiments, we noticed discrepancies between our LNO results and those presented in their article (Wang & Wang, 2024a).
Upon further investigation, we found that the datasets used in the LNO paper and in the original Transolver paper are not the same.
For example, in the Elasticity dataset, LNO was trained and tested on a $1000/1000$ split, whereas Transolver used a $1000/200$ split.
In the Darcy dataset, LNO employed a higher resolution of $241\times 241$, compared to $85\times 85$ in Transolver.
Because these differences make direct comparison unreliable, all models (including LNO and Transolver) were re-trained and evaluated on the standardized Transolver splits and resolutions to ensure fairness.

#### General Neural Operator Transformer (GNOT).

GNOT (Hao et al. (2023)) employs heterogeneous normalized attention, separately normalizing keys and values, to fuse multiple input fields on irregular meshes. A learnable geometric gate decomposes the domain and routes tokens to scale-specific expert MLPs. Linear cost attention plus this gating scales to large problems and surpasses earlier operator learners.

Following the hyperparameter recommendations outlined in the code of Wu et al. (2024),
GNOT is trained with 30% of the steps dedicated to warm-up, and gradient clipping with max_norm = $0.1$.
The recommended batch size for Transolver is 1 for elasticity, and 4 for the remaining 2D problems.
The batch size is set to 2 for elasticity, 4 for the remaining 2D benchmarks and 1 for the 3D benchmarks.

#### FLARE.

For all problems, we employ $B=8$ blocks with a feature dimension of $C=64$.
We set the number of residual layers and the number of key/value projection layers to 3, and vary the head dimension as
$D\in\{4,8\}$,
the number of latent tokens as
$M\in\{64,128,256\}$
The hyperparameters for each test case are presented in [Table 5](https://arxiv.org/html/2508.12594v2#A3.T5).

Figure: Figure 8: Execution times in FP32 for a single vanilla self-attention layer, a physics attention layer, and FLARE. The models are set to have approximately the same number of parameters as in [Section 4.1](https://arxiv.org/html/2508.12594v2#S4.SS1). This calculation is performed on a single H100 80GB GPU. Note that the curves for FLARE are somewhat overlapping.
Refer to caption: x7.png

## Appendix D Field-prediction on million-point geometries

All FLARE models in this study are trained in mixed precision to take advantage of the flash attention algorithm.
We train for 500 epochs with the OneCycleLR scheduler (Smith & Topin, 2019) where the first $5\%$ of epochs are spent warming up to a learning rate of $5\cdot 10^{-4}$ followed by cosine decay.

## Appendix E Discussion on the design principles of FLARE.

#### Latent tokens enable gather-scatter communication.

In FLARE, information flows through latent tokens by first gathering from the input sequence and then scattering back.
The encoding step can be understood as a gather (all-reduce) operation, where each latent token pools information from the input according to its learned query pattern.
Formally, for latent query $q_{m}$,

$$ $z_{m}=\sum_{n=1}^{N}\frac{\exp(q_{m}\cdot k_{n})}{\sum_{n^{\prime}=1}^{N}\exp(q_{m}\cdot k_{n^{\prime}})}\,v_{n},\quad m=1,\dots,M,$ (23) $$

$z_{m}$ aggregates input values $v_{n}$ with convex weights.
When the similarity scores are sharp, $z_{m}$ emphasizes a few dominant inputs; when they are flatter, $z_{m}$ acts like a specialized pooling token that averages a select set of tokens.
Across $M$ latents, this yields a compact set of global descriptors, each specializing in pooling different aspects of the input.

The decoding step is the dual scatter (broadcast) operation, but importantly, the broadcast is selective: each latent $z_{m}$ contributes only to the input tokens that assigned it high weight in the encoding step.
Concretely,

$$ $y_{n}=\sum_{m=1}^{M}\frac{\exp(k_{n}\cdot q_{m})}{\sum_{m^{\prime}=1}^{M}\exp(k_{n}\cdot q_{m^{\prime}})}\,z_{m},\quad n=1,\dots,N,$ (24) $$

the output token $y_{n}$ only receives substantial information from the latents whose query pattern matches its key $k_{n}$ strongly.
In this sense, each latent acts as both a selective pooling hub and broadcaster, routing information back along the pathways that originally recruited it.
Together, these gather–scatter stages form a low-rank, butterfly-like transformation: first contracting $N$ tokens into $M\ll N$ global features, and then expanding back to $N$ outputs, enabling efficient, yet expressive, global communication.

#### Symmetry in latent token communication.

The compression–expansion attention structure in FLARE is most effective when the encoding and decoding operators are structurally aligned.
Specifically, $W_{\text{encode},h}=\mathrm{softmax}(Q_{h}\cdot K_{h}^{T})$ and $W_{\text{decode},h}=\mathrm{softmax}(K_{h}\cdot Q_{h}^{T})$ are transposes of each other up to diagonal scaling.
We experimented with breaking this symmetry by using distinct query–key pairs for encoding and decoding, but observed no accuracy gains.
This suggests that the near-adjoint relationship between $W_{\text{encode}}$ and $W_{\text{decode}}$, both derived from the same parameters, may provide a form of mathematical optimality, ensuring stable information flow through latent tokens while reducing representational redundancy.

#### Tradeoff between query dynamics and key/value expressivity.

An important design choice in FLARE is to use fixed, input-independent queries $Q$, which constrain the flexibility of the attention pattern.
Detaching $Q$ from $X$ allows for the low-rank communication structure, but requires compensating expressivity in the key and value projections.
In practice, we find that deeper residual MLPs for $K$ and $V$ are crucial for capturing rich, feature interactions under this constraint.
Conversely, one could imagine making $Q$ dynamic by conditioning it on $X$, which would shift more of the modeling burden to the query side and potentially allow shallower $K,V$ projections.
Thus, FLARE embodies a clear tradeoff: fixing $Q$ encourages stability and efficiency, but places greater importance on the depth and expressivity of the key/value encoders.

## Appendix F Model analysis and ablations

Figure: Figure 9: (Left) effect of the number of residual layers in key/ value projection, and (right) of residual layers in residual block on test accuracy. In both cases, deeper networks lead to greater accuracy.
Refer to caption: x8.png

#### ResMLP depth: key/value projections.

A substantial distinction between vanilla self-attention and FLARE is the introduction of deep residual blocks for key and value projections in place of simple linear layers.
In standard self-attention, queries ($Q_{h}$) and keys ($K_{h}$) determine the global communication pattern through $W_{h}=\mathrm{softmax}(Q_{h}K_{h}^{T}/s)$, while values ($V_{h}$) carry the information to be communicated.
All three are typically computed as shallow linear projections of the input $X$.
In contrast, FLARE computes the attention pattern as $W_{h}=\mathrm{softmax}(K_{h}Q_{h}^{T})\cdot\mathrm{softmax}(Q_{h}K_{h}^{T})$, where the query embeddings $Q_{h}$ are learned parameters independent of the input.
This makes the attention pattern less dynamic, motivating architectural modifications to enhance flexibility.

To address this, we replace the linear key and value projections with deep residual MLPs.
Using residual networks for key and value encodings allows each token to learn richer and more structured features rather than shallow embeddings, which is particularly crucial in FLARE since the queries are fixed and cannot adapt to the input.
[Figure 9](https://arxiv.org/html/2508.12594v2#A6.F9) (left) shows the impact of varying the number of residual layers in key/value projections and within the residual block on test accuracy for the elasticity benchmark dataset.
We suspect that deeper key/value encodings lead to more meaningful and focused attention, encoding structured inductive priors beneficial to downstream prediction.

#### ResMLP depth: feedforward block.

A second difference between the standard attention block and FLARE block is that we replace the feed-forward block in vanilla self-attention with a deep residual MLP.
Preliminary experiments using standard feed-forward blocks led to training instabilities and poor convergence.
In contrast, residual MLPs consistently enabled stable training and allowed us to increase model capacity.
[Figure 9](https://arxiv.org/html/2508.12594v2#A6.F9) (right) indicates that increasing the number of residual layers leads to slight improvements in accuracy.
Based on these results, we use three residual layers in both key/value projections and the residual block, as this provides a good trade-off between model capacity and computational cost in all subsequent experiments.

Figure: Figure 10: Effect of head dimension ($D$) on test accuracy. We design FLARE to work optimally for $D=4$–$8$.
Refer to caption: x9.png

#### Number of heads.

A key hypothesis behind our design is multiple simultaneous (head-wise parallel) low-rank projections collectively approximate a full attention map.
Thus, for a fixed total feature dimension ($C$), using a smaller head dimension ($D$) should improve performance by allowing more such parallel projections.
To test this, we evaluate FLARE on the elasticity dataset with varying head counts while keeping the total feature dimension $C=64$ and number of blocks $B=8$ fixed.
As shown in [Figure 10](https://arxiv.org/html/2508.12594v2#A6.F10), FLARE achieves the lowest test error when $D=4$ or $D=8$, in contrast to standard transformers which often use $D=16$–$32$.
Notably, because our head dimensions are smaller, we use a scaling factor of $1$ instead of the typical $\sqrt{D}$ employed in transformers (Vaswani et al., 2017), as the dot product magnitudes are sufficiently small without additional scaling.

## Appendix G Benchmark dataset of additive manufacturing simulations

Figure: Figure 11: We simulate the LPBF process on selected geometries from the Autodesk segmentation dataset (Lambourne et al., 2021) to generate a benchmark dataset for AM calculations. Several geometries are presented in this gallery. The color indicates $Z$ (vertical) displacement field.
Refer to caption: x10.png

### G.1 Introduction & Background

In metal additive manufacturing (AM), subtle variations in design geometry and process parameter selection may result in undesirable part artifacts or even costly build failures.
Numerical simulations of laser powder bed fusion (LPBF), a popular additive manufacturing process, can be used to predict build failures, but this may take several minutes to hours depending on the part size.

Although AM enables the fabrication of a wide variety of new designs, the final product must nonetheless comply with the constraints of the underlying manufacturing process.
Metal AM parts often have anisotropic and spatially varying material properties that depend on geometric features (such as overhang, or support structure) and fabrication process parameters (such as laser energy density, hatch spacing, layer thickness, and raster path).
As-built parts with thin features often suffer distortions when manufactured with the LPBF process, which uses a high-power laser to selectively melt and fuse metal powder, layer-by-layer, to create complex, high-precision parts.
Thermal stresses, termed *residual stresses* (RS), accumulate in LPBF-fabricated parts as a result of rapid thermal cycling due to laser exposure.
These stresses can be severe to the point of inducing localized plastic deformations or delamination.
As such, due to these *residual deformations*, the final shape of the part may deviate from the designed geometry.

We present a high-fidelity thermomechanical RS calculation data set on the Fusion 360 segmentation dataset, a publicly available dataset of complex 3D geometries (Lambourne et al., 2021).
Numerical solvers for simulating the LPBF build process perform expensive quasi-static thermo-mechanical equations,

$$ $\underbrace{\rho C_{p}\frac{\mathrm{d}T}{\mathrm{d}t}=\gradient\cdot k\Delta T(\bm{x},t)+Q(\bm{x},t),}_{\text{thermal transport}}\hskip 20.00003pt\underbrace{\gradient\cdot\bm{\sigma}=0{\color[rgb]{1,1,1}\frac{1}{2}}\hskip 10.00002pt\bm{\sigma}=\bm{C}\bm{\varepsilon}_{e},}_{\text{stress equilibrium}}$ (25) $$

layer-by-layer within a finite element framework
(Denlinger et al., 2014; Liang et al., 2019; Autodesk, 2025), and are integrated in commercial software products (Autodesk, 2025).
These calculations take several minutes to hours, making them prohibitively expensive for part design scenarios that can involve hundreds of evaluations.

Ferguson et al. (2025) introduced a dataset of LPBF simulations in which multiple finite element calculations were performed on a collection of 3D shapes.
That dataset, however, was restricted to relatively coarse meshes with approximately 3,500 points per mesh.
The present work extends this line of research by considering larger and more refined meshes, with up to 50,000 grid points.

Figure: Figure 12: Summary of LPBF dataset statistics.
Refer to caption: x11.png

### G.2 Dataset Generation

To generate a dataset of LPBF simulations, we employ Autodesk NetFabb (Autodesk, 2025), a commercially available software tool for numerically simulating RS and associated physics.
We begin with the Fusion 360 segmentation dataset (Lambourne et al., 2021), and scale each shape to lie within
$[-30,30]\times[-30,30]\times[0,60]~\text{mm}$
such that it rests atop the build plate at $z\in[-25,0]~\text{mm}$; the parts are not otherwise rotated or transformed.
The simulation is then carried out for the
Renishaw AM250 machine
and Ti-6Al-4V material system deposited with
$40~\mu\text{m}$ thickness.
Other parameters are left as their default nominal values and no support structures are added (Ferguson et al., 2025).

In AM, the material is deposited layer by layer, and ideally, a high-fidelity simulation would model each layer individually.
However, this can be computationally expensive, especially for builds with hundreds or thousands of layers.
Layer lumping simplifies this process by combining multiple physical layers into a single *lumped* computational layer.
In our calculations, NetFabb applies layer-lumping with a lumped layer thickness of $2.5~\text{mm}$.

NetFabb re-meshes the geometry to contain axis-aligned hexahedral elements before simulating the build process.
Then, NetFabb generates a thermal and a mechanical history for each part
corresponding to lumped layer deposition steps during the build.
We obtain the displacement, elastic strain, von Mises stress, and temperature fields evaluated at all nodal locations throughout the build.
NetFabb also provides field values after the part has cooled down and detached from the build plate.

### G.3 Benchmark task

To evaluate neural surrogate models, we target a field prediction task that is both central to our dataset and broadly relevant to the AM community: predicting residual vertical ($Z$) displacement.
In LPBF, each layer is fused by a laser and followed by a recoater blade that spreads powder uniformly across the build area (Reijonen et al., 2024).
Overhanging features may cause vertical displacements that interfere with the blade’s path, potentially leading to collisions. Predicting the $Z$-displacement field can therefore help identify risk of blade collision failures.
Rapid estimation of displacement prior to a build, thus, is highly desirable for design troubleshooting, as severe distortion can render a part unusable.
Moreover, accurate prediction of nodal displacements can help anticipate build failures (Ferguson et al., 2025).

Since full-scale LPBF simulations are computationally expensive—taking minutes to hours—a fast surrogate model offers a valuable alternative for accelerating AM design.
Accordingly, we train our models to predict the $Z$-displacement at every node at the final time step.

More formally,
the input to a neural surrogate model is the volumetric axis-aligned hexahedral mesh describing the geometry.
This includes the point-coordinates (array of size $N\times 3$ where $N$ is the number of points) and, optionally, mesh connectivity information.
The corresponding label is the $Z$ displacement value at each point (array of size $N\times 1$).

While this dataset focuses on a steady-state prediction problem, future iterations of this benchmark could involve learning dynamic surrogate models that track the time-history of stress and deformation fields during the build process.

**Table 6: Summary statistics of our proposed LPBF dataset.**
|  | #Points | #Edges | Avg./ max aspect ratio | Max height (mm) | Max Displacement |
| --- | --- | --- | --- | --- | --- |
| Mean | 20,972 | 114,140 | 1.6421 / 1.6421 | 29.429 | 0.29526 |
| Std. | 12,476 | 68,308 | 0.43794 / 0.43794 | 23.246 | 0.21064 |
| Min | 736 | 2,860 | 1.0667 / 1.0667 | 0.60000 | 0.00048500 |
| 25% | 10,229 | 56,208 | 1.2800 / 1.2800 | 7.8000 | 0.13827 |
| 50% | 19,743 | 107,680 | 1.4733 / 1.4733 | 21.600 | 0.27075 |
| 75% | 30,503 | 166,250 | 1.9072 / 1.9072 | 60.000 | 0.41962 |
| Max | 47,542 | 249,930 | 2.9932 / 2.9933 | 60.000 | 0.99777 |

### G.4 Data filtering

The dataset contains a wide-varying range of shapes, making the data set general enough to train a strong data-driven field prediction model.
Out of $\sim$27,000 shapes, 19,732 were successfully simulated.
We analyze the first 3,500 successful simulations and filter them according to several statistics to design a balanced training and test set.
For example, we limit the learning problem to meshes with up to 50,000 points and up to 300,000 edges.
This is done to reduce memory usage which becomes a bottleneck when training on small GPUs.
We also filter meshes that have high aspect ratio elements as the FEM calculation could be unreliable on highly distorted geometries.
Our dataset processing code is available in the GitHub repository associated with this paper.
The statistics for the filtered data set are presented in [Table 6](https://arxiv.org/html/2508.12594v2#A7.T6) along with histograms in [Figure 12](https://arxiv.org/html/2508.12594v2#A7.F12).
A gallery of successful simulations is presented in [Figure 11](https://arxiv.org/html/2508.12594v2#A7.F11).

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1、The “Criteria” column in Table 1 lists several evaluation standards, but their specific meanings are not clearly explained. Please provide detailed descriptions for each criterion.
soundness: 2
strengths: 1、The idea of low-rank self-attention is intersting. This wrok provides a good explanation on architectures design for  the neural operator domain, especially classical works and methods based on Transformer architectures.
2、A new dataset, LPBF, is proposed, making a valuable contribution to the advancement of research in this field.
summary: This paper introduces FLARE, a novel low-rank self-attention mechanism designed for efficient neural surrogate modeling of partial differential equations (PDEs) on point clouds or unstructured meshes. By leveraging known properties of rank-deficient matrices, the authors employs a linear complexity attention mechanism with a learnable latent bottleneck to further improve accuracy and efficiency on PDE surrogate modeling tasks. Furthermore, FLARE enables end-to-end training on unstructured meshes with one million points without distributed computing or memory offloading and proposes a new a 3D field-prediction benchmark dataset named LPBF.
weaknesses: 1、The experimental datasets do not include one-dimensional or time-dependent PDE problems. It is recommended to add experiments demonstrating the model’s generality, for example by including shallow-water and reaction-diffusion equations from PDEBench as time-dependent cases.(Takamoto, Makoto, et al. "Pdebench: An extensive benchmark for scientific machine learning." Advances in Neural Information Processing Systems 35 (2022): 1596-1611.)
2、Since Mamba also achieves linear computational complexity, please compare and evaluate the proposed model against Mamba-based architectures, such as MambaNO and LaMO, highlighting their respective advantages and limitations.(Zheng, Jianwei, et al. "Alias-free mamba neural operator." Advances in Neural Information Processing Systems 37 (2024): 52962-52995.)(Tiwari, Karn, Niladri Dutta, and N. M. Krishnan. "Latent Mamba Operator for Partial Differential Equations." arXiv preprint arXiv:2505.19105 (2025).)

---

confidence: 4
contribution: 3
initial_rating: 4
presentation: 3
questions: Is there any design in the proposed method specifically for neural PDE scenarios? What if the proposed method applies to domains other than surrogate PDE solvers, e.g., LLM training?
soundness: 3
strengths: The proposed method achieves sota PDE surrogate performance, while is also capable of handling geometries with a million points.
summary: In this paper, a linear self-attention mechanism is proposed for applying to large unstructured meshes. It is achieved by projecting the input sequence to a much lower-dimentional latent sequence. It enables training with 1 million points. The proposed model achieves superior accuracy compared to neural PDE models, including Vanilla Transformer, PerceiverIO, GNOT, LNO, and Transolver, on various datasets, including a newly constructed additive manufacturing simulation dataset.
weaknesses: I think it is mandatory that in the experiments, this work also compares with other efficient attention methods proposed in general domains.

---

confidence: 4
contribution: 1
initial_rating: 4
presentation: 3
questions: Questions are formed as part of the weakness. Major concern is around the novelty of the work and the claims. I kindly request authors to focus on the novelty, and the claims (claims are specifically mentioned in the weakness) as they are clarifying the weaknesses.
soundness: 2
strengths: 1) Authors demonstrate a transformer based surrogate with global communication training directly on ~1M-point meshes on a single H100 80GB GPU using off-the-shelf fused attention kernels, and provide scaling curves. Scalability on experiments is something that has not gotten enough attention in the literature in this field, which makes these experiments the strong point.

2) They retrain all baselines (Perceiver IO, Transolver, LNO, etc.) under standardized splits, resolutions, hyperparameters, and model sizes. This makes the reported accuracy numbers more trustworthy.
summary: Self-attention is quadratic in sequence length, which makes it hard to apply to very large unstructured meshes. Authoers propose FLARE, a linear-complexity alternative. Each attention head routes global communication for N tokens through a learned latent bottleneck of M < N tokens. This induces a low-rank attention map (rank < M) and reduces cost to O(N·M). By choosing M, users can trade accuracy for efficiency.
weaknesses: 1) FLARE’s low-rank attention via a latent bottleneck is directly linked to the attentions of Perceiver/Perceiver IO (iterative cross-attention into a fixed-size latent that scales linearly with input size) and Linformer (self-attention is low-rank; approximate it to get O(N) complexity). There are various methods along the same direction which makes the novelty of the concept unclear. If the novelty lies in O(N^2) -> O(N.M) then the concept is not new. If the mechanism by which this procedure happens (encode-decode) please see issues 2 and 3. 

2) Authors do not run any meaningful ablations on the encode-decode mechanism which makes it unclear if this per-block encode decode is actually contributing to the results. The only ablations are minor (MLP depth, symmetry tie/untie), so the central architectural claim is not causally demonstrated. If the effect of encode-decode is not clear then do we fall back to the Preciever, or linformer style attention? 

3) The following claim is vert strong " However, the latent bottleneck in PerceiverIO can limit accuracy as the model may discard fine-grained features if the number of latent tokens is too low. ". Preciever style model keep the original tokens [M x ...] and conduct cross attention to these at the beginning of each block. How does this discard information?

4) Authors claim giving each head its own latent slice is a key idea for expressivity, but never ablate a shared-latent variant. The evidence is descriptive spectra plots, not causal. Ablations are missing to show what would have been the effect of a shared latent token. This is especially important since it's not something that has been explored or adopted widely with attention mechanism in general. 

5) The idea of weight sharing across cross attentions is also from Preciever paper so it's not fair to have the claim that PrevieverIO did not include the weight sharing (Table 1).

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 3
questions: - In Figure 1, the peak memory usage of vanilla attention is lower than physical attention and flare. Is this reasonable? Did you use any optimization trick? If yes, all baselines should use same optimization tricks.

- In Table 2, baseline models have different parameter counts. Did you compare the efficiency of different models with comparable parameter counts? This is important for fair comparison, especially for efficiency.

- In Figure 4, with different $M$, the peak GPU memory usage is similar. Does this imply that the memory bottleneck is not the attention operation?

- There is only spectral analysis of Flare in Section 3.3 and Appendix B. Without the same analysis of other baselines like Transolver or LNO, it is hard to get the conclusion “FLARE benefits from head-wise diversity”.
soundness: 2
strengths: - This paper improves the accuracy and efficiency of Transformer-based PDE surrogates.

- The proposed FLARE framework is clearly described and easily implemented.

- This paper provides a new dataset that is meaningful for the community.
summary: This paper introduces FLARE, a linear-complexity attention mechanism designed to overcome the quadratic cost of traditional attention on large unstructured meshes. FLARE routes attention through learnable latent tokens, effectively forming a low-rank attention matrix that balances efficiency and accuracy. Experiments show that FLARE outperforms advanced PDE surrogate models across multiple benchmarks on accuracy and efficiency. The authors also release a manufacturing dataset to further research in scientific machine learning.
weaknesses: - The contribution of the tech is incremental and limited. This work neither tackles a new important scientific problem nor proposes a fundamentally new solution. Regarding to the tech proposed, the core idea of latent tokens has already been explored in prior works such as PerceiverIO, Transolver, LNO and etc. The main contribution is a slight modification of the attention formulation but it is hard to see the advantage of this modification comparing to other models like Transolver from the technical analysis. For example, the attention is operated on sequences with length $M$ in Transolver and on sequences with length $N$ and $M$ in Flare. Why is Flare more GPU efficient if $N \gg M$ ? Are the improvements mainly due to the use of the SDPA kernel? More technical analysis and comparison should be included. Otherwise, the paper reads more like an engineering optimization report rather than a scientific research paper.

- Following the above item, in Table 1 and Line 74, the authors claim that Transolver cannot utilize the SDPA kernel. This is inaccurate. From the official repo of Transolver++, the slice self-attention of Transolver can be naturally implemented by SDPA kernel. This operation is the same in Transolver++ and Transolver, not involving formulation optimization. So the comparison in Table 1 is misleading and not convincing.

- The dataset is also listed as a main contribution. But there is very limited content about the proposed dataset in the main context. The writing needs better arrangement to justify its value and meaning.

- Lack of visualizations. Given that this is a work focus on solving scientific problems efficiently, visualizations of predicted vs. ground-truth physical fields (e.g., stress or flow distributions) are essential. If the simulation states are not physical meaningful, the improvements in numerical accuracy or computational efficiency are of little significance.


Decision: Reject