用少量隐向量实现低秩注意力,让Transformer处理百万级长序列
FLARE: Fast Low-rank Attention Routing Engine
- 通过两步标准注意力计算实现低秩信息路由,仅需少量隐向量
- 单卡支持百万点无结构网格,PDE代理模型精度达当前最优
- 适合长序列建模、科学计算与大规模几何数据任务
自注意力的二次复杂度限制了Transformer在长序列上的扩展性。我们提出快速低秩注意力路由引擎(FLARE),一种通过少量隐向量路由信息的标记混合算子。每层通过仅两次标准缩放点积注意力(SDPA)调用,以最小化编码-解码分解实现输入-输入标记混合矩阵,其秩不超过$M$。由于主导的${O}(NM)$计算完全由标准SDPA表达,FLARE兼容融合注意力内核,避免显式生成$M imes N$投影矩阵。FLARE还为每个注意力头分配独立的隐向量片段,形成头特定的低秩路径混合。实验表明,FLARE可在单个GPU上处理百万点无结构网格,在PDE代理基准上达到最先进准确率,并在长程范围竞技场(Long Range Arena)中超越通用高效注意力方法。我们额外发布了大规模增材制造基准数据集。代码已开源:https://github.com/vpuri3/FLARE.py。
原文摘要 · Abstract (English)
The quadratic complexity of self-attention limits the scalability of transformers on long sequences. We introduce Fast Low-rank Attention Routing Engine (FLARE), a token-mixing operator that realizes low-rank attention by routing information through a small set of latent tokens. Each layer induces an input-input token mixing matrix of rank at most $M$ via a minimal encode-decode factorization implemented using only two standard scaled dot-product attention (SDPA) calls. Because the dominant ${O}(NM)$ computation is expressed purely in terms of standard SDPA, FLARE is compatible with fused attention kernels and avoids materializing $M\times N$ projection matrices. FLARE further assigns disjoint latent slices to each attention head, yielding a mixture of head-specific low-rank pathways. Empirically, FLARE scales to one-million-point unstructured meshes on a single GPU, achieves state-of-the-art accuracy on PDE surrogate benchmarks, and outperforms general-purpose efficient-attention methods on the Long Range Arena suite. We additionally release a large-scale additive manufacturing benchmark dataset. Our code is available at https://github.com/vpuri3/FLARE.py.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。