【前馈三维重建】SAF3R:前馈式3D重建Transformer的动态稀疏注意力(加速)

目录

  • 摘要
  • 一、前馈式Transformer的全局注意力的分析
  • [二、 注意力头可(按行为)划分为四种类型](#二、 注意力头可(按行为)划分为四种类型)
  • [三、 部分全局注意力头具有高度动态性(依赖于输入)](#三、 部分全局注意力头具有高度动态性(依赖于输入))
  • 四、动态稀疏注意力框架
    • [1. 定制化稀疏注意力核 (Tailored Sparse Attention Kernels)](#1. 定制化稀疏注意力核 (Tailored Sparse Attention Kernels))
    • [2. 离线头分析 (Offline Head Profiling)](#2. 离线头分析 (Offline Head Profiling))
    • [3. 在线模式自适应 (Online Pattern Adaptation)](#3. 在线模式自适应 (Online Pattern Adaptation))
  • 实验结果

标题:SAF3R: Dynamic Sparse Attention for Feed-Forward 3D Reconstruction Transformers

美国匹兹堡大学;美国亚利桑那大学;上海同济大学

链接:https://github.com/jndeng/SAF3R

摘要

  前馈三维重建(F3R)Transformer近期取得了显著成果,但将其扩展至长图像序列仍面临挑战------交叉视图全局注意力机制的二次复杂度会迅速成为主要计算瓶颈。尽管近期研究尝试通过压缩或稀疏注意力机制来提升效率,但未能充分利用全局注意力机制固有的稀疏性和动态特性。

  本文对多种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 数据集的随机采样图像)

  1. 全局注意力模式在模型、层级和头之间具有异质性

  三阶段递进规律(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 K K 核 (Query-Probe Top- K K K Kernel): 针对"垂线头(Vertical-line heads)"。将计算限制在所有查询和少数关键键(Key tokens)之间。由于关键键的位置是动态且依赖输入的,因此需要通过轻量级预计算来实时估计键的位置。
  • DINO Top- K K K 核 (DINO Top- K K K Kernel ): 针对"对应头(Correspondence heads)"。利用 DINO 图像块特征作为对应关系估计器,预计算特征之间的余弦相似度,并选择最相似的 Top- K K K 标记作为键和值。
  • 均匀采样核 (Uniform Sampling Kernel): 针对"扫描头(Scanning heads)"。由于其注意力分布广泛且缺乏结构规律,采用固定步长的均匀采样策略。在保留 25%--50% 键的比例下表现良好。

2. 离线头分析 (Offline Head Profiling)

  提出一种渐进式局部替换搜索策略,目的是为每个注意力头分配最合适的稀疏核类别,以在维持模型性能的同时最大化稀疏性。

  • 搜索策略 : 从使用均匀采样核的基础配置 c base c_{\text{base}} cbase 开始。为每个头定义一个候选配置集 C h \mathcal{C}_h Ch,这些候选配置的稀疏度均等于或高于基线。
  • 误差评估公式 : 使用归一化均方误差(NMSE)来定义在配置 c c c 下,稀疏注意力输出 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}argminE(c)s.t.Ω(c)≤Ω(cbase)

3. 在线模式自适应 (Online Pattern Adaptation)

  由于 Query-Probe 和 DINO Top- K K K 头具有动态特性,本节引入了在线自适应机制,以便在推理阶段根据输入序列细化注意力模式

  • Query-Probe Top- K K K 核的加速 :为了避免计算完整的 O ( N 2 ) \mathcal{O}(N^2) O(N2) 预 Softmax 注意力矩阵,作者利用了内积的线性特性。设预 Softmax 注意力分数为 S = Q K ⊤ / D S = QK^\top / \sqrt{D} S=QK⊤/D 。S S S 按列计算平均值(即所有查询对某个键的平均相似度)等价于计算每个键与平均查询向量 q ˉ \bar{q} qˉ 之间的相似度 ,其中: q ˉ = 1 N ∑ i = 1 N q i \bar{q} = \frac{1}{N} \sum_{i=1}^N q_i qˉ=N1i=1∑Nqi因此,只需计算单个探针向量(Probe vector) q ˉ K ⊤ / D \bar{q}K^\top / \sqrt{D} qˉK⊤/D 即可对键进行排名,从而将复杂度降至 O ( N ) \mathcal{O}(N) O(N)。
  • DINO Top- K K K 核的缓存机制 :提取最后一个 DINO 模块的输出以计算所有图像块之间的全局余弦相似度,并存储每帧匹配到的 Top- K K K 索引。虽然这引入了一次 O ( N 2 ) \mathcal{O}(N^2) O(N2) 计算,但这些索引会被缓存并被后续所有的 DINO Top- K K K 核重复使用,从而有效摊销了整体计算成本。

实验结果

图6:单个不同序列长度场景的延迟与内存占用结果(在分辨率480×640下)。

图7:不同F3R模型的离线头部特征提取结果。经过特征提取后,每个全局注意力头将被分配至四种预定义核类型之一。

相关推荐
大模型码小白7 小时前
【AI】一文讲清 RAG:从大模型局限到企业级知识库落地流程
人工智能·深度学习·学习
m0_547486669 小时前
《人工智能导论:深度学习大模型基础》全套PPT课件2026
人工智能·深度学习·大模型
天天进步20159 小时前
Python全栈项目--基于深度学习的手写识别系统
开发语言·python·深度学习
牛企老板俱乐部11 小时前
广东手板3D打印三大核心产区 采购视角看珠海深圳东莞
3d
doubt。11 小时前
大模型prompt工程Zero-Shot与Few-Shot以及json格式
人工智能·深度学习·机器学习·语言模型·json·prompt
爱炼丹的James12 小时前
目标检测任务
人工智能·深度学习
春末的南方城市12 小时前
消费级显卡迎来实时视频生成!FastVideo 开源 FastWan-QAD,RTX 5090 实现 1.8 秒生成 5 秒 480P 视频!
人工智能·深度学习·计算机视觉·aigc·音视频
@LiX12 小时前
深度学习基本原理
人工智能·深度学习
其美杰布-富贵-李12 小时前
Loss Functions:Logistic Loss、Cross Entropy 与 Hinge Loss
深度学习·损失函数