Revisit to Segment: Working Memory Distillation for Reasoning Segmentation

Cilin Yan; Yilun Qiu; Wanyang Zhang; Rui Zu; Xiaolong Jiang; Jiayin Cai; Yao Hu

cs.CVSep 28, 2026

Revisit to Segment: Working Memory Distillation for Reasoning Segmentation

Authors: Cilin Yan, Yilun Qiu, Wanyang Zhang, Rui Zu, Xiaolong Jiang, Jiayin Cai, Yao Hu

Organizations: Xiaohongshu · NTU · PKU

Abstract

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.

Figures & tables

Figure 1: Working-memory-guided reasoning segmentation. Given an image and a query, the MLLM generates reasoning traces and localization proposals that form its working memory. It then revisits the same image and query with this memory as additional input, predicting boxes and points that prompt a frozen SAM2 model to produce the final mask.

Panel of Figure 2 — Figure 2: Reasoning segmentation without and with working memory (WM). (a) Quantitative results; gains are in percentage points. (b) A qualitative example of improved target identification.

Figure 3: Overview of SWiM. Top- $K$ rollouts ranked by segmentation quality form working memory. OPSD transfers token-level guidance from the working-memory-conditioned teacher to the student along the same sampled trajectories. GRPO provides outcome-based feedback on segmentation quality. Both objectives jointly update the student, which receives only the original image and query.

17,18 Method

Base model (Size)

RS test

RS-R

RS-X test

overall

P/F

C/KI

C/R

C/MH

Average

gIoU

cIoU

gIoU

cIoU

gIoU

cIoU

gIoU

cIoU

gIoU

cIoU

gIoU

cIoU

gIoU

cIoU

gIoU

cIoU

Supervised Fine-Tuning / Reinforcement Learning

LISA

Llama2-13B

51.5

51.3

52.5

53.3

25.1

26.0

27.9

30.6

28.2

29.2

26.8

25.8

13.8

16.8

43.0

43.5

SegLLM

LLaVA1.5-7B

52.4

48.4

–

READ

LLaVA1.5-13B

62.2

62.8

–

Table 1: Performance results on reasoning segmentation benchmarks. “WM” denotes working memory, and “MV” denotes the majority voting strategy adapted for segmentation. The best results in each column are highlighted in bold , and the second-best results are underlined . Higher values indicate better segmentation performance across all metrics.

Figure 4: Effects of training and inference configurations in SWiM. Results are averaged over RS, RS-R, and RS-X. Floating bars show gains in percentage points; solid bars show absolute scores.

GRPO	OPSD	RS		RS-R		RS-X		Average
GRPO	OPSD	gIoU	cIoU	gIoU	cIoU	gIoU	cIoU	gIoU	cIoU
		69.6	64.4	72.4	67.6	55.1	50.0	65.7	60.6
$\checkmark$		70.6	64.7	73.0	70.2	55.1	47.8	66.2	60.9
	$\checkmark$	70.1	65.9	72.3	68.4	54.8	48.4	65.7	60.9
$\checkmark$	$\checkmark$	70.8	65.8	73.0	70.8	55.6	49.3	66.5	62.0

Table 2: Ablation of training objectives. We compare GRPO, OPSD, and their joint optimization on RS, RS-R, and RS-X without working memory at inference.

Figure 5: Effect of the distillation weight

\lambda

. gIoU and cIoU are averaged over RS, RS-R, and RS-X.

Figure 6: Working memory versus majority voting on ReasonSeg. Each row shows a query, the input image, a greedy prediction, eight sampled candidate masks, a WM-guided prediction, and an MV-only prediction. WM-guided inference revisits the image and query using the corresponding responses as working memory, while MV directly aggregates the same eight masks.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Table 9

Method	RefCOCO	RefCOCO+	RefCOCOg
Method	testA	testA	test
LISA-7B	76.5	67.4	68.5
SegLLM	81.5	73.0	73.6
READ	80.2	73.7	71.4
CoReS	78.6	70.0	70.7
CoPRS-7B	85.3	80.3	76.2

