FLARE + High order attention ideas

Vedant Puri
https://vpuri3.github.io/

Data-Driven Modeling

Scalable neural surrogates for PDEs and beyond!

Surrogate models learn PDE solution operator from data

\mathcal{L}(\boldsymbol{x}, t, \boldsymbol{u}; \boldsymbol{\mu}) = 0

\mathcal{G}: \boldsymbol{\mu} \mapsto \boldsymbol{u}

Training

\mathcal{G}_\theta \approx \mathcal{G}

Inference

\mathcal{G}_\theta

Large training cost is amortized over several evaluations

Model learns to predict \(\boldsymbol{u}\) over a distribution of \(\boldsymbol{\mu}\)

Transformers [1] are state-of-the-art surrogate models

Message-passing on a dynamic all-to-all graph.

[1] Vaswani et al. — “Attention Is All You Need”, NeurIPS 2017

Quadratic (\(\mathcal{O}(N^2)\)) cost limits scalability

Quadratic \((\mathcal{O}(N^2))\) cost of attention limits scalability

Over \(20~\text{s}\) per gradient step on a mesh of 1m poins!

Goal: enable transformer models on large meshes.

[1] Vaswani et al. — “Attention Is All You Need”, NeurIPS 2017

\([1]\)

What are the limitations on communication patterns?

\Delta u = f

\begin{bmatrix} &&&\\ &&&\\ &&& \end{bmatrix} \cdot \begin{bmatrix} \\ \underline{u}\\ \\ \end{bmatrix} = \begin{bmatrix} \\ \underline{f}\\ \\ \end{bmatrix}

\begin{bmatrix} \\ \underline{u}\\ \\ \end{bmatrix} = \begin{bmatrix} &&&&\\ &&&&\\ &&&&\\ \end{bmatrix} \cdot \begin{bmatrix} \\ \underline{f}\\ \\ \end{bmatrix}

Solution operator requires global communication.

Forward operator is implemented with sparse, structured communication.

Need principled strategy for reducing communication cost.

[1] Math StackExchange — “Interpolation in Finite Element Method”

u = \int_\Omega G(x-y)dy

[1]

Are \(N \times N\) messages really necessary?

Smoothness implies redundancy in communication.

Are \(N \times N\) messages really necessary?

Smoothness implies redundancy in communication.

Method: club matching points to one cluster and communicate together.

FLARE: Fast Low-rank Attention Routing Engine

Encoding: introduce \(M\) latent clusters to pool messages from matching tokens

\(M\) learned queries

Decoding: map pooled messages to matching output tokens

FLARE: Fast Low-rank Attention Routing Engine

\(\mathcal{O}(2MN) \ll \mathcal{O}(N^2)\)

\(\text{rank}(W_\text{encode}\cdot W_\text{decode}) \leq M\)

\(>200\times\) speedup

\(\text{(} M \text{ tokens)}\)

\(\text{Latent}\)

[1] Vaswani et al. — “Attention Is All You Need”, NeurIPS 2017

\([1]\)

PDE surrogate benchmark problems

Relative \(L_2\) error \( (\times 10^{-3})\) (lower is better)

Pipe

Darcy

Elasticity

LPBF

DrivAerML

[1] Vaswani et al. — “Attention Is All You Need”, NeurIPS 2017

[2] Jaegle et al. — "PercieverIO: A General Architecture for Structured Inputs & Outputs", ICLR 2022

[3] Hao et al., — "GNOT: A General Neural Operator Transformer for Operator Learning", PMLR 2023

[4] Wang et al. —"Latent Neural Operator", NeurIPS 2024

[5] We et al. — "Transolver: A Fast Transformer Solver for PDEs on General Geometries", ICML 2024

Elasticity benchmark problem

Pipe flow, Darcy flow benchmark problems

Laser powder bed fusion benchmark problem

FLARE learns surrogate on a million-point mesh!

Largest experiment on a single GPU!

[1]

[1] Ashton et al. — “DrivAerML: High-Fidelity CFD Dataset for Road-Car Aerodynamics” (arXiv:2408.11969, 2024)

