FLARE++：低秩注意力与动态注意力路由

# FLARE++：具有动态注意力路由的低秩注意力
来源：https://arxiv.org/html/2608.11519
Yongjie Jessica Zhang & Levent Burak Kara
单位：Department of Mechanical Engineering
单位：Carnegie Mellon University
单位：Pittsburgh, Pennsylvania, USA

###### 摘要

完全自注意力\(36 (https://arxiv.org/html/2608.11519#bib.bib19)\)是不规则域上PDE代理模型的一种强力token混合器，但其二次成本限制了其在高分辨率问题上的应用。高效的潜在注意力模型，如快速低秩注意力路由引擎（FLARE）\(30 (https://arxiv.org/html/2608.11519#bib.bib42)\)，通过将所有\(N\)个token路由到\(M\ll N\)个可学习的潜在查询来避免这一成本，但这些查询是参数：一旦训练完成，相同的可学习查询模板就服务于每个输入。我们通过**FLARE++**（一种具有动态token路由的低秩注意力架构）消除了这一限制。FLARE++复用FLARE自身的编码器来构建其路由查询：可学习的潜在种子驱动一次额外的编码调用，将\(N\)个输入token聚合成\(M\)个输入条件化的查询，随后这些查询决定相同的token如何被压缩和重新分配。这保留了FLARE的显式低秩分解和线性\(\mathcal{O}(NM)\)复杂度，并且仅使用标准缩放点积注意力（SDPA）调用就表达了完整的路由操作。我们还提供了一种多GPU上下文并行实现，将输入token跨设备切分，而无需在任何单个设备上聚合完整的token序列。FLARE++在一组标准PDE代理基准上具有竞争力，平均比固定查询的FLARE提升\(24\%\)，并在Long Range Arena上平均准确率提升\(2.3\)个点。

## 1 引言

图1：FLARE++混合器。可学习的种子\(\widetilde{Q}_{h}\)合成\(M\)个输入条件化的路由查询\(Q_{h}(X)\)（左），这些查询将\(N\)个token聚合成\(M\)个潜在值并重新分配它们（右）。每个分数矩阵内的箭头标记其softmax归一化的轴：\(W_{\mathrm{enc},h}\)在\(N\)个token上，\(W_{\mathrm{dec},h}\)在\(M\)条路由上。下方：所有\(H\)个头的混合器及其成本，以及其两次路由调用组合而成的秩\(M\)算子。SDPA是融合的，因此\(S_{h}\)永远不会到达内存，空间保持\(\mathcal{O}(NC)\)。

自注意力已成为PDE代理模型的主导架构，因为它允许每个离散点与所有其他点通信。在我们考虑的每个代理模型中，每个离散点都被嵌入为自身的token，因此包含\(N\)个点的网格就是一个包含\(N\)个token的序列；我们在全文使用后一个术语，因为它是混合器作用的实体。这种通用性的代价是一张\(N\times N\)的通信矩阵，因此在token数量上需要\(\mathcal{O}(N^{2})\)的工作量\(36 (https://arxiv.org/html/2608.11519#bib.bib19)\)，这在工程问题实际产生的网格上是无法实现的。一系列高效的token混合器改为通过\(M\ll N\)个潜在token进行通信。我们基于FLARE\(30 (https://arxiv.org/html/2608.11519#bib.bib42)\)构建，它*仅*将潜在token用于路由：对于每个注意力头，\(M\)个可学习的潜在查询通过一次缩放点积注意力（SDPA）调用将\(N\)个输入token聚合成\(M\)个潜在token，第二次调用则反转方向，将这些值分发回\(N\)个输入位置。这两次调用构成一种编码-解码分解，对于每个固定输入，诱导出一个秩至多为\(M\)的显式输入到输入注意力矩阵，并且可以完全使用融合的SDPA实现。潜在工作空间模型（如Transolver\(39 (https://arxiv.org/html/2608.11519#bib.bib9)\)）则使用自身的自注意力阶段来处理潜在序列。我们基于FLARE，因为它将路由本身隔离为token混合操作，这使我们能够在不引入单独潜在处理阶段的情况下使路由依赖于输入。这类混合器，包括PerceiverIO\(13 (https://arxiv.org/html/2608.11519#bib.bib8)\)、LNO\(38 (https://arxiv.org/html/2608.11519#bib.bib10)\)、Transolver系列\(39 (https://arxiv.org/html/2608.11519#bib.bib9);23 (https://arxiv.org/html/2608.11519#bib.bib36);44 (https://arxiv.org/html/2608.11519#bib.bib43)\)以及FLARE，都分为两个对象：一个*压缩模板*，即决定信息如何被收集和重新分配的\(M\)槽结构，以及该模板所应用的场。在所有方法中，压缩表示都依赖于场，这是显然的，并非它们的区别所在。它们的区别在于模板的来源：在FLARE中，定义模板的可学习查询是参数，而其他模型使用可学习的逐点映射或固定的场依赖规则来构建其模板。在FLARE中，训练只学习每层-头\(M\)个查询模板一次，之后这相同的模板服务于每个几何形状和每个边界条件：路由权重会响应当前键，但定义模板的查询不会。这是算子中唯一从不看到其所压缩的场的部分，也正是它决定了秩\(M\)瓶颈如何被使用。

我们提出**FLARE++**，它从输入token构建压缩模板，而不是固定模板。这种动态查询构建复用了FLARE自身的编码机制：可学习的潜在种子充当一次额外编码调用的查询，将\(N\)个输入token聚合成\(M\)个向量，随后这\(M\)个向量被用作编码-解码对的路由查询，该编码-解码对压缩并重新分配场。由于查询是由FLARE已经执行的相同编码调用产生的，它们的构建继承了其高效的融合SDPA实现，无需新的内核。改动仅限于混合器内部，更具体地说是模板的构建：对于每个固定输入，诱导路由矩阵的秩仍然至多为\(M\)，复杂度仍然为\(\mathcal{O}(NM)\)。残差流不受影响，承载相同数量的块、相同的宽度、相同的残差更新，因此下文报告的优异结果都不是通过加深或加宽网络换来的。如图1 (https://arxiv.org/html/2608.11519#S1.F1)所示，FLARE++仅用一次额外的SDPA调用替换了静态查询参数，其他一切不变。

我们在匹配的骨干下，将FLARE++与FLARE和Transolver系列在标准PDE代理基准上进行比较，发现FLARE++在所有五个基准上都取得了最低的相对\(L^{2}\)误差（表1 (https://arxiv.org/html/2608.11519#S1.T1)），平均将固定模板FLARE的误差降低\(24\%\)，将Transolver-3的误差降低\(31\%\)。与FLARE在潜在预算\(M\)和残差深度\(B\)上的联合消融实验发现，在我们测量的每个配置中，动态路由都优于固定模板FLARE（第5.2节 (https://arxiv.org/html/2608.11519#S5.SS2)）。此外，FLARE++在测得的潜在预算范围内持续改善，而FLARE在该范围内趋于饱和。更进一步，动态路由可以替代深度，FLARE++在更浅的残差深度下就能达到比FLARE更低的误差。在PDE代理之外，同样的替代改善了Long Range Arena的每个任务，并将FLARE的平均分提升了\(2.3\)个点。动态路由并非没有代价，在匹配深度和潜在预算下，其单步时间为FLARE的\(1.3\)–\(1.5\times\)（第C.3节 (https://arxiv.org/html/2608.11519#A3.SS3)），但它通过用更少的块达到给定精度来部分弥补这一代价。为缓解该成本，我们提供了一种精确的token切分实现，将输入token跨设备切分，无需在任何单个设备上聚合完整token序列，并发现并行效率在时间和内存上都保持或接近1。我们在下文中总结贡献。参考说明

图2：二维基准。顶部：每个代理模型接收的输入，即Elasticity上的点云、Darcy上的渗透率场，以及Airfoil和Pipe上的贴体网格。底部：目标场。每个基准一个测试用例，按真实纵横比绘制；Airfoil被裁剪到物体附近，其网格延伸到远场。

表1：标准PDE基准上的测试相对\(L_{2}\)误差（%）；粗体和下划线分别标记最佳和次佳。完全自注意力位于规则线之上，仅作为无限制路由的参考，不参与排名，因为在这些分辨率下它不是候选代理模型。\(\sim\)标记在预算内训练过慢而不可行的情况。

| Blocks \(B\) | FLARE | FLARE++ | \(\Delta\) |
|---|---|---|---|
| 2 | 1.76 | 0.96 | −45% |
| 4 | 0.93 | 0.52 | −44% |
| 8 | 0.68 | 0.38 | −44% |

| Blocks \(B\) | FLARE | FLARE++ | \(\Delta\) |
|---|---|---|---|
| 2 | 1.66 | 1.04 | −38% |
| 4 | 1.04 | 0.75 | −28% |
| 8 | 0.76 | 0.59 | −22% |

图3：在联合\((M,B)\)网格上的固定与动态路由；FLARE为虚线空心标记，FLARE++为实线实心标记。顶部：Elasticity随深度的变化（颜色代表潜在预算\(M\)）以及Darcy随潜在预算的变化（颜色代表深度\(B\)）。在固定误差下水平读取，FLARE++在深度减半时可匹配FLARE；与Elasticity不同，Darcy持续将更大的预算转化为精度，即30 (https://arxiv.org/html/2608.11519#bib.bib42)所报告的秩受限行为。底部：相同的运行结果随深度的变化，\(\Delta\)为相对误差降低，Elasticity上对\(M\)取平均，Darcy上取\(M=128\)。

- • 通过合成的压缩模板实现低秩自注意力。FLARE++使用SDPA从输入构建\(M\)个路由查询，因此输入token决定了它们自身如何被聚合成紧凑的潜在表示并从中重新分配。它保留了FLARE的显式秩\(M\)编码-解码算子、每个头独立的通路、\(\mathcal{O}(NM)\)复杂度以及融合SDPA实现，并且不增加残差流的深度或宽度。
- • 多GPU上下文并行。一种精确的token切分实现将逐点激活和注意力计算分布到多个加速器上，仅通信潜在输出和softmax统计量，因此集合通信负载与输入token数量无关，解码也不需要全收集。在四个rank上，并行效率在时间和内存上都保持或接近1。
- • 准确性和成本的评估。我们在匹配骨干下，将动态路由与固定查询FLARE、Transolver和完全自注意力在标准PDE基准和Long Range Arena上进行比较，并测量每种混合器在token数量跨越三个数量级时的单GPU时间和内存成本，以及多GPU并行效率。我们同时报告该机制有效和无效的情况。

## 2 相关工作

#### 不规则和复杂域上的神经算子。
神经算子学习函数空间之间的映射，已从规则网格公式发展为接受不规则点集和复杂几何的模型\(19 (https://arxiv.org/html/2608.11519#bib.bib12);22 (https://arxiv.org/html/2608.11519#bib.bib13);17 (https://arxiv.org/html/2608.11519#bib.bib15);20 (https://arxiv.org/html/2608.11519#bib.bib14)\)。傅里叶神经算子在规则网格上提供高效的全局混合\(19 (https://arxiv.org/html/2608.11519#bib.bib12)\)，而图和点云算子将学习的PDE映射扩展到非结构化离散\(29 (https://arxiv.org/html/2608.11519#bib.bib11);20 (https://arxiv.org/html/2608.11519#bib.bib14)\)。GNOT将transformer风格的算子学习应用于不规则网格和多个输入函数\(12 (https://arxiv.org/html/2608.11519#bib.bib16)\)；GINO和GINOT在评估查询点处的场之前先对几何进行编码\(20 (https://arxiv.org/html/2608.11519#bib.bib14);21 (https://arxiv.org/html/2608.11519#bib.bib41)\)；区域图算子构建多尺度通信图\(26 (https://arxiv.org/html/2608.11519#bib.bib40)\)。我们的设置遵循这条工作线，但重点放在全局算子内部的token混合器上。

#### 用于PDE代理的完全与高效注意力。
完全自注意力提供无限制的全局通信，但其时间复杂度在token数量上是二次的，而融合实现可以通过避免物化\(N\times N\)分数矩阵来保持峰值内存线性\(36 (https://arxiv.org/html/2608.11519#bib.bib19);9 (https://arxiv.org/html/2608.11519#bib.bib23)\)。Perceiver确立了这类高效注意力模型使用的潜在空间处理范式：一个固定大小的潜在数组与可变长度输入进行交叉注意力，然后使用潜在自注意力进行处理\(14 (https://arxiv.org/html/2608.11519#bib.bib7);13 (https://arxiv.org/html/2608.11519#bib.bib8)\)。LNO将该潜在工作空间范式改编到PDE代理建模，将离散场投影到固定长度表示、在潜在空间中处理、再解码回物理域\(38 (https://arxiv.org/html/2608.11519#bib.bib10)\)。Transolver引入了物理感知切片token，并在每个块内重复投影、潜在自注意力和反投影；Transolver++将该构造扩展到更大的几何体\(39 (https://arxiv.org/html/2608.11519#bib.bib9);23 (https://arxiv.org/html/2608.11519#bib.bib36)\)。Transolver-3将该潜在工作空间系列扩展到工业规模几何体\(44 (https://arxiv.org/html/2608.11519#bib.bib43)\)。FLARE则仅使用潜在查询来为每个固定输入定义一个显式的编码-解码路由矩阵，其秩至多为\(M\)，由两次SDPA调用实现，无需潜在自注意力\(30 (https://arxiv.org/html/2608.11519#bib.bib42)\)。FLARE++保留了这一分解，同时从当前输入合成其路由查询。

#### 决定压缩模板的是什么。