SAMTok: Representing Any Mask with Two Words

AI summary

Overview

Research area: Computer vision and multimodal large language models — specifically region-level (pixel-wise) understanding and generation.

Technical level: Advanced. The paper assumes familiarity with MLLMs, vector quantization, SAM-style segmentation models, and reinforcement learning (GRPO).

Scope: The paper proposes a discrete mask tokenizer that compresses any 2D region mask into exactly two text-like tokens, letting standard MLLMs learn pixel-level segmentation tasks using only ordinary next-token prediction and text-based reinforcement learning.

What This Paper Is About

Most multimodal LLMs can talk about images but struggle to point at, or produce, precise pixel-level regions. Existing methods solve this with special region encoders, dedicated segmentation decoders, and custom losses that make training complicated and incompatible with text-only data. SAMTok asks whether masks can simply be turned into words — two special tokens per mask — so that a base MLLM like QwenVL learns segmentation the same way it learns language, with no architectural changes.

Key Contributions

1. A new paradigm for mask modeling in MLLMs. Region masks are treated as a new "language," allowing mask understanding and generation to be learned through standard next-token prediction rather than task-specific losses or architecture modifications.

2. SAMTok, a high-fidelity discrete mask tokenizer. Built on SAM2 with a mask encoder and a two-stage residual vector quantizer, it converts any region mask into two discrete tokens and reconstructs it accurately. It is trained on 209M masks spanning part, object, entity, and semantic granularities.

3. A unified mask-token interface and large-scale dataset. Roughly 5M mask understanding and generation conversations are reformatted into purely textual dialogue using SAMTok tokens, enabling co-training of region captioning, region VQA, grounded conversation generation, referring segmentation, scene graph parsing, and multi-round interactive segmentation under a single loss.

4. A textual answer-matching reward for reinforcement learning on mask generation. Because masks are text tokens, reward can be computed by simple string matching over predicted special tokens, eliminating the need for de-tokenization, IoU computation, or external segmentation models. This is described as the first successful use of purely text-based rewards to optimize mask generation.

Main Findings

• Strong performance across many benchmarks without specialized losses. QwenVL-SAMTok reaches state-of-the-art or comparable results on region captioning, region VQA, GCG, referring segmentation, scene graph parsing, and multi-round interactive segmentation, despite being trained only with next-token prediction loss.

• Grounded conversation generation (GCG) improvements. New SOTA over prior best on the validation set by +1.3% METEOR, +5.5% CIDEr, +5.3% AP50, +5.2% mIoU, and +4.7% Recall; similar consistent gains on the test set.

• Multi-round interactive segmentation gains. +7.7% average cIoU over the prior best on MR-RefCOCO/+/g and +10.7% average on MR-PACO, showing that SAMTok handles both mask input and mask output over long interaction chains.

• Text-to-mask (GRES) results. +1.5% average gIoU and +4.3% average N-acc over the previous best; competitive on RefCOCO/+/g, with new SOTA among sub-4B models.

• Large RL gains from purely textual rewards. On GRES, GRPO delivers average gains of +6.8% gIoU, +4.9% cIoU, and +18.9% N-acc; on the validation split specifically, +8.9% gIoU and +21.0% N-acc. On GCG, +4.5% AP50, +2.0% mIoU, +6.6% Recall on average. Caption metrics dipped slightly because no caption-quality reward was used.

• Mask tokens beat text-box interfaces for grounding. Converting predicted mask tokens to boxes improves RefCOCO/+/g accuracy at both 3B and 7B scales (for example, +10.3% on RefCOCO+ test B at 3B), while keeping the same natural-language interaction.

• Strong region captioning and zero-shot generalization. On DLC-Bench, 65.6 average vs. 67.3 for the expert model DAM and 41.2 for base Qwen2.5VL-7B. On MDVP-Bench, it outperforms DAM on three of four scene-type metrics, and it is competitive on video region captioning (VideoRefer-D).

Methodology in Plain English

The researchers started from SAM2, the strong segmentation foundation model, and modified it into a two-way translator.

Encoding. Given an image and a region mask, the mask is fed into SAM's prompt encoder, which produces dense prompt embeddings. These are added to image features from SAM's image backbone and passed through the mask decoder (with its final prediction head removed) to produce a single continuous mask embedding.

