【前馈三维重建】SAF3R：前馈式3D重建Transformer的动态稀疏注意力（加速）

目录摘要一、前馈式Transformer的全局注意力的分析二、 注意力头可按行为划分为四种类型三、 部分全局注意力头具有高度动态性依赖于输入四、动态稀疏注意力框架1. 定制化稀疏注意力核 (Tailored Sparse Attention Kernels)2. 离线头分析 (Offline Head Profiling)3. 在线模式自适应 (Online Pattern Adaptation)实验结果标题SAF3R: Dynamic Sparse Attention for Feed-Forward 3D Reconstruction Transformers美国匹兹堡大学美国亚利桑那大学上海同济大学链接https://github.com/jndeng/SAF3R摘要前馈三维重建F3RTransformer近期取得了显著成果但将其扩展至长图像序列仍面临挑战——交叉视图全局注意力机制的二次复杂度会迅速成为主要计算瓶颈。尽管近期研究尝试通过压缩或稀疏注意力机制来提升效率但未能充分利用全局注意力机制固有的稀疏性和动态特性。本文对多种F3R Transformer中的全局注意力机制进行了全面分析揭示了注意力模式在不同层和注意力头之间具有高度异质性、动态变化及极高的稀疏性。基于这些发现我们提出了SAF3R——一个专为F3R Transformer设计的无需训练的动态稀疏注意力框架。SAF3R将定制化的稀疏注意力机制与离线头特征分析及高效的在线适配策略相结合以匹配输入依赖的注意力行为。大量实验表明SAF3R能在保持相机位姿估计与三维重建质量的同时实现较高的稀疏率相较于现有方法在F3R Transformer上实现了显著的端到端加速效果。图1左图SAF3R可加速前馈式三维重建F3R模型的运行速度同时仍能保持重建质量右图SAF3R在相机位姿估计任务上优于现有的高效F3R方法并实现了与原始模型相近的性能。图2现有F3R Transformer所共用的架构概览。图像块首先由DINO编码器进行编码随后通过N个交替的局部与全局注意力模块进行处理每个模块均包含多头自注意力MHA和前馈网络FFN一、前馈式Transformer的全局注意力的分析采用多头Head-level分析方法分析四个代表性 F3R 模型VGGT、π 3 \pi^3π3、MapAnything 和 DA3中的跨视角全局注意力使用 7Scenes 数据集的随机采样图像全局注意力模式在模型、层级和头之间具有异质性三阶段递进规律Stage-wise progression浅层Shallow layers 主要进行局部聚合或聚焦于少数显著图像块表现出较高的集中度和静态稀疏性跨视角交互较少。中层Middle layers 建立跨视角的点对点对应关系稀疏度依然较高但呈现出与内容相关的语义结构。深层Later layers 注意力变得更加弥散Diffuse主要用于细化特征表示以及相对于锚点帧Anchor frame的相机姿态。注意力头的粒度差异 即便在同一层内不同注意力头Heads的模式差异也极大同时不同 F3R 模型间的注意力行为也有所不同例如 DA3 未出现 VGGT 中的位置编码诱导的局部注意力模式。图3跨层(G)与跨头(H)的代表性全局注意力模式。注意力图中的灰色虚线用于区分来自不同图像的标记。类似的观察结果也适用于MapAnything和π3二、 注意力头可按行为划分为四种类型位置头(Position heads)由位置编码驱动而非语义特征注意力模式固定如跨图像块固定或主要关注定义相机坐标系的锚点图像Anchor image。垂线头(Vertical-line heads)大多数查询Queries持续关注少数关键键Key tokens如显著图像块或汇聚节点 Sink tokens在注意力图上形成明显的“垂直线”。对应头(Correspondence heads)由语义特征驱动每个查询选择性地关注对应相同 3D 位置的键具有极高的稀疏性。扫描头(Scanning heads)占绝大多数注意力模式呈现多样化分布在所有图像的各个位置稀疏度通常较低。三、 部分全局注意力头具有高度动态性依赖于输入静态与动态行为并存 部分头的注意力图跨输入保持一致而另一部分头的具体注意力模式则依赖于输入的视觉内容。结论启发 稀疏注意力模式Sparse attention patterns不应被固定而应根据输入内容动态调整。图4(a)和(b)展示了对应头的高稀疏性而©则展示了垂直线头的内容依赖型动态特性。四、动态稀疏注意力框架1. 定制化稀疏注意力核 (Tailored Sparse Attention Kernels)基于前文对注意力头Attention Head的分类为不同类型的头设计了特定的稀疏注意力核将注意力计算复杂度从O ( N 2 ) \mathcal{O}(N^2)O(N2)降低至O ( N ) \mathcal{O}(N)O(N)。主要包含四种核静态核 (Static Kernel) 针对“位置头Positional heads”。由于其与内容无关采用固定的稀疏注意力模式例如允许所有查询Queries关注指定的锚点帧或广播局部注意力图。Query-Probe Top-K KK核 (Query-Probe Top-K KKKernel) 针对“垂线头Vertical-line heads”。将计算限制在所有查询和少数关键键Key tokens之间。由于关键键的位置是动态且依赖输入的因此需要通过轻量级预计算来实时估计键的位置。DINO Top-K KK核 (DINO Top-K KKKernel) 针对“对应头Correspondence heads”。利用 DINO 图像块特征作为对应关系估计器预计算特征之间的余弦相似度并选择最相似的 Top-K KK标记作为键和值。均匀采样核 (Uniform Sampling Kernel) 针对“扫描头Scanning heads”。由于其注意力分布广泛且缺乏结构规律采用固定步长的均匀采样策略。在保留 25%–50% 键的比例下表现良好。2. 离线头分析 (Offline Head Profiling)提出一种渐进式局部替换搜索策略目的是为每个注意力头分配最合适的稀疏核类别以在维持模型性能的同时最大化稀疏性。搜索策略 从使用均匀采样核的基础配置c base c_{\text{base}}cbase​开始。为每个头定义一个候选配置集C h \mathcal{C}_hCh​这些候选配置的稀疏度均等于或高于基线。误差评估公式 使用归一化均方误差NMSE来定义在配置c cc下稀疏注意力输出O sparse ( h , c ) \mathbf{O}_{\text{sparse}}^{(h, c)}Osparse(h,c)​相比于精确的全注意力输出O full ( h ) \mathbf{O}_{\text{full}}^{(h)}Ofull(h)​的近似误差E ( c ) E D calib [ ∥ O full ( h ) − O sparse ( h , c ) ∥ 2 2 ∥ O full ( h ) ∥ 2 2 ] E(c) \mathbb{E}_{\mathcal{D}_{\text{calib}}} \left[ \frac{\Vert{} \mathbf{O}_{\text{full}}^{(h)} - \mathbf{O}_{\text{sparse}}^{(h, c)} \Vert{}_2^2}{\Vert{} \mathbf{O}_{\text{full}}^{(h)} \Vert{}_2^2} \right]E(c)EDcalib​​[∥Ofull(h)​∥22​∥Ofull(h)​−Osparse(h,c)​∥22​​]其中D calib \mathcal{D}_{\text{calib}}Dcalib​是用于校准的验证数据集优化目标公式 在不超过基线计算复杂度Ω ( c base ) \Omega(c_{\text{base}})Ω(cbase​)的约束下选择使近似误差最小化的配置c ∗ c^*c∗c ∗ arg ⁡ min ⁡ c ∈ C h ∪ { c base } E ( c ) s.t. Ω ( c ) ≤ Ω ( c base ) c^* \underset{c \in \mathcal{C}_h \cup \{c_{\text{base}}\}}{\arg\min} E(c) \quad \text{s.t.} \quad \Omega(c) \le \Omega(c_{\text{base}})c∗c∈Ch​∪{cbase​}argmin​E(c)s.t.Ω(c)≤Ω(cbase​)3. 在线模式自适应 (Online Pattern Adaptation)由于 Query-Probe 和 DINO Top-K KK头具有动态特性本节引入了在线自适应机制以便在推理阶段根据输入序列细化注意力模式。Query-Probe Top-K KK核的加速为了避免计算完整的O ( N 2 ) \mathcal{O}(N^2)O(N2)预 Softmax 注意力矩阵作者利用了内积的线性特性。设预 Softmax 注意力分数为S Q K ⊤ / D S QK^\top / \sqrt{D}SQK⊤/D​。S SS按列计算平均值即所有查询对某个键的平均相似度等价于计算每个键与平均查询向量q ˉ \bar{q}qˉ​之间的相似度其中q ˉ 1 N ∑ i 1 N q i \bar{q} \frac{1}{N} \sum_{i1}^N q_iqˉ​N1​i1∑N​qi​因此只需计算单个探针向量Probe vectorq ˉ K ⊤ / D \bar{q}K^\top / \sqrt{D}qˉ​K⊤/D​即可对键进行排名从而将复杂度降至O ( N ) \mathcal{O}(N)O(N)。DINO Top-K KK核的缓存机制提取最后一个 DINO 模块的输出以计算所有图像块之间的全局余弦相似度并存储每帧匹配到的 Top-K KK索引。虽然这引入了一次O ( N 2 ) \mathcal{O}(N^2)O(N2)计算但这些索引会被缓存并被后续所有的 DINO Top-K KK核重复使用从而有效摊销了整体计算成本。实验结果图6单个不同序列长度场景的延迟与内存占用结果(在分辨率480×640下。图7不同F3R模型的离线头部特征提取结果。经过特征提取后每个全局注意力头将被分配至四种预定义核类型之一。