Appendix

Table 5: Performance on referring expression segmentation benchmarks. We report cIoU results.

Training		Inference		RS		RS-R		RS-X		Average
WM	OPSD	WM	MV	gIoU	cIoU	gIoU	cIoU	gIoU	cIoU	gIoU	cIoU
				68.6	60.2	71.3	65.7	54.0	49.2	64.6	58.4
$\checkmark$				69.6	64.4	72.4	67.6	55.1	50.0	65.7	60.6
$\checkmark$	$\checkmark$			70.8	65.8	73.0	70.8	55.6	49.3	66.5	62.0
$\checkmark$	$\checkmark$	$\checkmark$		71.1	67.0	72.7	70.8	57.2	51.0	67.0	62.9
$\checkmark$	$\checkmark$	$\checkmark$	$\checkmark$	71.6	67.3	73.5	70.3	58.1	51.5	67.8	63.0

Appendix

Table 6: Progressive evaluation of the core components of SWiM. We report results on RS, RS-R, and RS-X using the Qwen3-VL-8B backbone.

WM	GT	Random	RS		RS-R		RS-X		Average
WM	GT	Random	gIoU	cIoU	gIoU	cIoU	gIoU	cIoU	gIoU	cIoU
		$\checkmark$	55.4	52.4	56.0	56.0	43.0	37.7	51.5	48.7
$\checkmark$			70.8	65.8	73.0	70.8	55.6	49.3	66.5	62.0
	$\checkmark$		70.7	65.5	73.2	70.1	54.4	48.3	66.1	61.3
$\checkmark$	$\checkmark$		70.7	67.1	72.6	70.2	55.6	48.2	66.3	61.8

Appendix

Table 7: Comparison of teacher context configurations. WM denotes working memory, GT denotes ground-truth visual hints, and Random denotes surrogate annotations generated by SAM2 from random point prompts.

Loss	RS		RS-R		RS-X test		Average
Loss	gIoU	cIoU	gIoU	cIoU	gIoU	cIoU	gIoU	cIoU
Reverse KL	70.5	64.4	73.1	70.8	54.0	46.0	65.9	60.4
Forward KL	70.2	63.9	73.1	70.4	53.8	48.6	65.7	61.0
JSD	70.8	65.8	73.0	70.8	55.6	49.3	66.5	62.0

Appendix

Table 8: Comparison of OPSD divergence objectives. All variants use working memory as the teacher context during training and are evaluated without working memory or majority voting.

$\lambda$	RS		RS-R		RS-X		Average
$\lambda$	gIoU	cIoU	gIoU	cIoU	gIoU	cIoU	gIoU	cIoU
0.1	70.8	65.8	73.0	70.8	55.6	49.3	66.5	62.0
0.5	70.5	65.0	72.7	68.8	55.0	48.5	66.1	60.8
1.0	70.8	63.4	73.0	71.1	54.4	49.5	66.1	61.3
2.0	70.8	65.7	72.9	69.9	54.9	49.9	66.2	61.9
10.0	70.3	65.1	72.4	69.6	54.0	48.0	65.6	60.9

Appendix

Table 9: Effect of the distillation weight

\lambda

. All settings are evaluated without working memory.

Figure 8: Question–answer visualizations on ReasonSeg-X. Each panel shows the query, original image, ground-truth mask, and predicted mask with the model-generated box and point, followed by verbatim thinking excerpts and answer labels. Green overlays denote masks, orange rectangles denote predicted boxes, and pink markers denote predicted points. IoUs are per-example scores. The examples cover route planning, visible constraints on movement, physical obstruction, and role understanding.

Figure 9: Qualitative comparison with StAR on the ReasonSeg series. Each example shows the image and query alongside predictions from our reproduced StAR baseline and SWiM, both using Qwen3-VL-8B. These examples highlight improvements in target identification, region coverage, and small-target localization.