FLARE generalizes beyond PDE tasks

Pathfinder

\texttt{INPUT:\, [MAX 4 3 [MIN 2 3 ] 1 0 [MEDIAN 1 5 8 9, 2]] \,OUTPUT: 5}

Listops

\texttt{INPUT:\, [MAX 7 [MEDIAN 1 2 3 ] [MAX 9 2 2] [MIN 2 8]] \,OUTPUT: 9}

Image classification

Text sentiment analysis

[7]

[8]

[1]

[5]Choromanski et al. — "Rethinking Attention with Performers", ICLR 2021

[6] Tay, Y. et al. — “Long Range Arena: A Benchmark for Efficient Transformers” (arXiv 2020)

[7] Centric Consulting — “Sentiment Analysis: Way Beyond Polarity” (blog)
[8] Krizhevsky — CIFAR dataset homepage

[6]

Accuracy \((\%)\) (higher is better)

[1] Vaswani et al. — “Attention Is All You Need”, NeurIPS 2017

[2] Katharopoulos et al. — "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention", ICML 2020

[3] Wang et al. — "Linformer: Self-attention with linear complexity", arXiv:2006.04768 2020

[4] Qin et al. — "The devil in linear transformer", arXiv:2210.10340 2022

Key idea: Modulate token-mixing with conditioning tokens

\begin{bmatrix} \mathbf{x}\\ \end{bmatrix}

Cross FLARE

\times B

\begin{bmatrix} \mathbf{u}_{t+1} \end{bmatrix}

\begin{bmatrix} t & \mathbf{u}_{t-k:t} \end{bmatrix}

Conditioning mechanism for FLARE

We propose to handle token mixing and conditioning in one unified block

\(\mathcal{O}(2MN + MC) \) complexity

Notes on high-order attention

Linear attention replaces \(N\times N\) connections with fixed-size state

Y = \text{softmax} \left(Q \cdot K^T\right) \cdot V

y_i = \frac{\sum_{j}\exp(q_i^T \cdot k_j) \cdot v_j}{\sum_{\ell} \exp(q_i^T k_\ell)}

Q = X \cdot W^q,\hspace{1em} K = X \cdot W^k,\hspace{1em} V = X \cdot W^v \in \mathbb{R}^{N\times D}

Full self-attention

Linear attention

Y = \text{row\,norm}({Q} \cdot {K}^T) \cdot V

y_i = \frac{\sum_{j}({q}_i^T \cdot {k}_j) v_j}{\sum_{\ell} {q}_i^T {k}_\ell}

=\frac{{q}_i^T \sum_{j} {k}_j v_j^T}{{q}_i^T\sum_\ell {k}_\ell}

Y = Q \cdot (K^T \cdot V)

Token mixing mediated through \(S\in\mathbb{R}^{D\times D}\)

Q = \phi(Q), \hspace{1em} K = \phi(K)

Proposal 1: Multilinear kernelized attention from 2-simplicial attn

Self attention is built on the bilinear form

\langle x_1, x_2 \rangle = \sum_{d=1}^D x_1[d] x_2[d]

2-simplicial attention generalizes it to trilinear forms

\langle x_1, x_2, x_3 \rangle = \sum_{d=1}^D x_1[d] x_2[d] x_3[d]

Then 2-simplicial attention proceeds as follows: consider \(Q, K^1, K^2, V^1, V^2 \in \mathbb{R}^{N \times D}\).

S_{ijk} = \sum_{d'=1}^D Q_{id'} K^1_{jd'} K^2_{kd'} \hspace{2em} A = \text{softmax}\left( S, \text{dims=(-1,-2)}\right) \hspace{2em} Y_{id} = \sum_{j,k=1}^N A_{ijk} V^1_{jd} V^2_{kd}.

We review linear attention and note that removing softmax allows for exploiting associativity:

Y_{id} = \sum_{d'=1}^{D} \sum_{j=1}^N Q_{id'} K_{jd'} V_{jd} = \sum_{d'=1}^{D} Q_{id'} \left(\sum_{j=1}^N K_{jd'} V_{jd} \right)

