2025-11-17T15:49:13.397134

FLARE: Fast Low-rank Attention Routing Engine

Puri, Joglekar, Ferguson et al.

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$ tokens by projecting the input sequence onto a fixed length latent sequence of $M \ll N$ tokens using learnable query tokens. By routing attention through a bottleneck sequence, FLARE learns a low-rank form of attention that can be applied at $O(NM)$ cost. 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 release a new additive manufacturing dataset to spur further research. Our code is available at https://github.com/vpuri3/FLARE.py.

academic

FLARE: Fast Low-rank Attention Routing Engine

基本信息

论文ID: 2508.12594
标题: FLARE: Fast Low-rank Attention Routing Engine
作者: Vedant Puri, Aditya Joglekar, Kevin Ferguson, Yu-hsuan Chen, Yongjie Jessica Zhang, Levent Burak Kara (Carnegie Mellon University)
分类: cs.LG (Machine Learning)
发表时间: 2025年10月15日 (arXiv v2)
论文链接: https://arxiv.org/abs/2508.12594

摘要

传统自注意力机制的二次复杂度限制了其在大规模非结构化网格上的适用性和可扩展性。本文提出了快速低秩注意力路由引擎(FLARE)，这是一种线性复杂度的自注意力机制，通过固定长度的潜在序列路由注意力。每个注意力头通过使用可学习的查询令牌将输入序列投影到长度为M≪N的固定长度潜在序列上，实现N个令牌间的全局通信。通过瓶颈序列路由注意力，FLARE学习低秩形式的注意力，可以O(NM)的代价应用。FLARE不仅能扩展到前所未有的问题规模，而且在多个基准测试中相比最先进的神经PDE代理模型提供了更优的准确性。

研究背景与动机

问题背景

核心问题：传统Transformer的自注意力机制具有O(N²)的时间和内存复杂度，这严重限制了其在大规模非结构化网格(如物理仿真中的点云和网格)上的应用。
应用重要性：在偏微分方程(PDE)代理建模中，每个3D点云中的点被视为一个令牌，包含几何和物理量(如坐标、法向量、材料属性)等特征。高保真物理系统仿真成本过高，机器学习代理模型提供了快速近似的替代方案。
现有方法局限性：
- PerceiverIO：仅执行单次编码和解码，潜在瓶颈可能限制准确性
- Transolver：跨头共享投影权重，无法利用现有GPU内核进行缩放点积注意力
- LNO：仅应用单次投影，缺乏深层模型能力
研究动机：开发一种能够保持全局通信能力但具有线性复杂度的注意力机制，使Transformer能够处理百万级点的几何体。

核心贡献

线性复杂度令牌混合：提出FLARE自注意力机制，通过低秩投影和重构替代完整自注意力，实现线性复杂度。
卓越准确性：在多个PDE基准测试中，FLARE以更少参数和更低计算复杂度实现了优于领先神经代理模型的预测准确性。
前所未有的可扩展性：FLARE完全基于标准融合注意力原语构建，确保高GPU利用率，支持百万点非结构化网格的端到端训练。
新基准数据集：发布大规模高分辨率金属增材制造数据集，用于残余位移预测研究。

方法详解

任务定义

给定输入序列X ∈ R^(N×C)，其中N为令牌数，C为特征维度，FLARE旨在学习一个线性复杂度的注意力机制，实现高效的全局令牌间通信。

模型架构

FLARE核心机制

FLARE引入M≪N个可学习的潜在令牌作为信息交换的瓶颈，包含两个阶段：

编码阶段：输入序列通过交叉注意力投影到潜在令牌
```
Z_h = SDPA(Q_h, K_h, V_h, s=1)
```
其中Q_h ∈ R^(M×D)为可学习查询矩阵，K_h, V_h ∈ R^(N×D)
解码阶段：潜在令牌投影回输入序列
```
Y_h = SDPA(K_h, Q_h, Z_h, s=1)
```

低秩通信矩阵

整个过程等价于：

Y_h = (W_decode,h · W_encode,h) · V_h

其中：

W_encode,h = softmax(Q_h · K_h^T) ∈ R^(M×N)
W_decode,h = softmax(K_h · Q_h^T) ∈ R^(N×M)
W_h = W_decode,h · W_encode,h ∈ R^(N×N)为秩最多为M的全局通信矩阵

FLARE块结构

X = X + FLARE(LayerNorm(X))
X = X + ResMLP(LayerNorm(X))

技术创新点

头间独立投影：与Transolver共享投影权重不同，FLARE为每个头分配不同的潜在令牌切片，使每个头学习独立的注意力关系。
深度残差MLP：使用深度残差网络进行键/值投影，相比简单线性层能学习更高阶特征交互。
对称编解码设计：编码和解码操作的对称性促进稳定的信息流。
兼容融合内核：完全基于标准SDPA操作，可利用Flash Attention等优化算法。

实验设置

数据集

论文评估了6个基准数据集和1个新提出的数据集：

数据集	维度	网格类型	点数	输入/输出特征	训练/测试样本
Elasticity	2D	非结构化	972	2/1	1000/200
Darcy	2D	结构化	7,225	2/1	1000/200
Airfoil	2D	结构化	11,271	2/1	1000/200
Pipe	2D	结构化	16,641	2/1	1000/200
DrivAerML-40k	3D	非结构化	40,000	3/1	387/97
LPBF	3D	非结构化	1,000-50,000	3/1	1100/290

评价指标

主要使用相对L2误差：

Relative L2 = ||û - u||₂ / ||u||₂

对比方法

通用注意力模型：Vanilla Transformer, PerceiverIO
基于注意力的PDE代理：Transolver, LNO
神经算子：GNOT

实现细节

优化器：AdamW (β₁=0.9, β₂=0.999)
学习率调度：OneCycleLR，峰值学习率10⁻³
训练轮数：2D问题500轮，LPBF 250轮
批量大小：2D问题为2，3D问题为1

实验结果

主要结果

FLARE在所有基准测试中均取得最优或次优结果：

模型	Elasticity	Darcy	Airfoil	Pipe	DrivAerML-40k	LPBF
Vanilla Transformer	5.37	4.38	6.28	∼	∼	∼
PerceiverIO	23.4	21.5	162	7.14	760	56.3
GNOT	13.3	16.9	103	5.89	115	24.3
LNO	9.25	7.64	17.8	8.10	146	24.7
Transolver w/o conv	6.40	18.6	8.24	4.87	70.5	20.4
Transolver with conv	\	5.94	5.50	3.90	\	\
FLARE (ours)	3.38	5.10	4.28	2.85	60.8	18.5

注：数值为相对L2误差(×10⁻³)

百万点几何体实验

FLARE成功在单个H100 GPU上训练百万点DrivAerML数据集，这是首个在不使用内存卸载或分布式计算的情况下处理百万点的基于注意力的神经代理模型。

消融实验

块数(B)和潜在令牌数(M)的影响：
- 增加块数持续降低相对误差
- 增加M通常改善性能，但趋势不严格单调
- 不同问题对秩的需求不同
时间和内存复杂度：
- FLARE比vanilla attention快200倍以上
- 内存使用略高于vanilla attention但远低于Physics Attention