Quantization. That continuous embedding is compressed into two discrete codes using residual vector quantization: the nearest codebook entry is looked up, then the leftover residual is looked up again. Two codes, one compact vocabulary.

Decoding. The two codes are turned back into embeddings, summed, and fed into a full SAM model as if they were prompts, reconstructing the 2D mask.

Training the tokenizer. SAMTok is trained on 209M masks with a reconstruction loss (cross-entropy plus Dice) and a commitment loss, with stop-gradients on the codebook.

Training the MLLM. Each mask vocabulary entry becomes a special token added to the MLLM's vocabulary. Any region mask is written as two such tokens, so region captioning, grounding, referring segmentation, GCG, and interactive segmentation all become ordinary image-plus-text dialogues. The base QwenVL model is fine-tuned with standard next-token prediction, with no new modules.

Reinforcement learning. For mask generation, the model rolls out responses, unique mask tokens are extracted, and a reward is computed as the number of predicted mask tokens matching the ground-truth answer divided by the larger of the predicted or ground-truth mask counts. GRPO then optimizes this reward, fully in text space.

Why This Matters

Impact on research. The paper removes a long-standing friction point: pixel-wise capability has typically required bolting segmentation heads onto LLMs and mixing incompatible losses. SAMTok shows that a compact discrete mask vocabulary lets region tasks ride on the same training and RL pipelines used for text, which could simplify how future multimodal systems are built and scaled.

Real-world applications:

• Interactive image and video editing where users describe a region in conversation and expect a precise mask back.
• Assistive and conversational agents that must ground references such as "that object behind the chair" across multiple dialogue turns.
• Document, screenshot, and UI understanding, where the paper shows strong zero-shot region captioning.
• Robotics and autonomous systems that need pixel-accurate, language-driven spatial reasoning about the physical world.

Industry relevance. The approach uses off-the-shelf SAM2 and standard QwenVL models with no custom architecture, which lowers engineering cost for teams wanting to add segmentation to an existing MLLM. The text-only reward design also means RL infrastructure used for language models can be reused directly for mask generation, avoiding extra segmentation models in the loop.

Future Directions

• Scaling the mask tokenizer. Only two tokens per mask are used; whether larger codebooks or more tokens improve fine-grained fidelity, especially for parts and thin structures, is open.

• Reward design for both mask and caption quality. The GCG caption metrics dropped under RL because only mask tokens were rewarded. A combined text-and-mask reward could lift both dimensions together.

• Extending to video and 3D. The paper already touches video region captioning; unified spatio-temporal mask tokens for tracking and 3D scene understanding are a natural extension.

• Generalization of the discrete mask vocabulary. Questions remain about how well the codebook transfers across domains, image resolutions, and unseen mask distributions, and whether multiple tokenizers can be merged.

Target Audience

Researchers and engineers working on multimodal LLMs, region-level vision-language tasks, or segmentation, particularly those interested in discrete representations and reinforcement learning for generation. It is most useful for readers already comfortable with SAM-style segmentation, vector quantization, and GRPO, though the conceptual framing — masks as a language — is accessible to a broader computer vision audience.

Authors’ abstract

Pixel-wise capabilities are essential for building interactive intelligent systems. However, pixel-wise multi-modal LLMs (MLLMs) remain difficult to scale due to complex region-level encoders, specialized segmentation decoders, and incompatible training objectives. To address these challenges, we present SAMTok, a discrete mask tokenizer that converts any region mask into two special tokens and reconstructs the mask using these tokens with high fidelity. By treating masks as new language tokens, SAMTok enables base MLLMs (such as the QwenVL series) to learn pixel-wise capabilities through standard next-token prediction and simple reinforcement learning, without architectural modifications and specialized loss design. SAMTok builds on SAM2 and is trained on 209M diverse masks using a mask encoder and residual vector quantizer to produce discrete, compact, and information-rich tokens. With 5M SAMTok-formatted mask understanding and generation data samples, QwenVL-SAMTok attains state-of-the-art or comparable results on region captioning, region VQA, grounded conversation, referring segmentation, scene graph parsing, and multi-round interactive segmentation. We further introduce a textual answer-matching reward that enables efficient reinforcement learning for mask generation, delivering substantial improvements on GRES and GCG benchmarks. Our results demonstrate a scalable and straightforward paradigm for equipping MLLMs with strong pixel-wise capabilities. Our code and models are available.