Revisit to Segment: Working Memory Distillation for Reasoning Segmentation
作者: Cilin Yan, Yilun Qiu, Wanyang Zhang, Rui Zu, Xiaolong Jiang, Yao Hu, Jiayin Cai
分类: cs.CV
发布日期: 2026-09-28
💡 一句话要点
提出SWiM框架以增强推理分割能力
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 推理分割 工作记忆 蒸馏训练 多模态模型 深度学习
📋 核心要点
- 现有的多模态大型语言模型在图像分割任务中面临推理能力不足的问题，尤其是在处理复杂场景时。
- 本文提出SWiM框架，通过蒸馏重访过程中的工作记忆，增强模型在推理分割任务中的表现，支持有无工作记忆的推理。
- 实验结果表明，SWiM在多个推理分割基准上达到了最先进的性能，显著提升了分割质量。
📝 摘要（中文）
多模态大型语言模型（MLLMs）通过推理视觉内容和预测目标位置来进行图像分割。研究表明，MLLMs在重访同一图像和查询时，利用自生成的工作记忆作为上下文，能够提升推理分割的效果。为此，本文提出了具有工作记忆蒸馏的推理分割器SWiM，旨在通过蒸馏重访先前尝试中获得的指导，增强模型的推理分割能力。SWiM根据分割质量选择回放，构建工作记忆，并利用记忆条件模型作为教师，提供基于标记的分布监督。通过在策略自蒸馏和基于结果的强化学习的联合优化，SWiM在推理分割基准测试中实现了最先进的性能，验证了工作记忆蒸馏的有效性。
🔬 方法详解
问题定义：本文旨在解决现有多模态大型语言模型在图像分割任务中推理能力不足的问题，尤其是在复杂场景下的表现不佳。现有方法缺乏有效的上下文利用，导致推理结果不够准确。
核心思路：SWiM框架通过蒸馏重访过程中的工作记忆，增强模型的推理能力。该方法利用自生成的工作记忆作为上下文，提升模型在推理分割任务中的表现。
技术框架：SWiM的整体架构包括工作记忆的构建、教师模型的设计和学生模型的训练。首先，根据分割质量选择回放以构建工作记忆，然后使用记忆条件模型作为教师，提供分布监督。最后，学生模型在原始图像和查询的基础上进行训练。
关键创新：SWiM的主要创新在于工作记忆蒸馏的引入，使得模型能够在推理时有效利用历史信息，显著提升了分割质量。这一方法与传统的单一模型推理方法有本质区别。
关键设计：SWiM采用了联合优化策略，包括基于结果的强化学习和在策略自蒸馏的损失函数设计。模型结构上，教师模型和学生模型的设计均考虑了分割任务的特性，确保了信息的有效传递。
🖼️ 关键图片
📊 实验亮点
SWiM在多个推理分割基准测试中表现出色，达到了最先进的性能，具体提升幅度超过了现有基线模型，验证了工作记忆蒸馏的有效性和实用性。
🎯 应用场景
该研究的潜在应用领域包括自动驾驶、医学影像分析和机器人视觉等。通过提升图像分割的推理能力，SWiM框架能够在复杂场景下提供更准确的分割结果，进而推动相关领域的技术进步与应用落地。
📄 摘要（原文）
Multimodal large language models (MLLMs) have approached image segmentation by reasoning about visual content and predicting target locations. Their generated responses contain reasoning traces and localization proposals that can serve as working memory when revisiting the same image and query. Our exploration reveals that MLLMs benefit from using this self-generated working memory as context, leading to enhanced reasoning segmentation. Motivated by this finding, we seek to strengthen the backbone model's reasoning segmentation capabilities by distilling the guidance gained from revisiting prior attempts, enabling it to benefit with or without working memory at inference time. To this end, we propose Reasoning Segmenter with Working Memory (SWiM), a working-memory distillation framework for reasoning segmentation. Specifically, SWiM selects rollouts based on segmentation quality to construct working memory and uses the memory-conditioned model as a teacher. The teacher provides token-level distributional supervision along student-generated trajectories, while the student receives only the original image and query. Joint optimization of on-policy self-distillation and outcome-based reinforcement learning combines working-memory guidance with direct feedback on segmentation quality. Extensive experiments on reasoning segmentation benchmarks demonstrate that SWiM achieves state-of-the-art performance, validating the effectiveness of working-memory distillation.