We linearize 2-simplicial attention

A_{ijk} = \sum_{d'=1}^D Q_{id'} K^1_{jd'} K^2_{kd'} \hspace{2em} Y_{id} = \sum_{j,k=1}^N A_{ijk} V^1_{jd} V^2_{kd}.

This implies:

S_1 = K_1^T \cdot V_1,\quad S_2 = K_2^T \cdot V_2 \in \mathbb{R}^{D\times D}

Y = Q \cdot (S_1 \odot S_2)

We can generalize this to arbitrary order

Y = Q \cdot (S_1 \odot \cdots \odot S_L)

S_1 = K_1^T \cdot V_1,\quad\ldots,\quad S_L = K_L^T \cdot V_L \in \mathbb{R}^{D\times D}

Proposal 2: Higher order memory

Instead of the state being \(D \times D\), have the state be \(D \times D \times D\). We call it Triple Attention

S_{ijk} = \sum_{n=1}^{N} K^{1}_{ni} V_{nj} K^{2}_{nk} \hspace{2em} Y_{nj} = \sum_{i,k=1}^{D} Q^{1}_{ni} S_{ijk} Q^{2}_{nk}

Similarly, we propose Quad-Attention with \(D \times D \times D \times D\) state.

S_{ijk\ell} = \sum_{n=1}^{N} K^{1}_{ni} V_{nj} K^{2}_{nk} K^{3}{n\ell} \hspace{2em} Y_{nj} = \sum_{i,k,\ell=1}^{D} Q^{1}_{ni} S_{ijk} Q^{2}_{nk} Q^{3}_{n\ell}

Triple Attention or Multi-linear attention

FEATURES

Considers N-tuples of tokens at a time.
More expressive than standard attention, linear attention
Easily parallelizable across multiple GPUs
Kernel-based interpretation
As efficient and accurate as FLARE (SOTA)

DEMONSTRATIONS

Encoder transformer
- PDE Surrogate modeling, Long-range arena
Decoder transformer
- Next-token prediction/ language modeling

Triple Attention Scaling Study

Notes

Hyper-graph attention

Triple/third order attention is a starting point for exploring high-order attention but it really shouldn't be the end point. Real world tasks may need arbitrary order compositional reasoning.
However, the methods we have explored thus far consider fixed-length N-Tuples. This is where the idea of looping becomes interesting: can we design a method that builds up high-order reasoning with low-order primitives?

Relation to linearized attention

In principle, the state \(S \in \mathbb{R}^{D \times D}\) in linear attention can store information from several tokens. As such, does that negate the need for high-order attention methods? I don't know.

Proposed Work

Extend FLARE to transient problems

Motivation: preempt build failures in metal additive manufacturing

Laser Powder Bed Fusion (LPBF)

Dataset of 20k LPBF calculations

Goal: develop fast surrogate model to predict warpage during build

\rho C_p \frac{dT}{dt} = \nabla \cdot k\Delta T(\mathbf{x}, t) + Q(\mathbf{x}, t)

\nabla \cdot \phi = 0, \,\, \sigma = C\varepsilon_e

Governing equations

End results could be deployed as a valuable design tool for metal AM.

[1]

[2]

[1] Nature Scientific Data — High-resolution dataset (2025)
[2] TechXplore — “Synergetic optimization reduces residual warpage in LPBF” (2022)

Proposed aims

Aim 1: Advance FLARE for enhanced surrogate modeling

Aim 2: Develop decoder version of FLARE

AIM 1(a): rank-adaptive FLARE

AIM 1(b): conditioning mechanism for FLARE

AIM 1: advance FLARE for enhanced surrogate learning

AIM 1(a): rank-adaptive FLARE for faster training

Complexity scales with latents (\(M\)): \(\mathcal{O}(2MN)\)

Accuracy increases with \(M\)

Method: progressively increase latents (\(M\)) through training.

Challenge: Minimize loss spikes, training instabilities.

Aim 1(b) background on conditioning in transformers

