EdgeTAM: On-Device Track Anything Model

📄 arXiv: 2501.07256v1 📥 PDF

作者: Chong Zhou, Chenchen Zhu, Yunyang Xiong, Saksham Suri, Fanyi Xiao, Lemeng Wu, Raghuraman Krishnamoorthi, Bo Dai, Chen Change Loy, Vikas Chandra, Bilge Soran

分类: cs.CV

发布日期: 2025-01-13

备注: Code will be released at https://github.com/facebookresearch/EdgeTAM

💡 一句话要点

提出EdgeTAM，通过2D空间感知器加速SAM 2，实现移动端视频分割。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 视频分割 移动端部署 模型压缩 空间感知器 Transformer 蒸馏训练

📋 核心要点

  1. 现有SAM优化方法主要集中于图像编码器压缩，忽略了SAM 2中记忆注意力模块的延迟瓶颈。
  2. EdgeTAM利用2D空间感知器，通过轻量级Transformer和可学习查询，高效编码帧级记忆，降低计算成本。
  3. EdgeTAM在iPhone 15 Pro Max上实现了16 FPS，并在多个视频分割数据集上取得了优异的J&F指标。

📝 摘要（中文）

本文旨在提高Segment Anything Model 2 (SAM 2)的效率，使其能够在移动设备上运行，同时保持可比的性能。尽管已有工作致力于优化SAM的效率，但它们主要集中在压缩图像编码器上。本文的基准测试表明，新引入的记忆注意力模块也是延迟瓶颈。为此，本文提出了EdgeTAM，它利用一种新颖的2D空间感知器来降低计算成本。具体来说，提出的2D空间感知器使用包含一组固定可学习查询的轻量级Transformer对密集存储的帧级记忆进行编码。考虑到视频分割是一个密集预测任务，保持记忆的空间结构至关重要，因此查询被分成全局级别和补丁级别组。本文还提出了一种蒸馏流程，进一步提高了性能，而没有推理开销。结果表明，EdgeTAM在iPhone 15 Pro Max上以16 FPS的运行速度，在DAVIS 2017、MOSE、SA-V val和SA-V test上分别实现了87.7、70.0、72.3和71.7的J&F指标。

🔬 方法详解

问题定义：SAM 2在视频分割任务中表现出色，但计算量巨大，难以在移动设备等资源受限的平台上部署。现有SAM优化方法主要关注图像编码器的压缩，忽略了SAM 2中新引入的记忆注意力模块带来的延迟瓶颈。因此，需要一种更高效的SAM 2实现，使其能够在移动设备上运行，同时保持良好的分割性能。

核心思路：EdgeTAM的核心思路是利用一种新颖的2D空间感知器来高效编码帧级记忆，从而降低计算成本。通过轻量级的Transformer和可学习的查询，可以在保持空间结构信息的同时，显著减少需要处理的token数量。此外，采用蒸馏训练进一步提升模型性能，且不增加推理负担。

技术框架：EdgeTAM的整体框架包括：1) 图像编码器（可以采用现有的轻量化SAM编码器）；2) 2D空间感知器：用于编码帧级记忆，降低计算量；3) 提示编码器：处理用户提供的提示信息；4) 解码器：根据编码后的图像特征、记忆和提示信息，生成分割掩码。整个流程可以端到端训练。

关键创新：EdgeTAM的关键创新在于提出的2D空间感知器。与传统的注意力机制需要计算所有token之间的关系不同，2D空间感知器使用一组固定的可学习查询来提取记忆中的关键信息。此外，为了保持空间结构信息，查询被分为全局级别和补丁级别两组，分别关注全局上下文和局部细节。这种设计显著降低了计算复杂度，同时保持了分割性能。

关键设计：2D空间感知器使用轻量级的Transformer结构，包含多个Transformer层。查询的数量是一个关键参数，需要根据计算资源和性能要求进行调整。损失函数包括分割损失（例如Dice Loss或Cross-Entropy Loss）和蒸馏损失。蒸馏损失用于将大型模型的知识迁移到EdgeTAM中，进一步提升性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

EdgeTAM在iPhone 15 Pro Max上实现了16 FPS的运行速度，同时在DAVIS 2017、MOSE、SA-V val和SA-V test数据集上分别取得了87.7、70.0、72.3和71.7的J&F指标。这些结果表明，EdgeTAM在移动端设备上实现了高效且高质量的视频分割。

🎯 应用场景

EdgeTAM可应用于移动端视频编辑、AR/VR应用、智能监控、自动驾驶等领域。它能够在资源受限的设备上实现实时的视频对象分割，为用户提供更智能、更便捷的交互体验，并降低部署成本。

📄 摘要（原文）

On top of Segment Anything Model (SAM), SAM 2 further extends its capability from image to video inputs through a memory bank mechanism and obtains a remarkable performance compared with previous methods, making it a foundation model for video segmentation task. In this paper, we aim at making SAM 2 much more efficient so that it even runs on mobile devices while maintaining a comparable performance. Despite several works optimizing SAM for better efficiency, we find they are not sufficient for SAM 2 because they all focus on compressing the image encoder, while our benchmark shows that the newly introduced memory attention blocks are also the latency bottleneck. Given this observation, we propose EdgeTAM, which leverages a novel 2D Spatial Perceiver to reduce the computational cost. In particular, the proposed 2D Spatial Perceiver encodes the densely stored frame-level memories with a lightweight Transformer that contains a fixed set of learnable queries. Given that video segmentation is a dense prediction task, we find preserving the spatial structure of the memories is essential so that the queries are split into global-level and patch-level groups. We also propose a distillation pipeline that further improves the performance without inference overhead. As a result, EdgeTAM achieves 87.7, 70.0, 72.3, and 71.7 J&F on DAVIS 2017, MOSE, SA-V val, and SA-V test, while running at 16 FPS on iPhone 15 Pro Max.