Figure 10: Two types of low-IoU predictions on the ReasonSeg series. All predictions are generated by SWiM with the Qwen3-VL-8B backbone using working memory at inference. The top row shows target-granularity ambiguities: handrail versus railing, and control unit versus knob. The bottom row shows incomplete SAM2 masks despite plausible object descriptions and localization prompts: a barbell and glass windows. Each panel includes the original image, GT mask, prediction with box and point, and verbatim thinking excerpts with answer labels. These examples diagnose possible error sources; they do not establish annotation errors or a causal attribution to the mask decoder.

Explore similar work

Reason Twice: Segmentation via Candidate Discovery and Comparative Reasoning
Jun 8, 2026Xinyan Gao, Haoran Hao, Xiangyu YueNeural Mask Estimation Multimodal Large Language Models
MMLab, The Chinese University of Hong Kong · Nanjing University
DGSeg: Dynamic Gating of Semantic-Spatial Guided Predictions for Reasoning Segmentation
Jul 6, 2026Ruizhe Zeng, Siyu Cao, Lu Zhang +1Vision-Language Model Grounding Object-Centric Cues
State Key Laboratory of Multimodal Artificial Intelligence Systems, Institute of Automation, Chinese Academy of Sciences, Beijing 100190, China · School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing 100049, China · Nanjing Artificial Intelligence Research of IA, Nanjing, 211100, China
CR-Seg: Attention-Guided and CoT-Enhanced Coarse-to-Refined Reasoning Segmentation
Jun 2, 2026Yifan Cao, Xiaocui Yang, Faxian Wan +3Chain-of-Thought Reasoning
School of Computer Science and Engineering, Northeastern University Shenyang 110819, China

Jun 8, 2026, cs.CVJ/K move · Enter open · S save

Reason Twice: Segmentation via Candidate Discovery and Comparative Reasoning

Xinyan Gao, Haoran Hao, Xiangyu Yue

MMLab, The Chinese University of Hong Kong · Nanjing University

Neural Mask Estimation Multimodal Large Language Models Comparison Discovery

Open paper AskPDF arXiv

The rapid development of pretrained foundation models has enabled more general image segmentation. Multimodal large language models (MLLMs) have been widely explored for image segmentation with complex queries that require high-level reasoning. Despite promising progress, existing methods are often constrained by limited training data and the gap between MLLMs and mask generation modules. To better transfer MLLMs' perception and reasoning ability to complex reasoning-based segmentation tasks, we propose a two-stage framework Rea2Seg for mask generation and selection. Specifically, the framework first identifies potential regions as candidate masks based on the attention maps of a segmentation MLLM. It then employs an MLLM to reason over the question and candidate masks and assign scores to each mask. The final segmentation result is obtained by reranking the candidates and selecting the highest-scoring mask, reformulating image segmentation as candidate discovery followed by discriminative mask selection. We also notice that a large portion of questions in existing benchmarks focus on commonsense reasoning, and these questions usually do not fully require joint visual observation and reasoning. To address this issue, we introduce a new benchmark called ReasonSeg-SGDR that comprehensively evaluates a model's perception, grounding, and reasoning abilities across multiple dimensions, including discriminative recognition, spatial reasoning, geometric reasoning, and multi-step reasoning, with fine-grained mask generation. In addition, we collect training data to enhance MLLMs' ability to jointly understand multimodal queries and candidate masks, and to assign scores through reasoning. Experimental results on the proposed benchmark and ReasonSeg demonstrate the effectiveness of the unified mask generation and selection framework.

Abstract

Explore similar work

Reason Twice: Segmentation via Candidate Discovery and Comparative Reasoning

DGSeg: Dynamic Gating of Semantic-Spatial Guided Predictions for Reasoning Segmentation

CR-Seg: Attention-Guided and CoT-Enhanced Coarse-to-Refined Reasoning Segmentation