\left( t,\, \mathbf{x},\, \mathbf{u}_{t-k},\, \cdots, \mathbf{u}_{t} \right) \mapsto \mathbf{u}_{t+1}

Token mixing [1] (\(\mathcal{O}(N^2)\))

Conditioning [1] (\(\mathcal{O}(N\cdot C)\))

\begin{bmatrix} {t} \\ \mathbf{x}\\ \mathbf{u}_{t-k:t} \end{bmatrix}

Token mixing

\times B

\begin{bmatrix} \mathbf{u}_{t+1} \end{bmatrix}

\begin{bmatrix} \mathbf{x}\\ \end{bmatrix}

Token mixing

\begin{bmatrix} \mathbf{u}_{t+1} \end{bmatrix}

Conditioning

\begin{bmatrix} t & \mathbf{u}_{t-k:t} \end{bmatrix}

\times B

[1] Vaswani et al. — “Attention Is All You Need”, NeurIPS 2017

Key idea: Modulate token-mixing with conditioning tokens

\begin{bmatrix} \mathbf{x}\\ \end{bmatrix}

Cross FLARE

\times B

\begin{bmatrix} \mathbf{u}_{t+1} \end{bmatrix}

\begin{bmatrix} t & \mathbf{u}_{t-k:t} \end{bmatrix}

Aim 1(b) conditioning mechanism for FLARE

We propose to handle token mixing and conditioning in one unified block

\(\mathcal{O}(2MN + MC) \) complexity

Aim 2: background on next-token prediction transformers [1]

y_t = \frac{\sum_{\tau = 1}^t\exp\left(q_t \cdot k_\tau \right) v_\tau}{\sum_{\tau = 1}^t \exp \left(q_t \cdot k_\tau \right)}

All previous key/value \(\{k_\tau, v_\tau \}_{\tau \leq t}\) must be cached on the GPU.

Major memory and latency bottleneck!

[1] Vaswani et al. — “Attention Is All You Need”, NeurIPS 2017

Training algorithm (causal masking)

Inference algorithm (recurrence relation)

Dot-products need to be recomputed for every \(q_t\).

\(\mathcal{O}(N^2)\) complexity.

Aim 2: Develop decoder version of FLARE

Linear time auto-regressive attention.

Fixed memory footprint (only store \(\mathcal{O}(M)\) cache).

Flexible latent capacity.

Advantages

Required components

Fused GPU kernels for training and inference.

Bespoke training algorithm for causal FLARE.

Extensive benchmarking and evaluation.

Z_t = \text{online\_softmax}(Z_{t-1}, k_t, v_t)\\ y_t = \text{softmax}(Q^T \cdot k_t)^T\cdot Z_t

Inference algorithm (recurrence rule)

Next-token prediction with FLARE

Efficient transformers models

Adaptive Layer Norm in Diffusion Transformer allows for token mixing + time-conditioning in one go

This is only possible with a single token as conditioning vector, and won't work when you want to condition on a sequence.

Linear transformers only store state \( S \in \mathbb{R}^{D \times D} \) but their performance is not on par with softmax attention

Linear transformers replace the softmax kernel with a feature map \(\phi(\cdot)\) such that

This factorization allows causal attention to be computed recurrently:

\mathrm{softmax}(QK^\top) V \approx \phi(Q)\,\big(\phi(K)^\top V\big)

S_t = S_{t-1} + \phi(k_t)^\top v_t, \qquad \mathbf{y}_t = \phi(q_t)\, S_t,

Chunkwise training for linear transformers

https://manifestai.com/articles/linear-transformers-are-faster/

Premise: strong encoder model --> strong LLM

Next-token prediction model
- FLARE, Triple Attention
- Write CUDA kernels --> get scaling plots
- Test on language tasks
Extend FLARE
- Allow model to increase/ decrease \(M\) during training
- Create efficient conditioning mechanism (time-series PDE problems)
FOCUS ON NEW CONTRIBUTIONS and how we can differentiate ourselves from SOTA
explain novelty compared to SOTA