SAMTok: Representing Any Mask with Two Words

ArXiv: 2601.16093

🎯 Pitch

SAMTok compresses any segmentation mask into just two text tokens, enabling base multimodal LLMs to perform pixel-level tasks via standard next-token prediction with no architecture changes. This tokenization unlocks reinforcement learning for mask generation using only text-matching rewards, yielding massive gains on GRES and GCG—demonstrating that RL on visual tasks works only when masks become language tokens, not continuous embeddings.

1. Executive Summary

This paper introduces SAMTok, a discrete mask tokenizer that converts any region mask into two special text tokens and reconstructs the mask from those tokens, enabling base multimodal LLMs to learn pixel-wise capabilities through standard next-token prediction and reinforcement learning without architectural modifications or specialized segmentation decoders. Trained on 209M diverse masks, SAMTok is integrated with the QwenVL series and finetuned on 5M mask-understanding and generation data samples, achieving state-of-the-art or comparable results across region captioning, region VQA, grounded conversation, referring segmentation, scene graph parsing, and multi-round interactive segmentation. The discrete textual representation of masks enables a purely textual answer-matching reward for GRPO-based reinforcement learning on mask generation, yielding 8.9% gIoU and 21.0% N-acc gains on GRES, and 4.7% AP50 and 6.6% Recall gains on GCG — surpassing prior SOTA methods by 4.3% gIoU, 8.3% N-acc, 8.3% AP50, and 8.4% Recall respectively — establishing that RL-based mask generation optimization works with text-only rewards only when masks are represented as language tokens rather than continuous embeddings.

2. Context and Motivation

The Core Problem: Pixel-Wise Capabilities Are Hard to Scale in Multimodal LLMs

The fundamental question this paper addresses is both concrete and architecturally significant: how can we equip general-purpose multimodal large language models with pixel-level visual understanding — the ability to both accept and generate precise region masks — without adding task-specific components, loss functions, or training procedures? This matters because pixel-wise capabilities (grounding language to specific image regions, generating segmentation masks from natural language queries, understanding what a marked region contains) are essential for building interactive intelligent systems — from visual assistants that can collaboratively edit images, to robots that must reason about spatial relationships between objects, to accessibility tools that describe specific parts of a scene to users with visual impairments.

The problem is not that MLLMs lack pixel-wise capabilities entirely. Rather, the problem is that existing approaches to adding these capabilities have created a fragmented landscape where architectural complexity, training incompatibility, and representational inefficiency collectively prevent scaling. The authors identify four specific challenges that motivated their work (Section 1):

1. No unified design for mask input and mask output. Current state-of-the-art pixel-wise MLLMs treat mask understanding (accepting a region mask as input to answer questions about it) and mask generation (producing a segmentation mask as output) as fundamentally different operations requiring different mechanisms. For mask input, models rely on complex region-level feature pooling — extracting ROI features from feature maps, encoding masks as dense prompt embeddings, or using specialized region encoders. For mask output, models depend on carefully designed segmentation decoders that convert hidden states back into 2D masks, often requiring auxiliary losses (Dice loss, binary cross-entropy) that are incompatible with the standard language modeling objective. This bifurcation means that tasks requiring both input and output masks simultaneously — such as multi-round interactive segmentation, where the model must understand previously segmented regions and generate new ones — become engineering challenges that require coordinating two distinct subsystems.

"Existing models rely on complex region-level feature pooling designs [30, 62, 85, 87], while mask output depends on carefully designed segmentation decoders [25, 53, 73, 85, 93]."

2. Continuous mask representations block reinforcement learning. Perhaps the most subtle but practically consequential limitation is that current SOTA pixel-wise MLLMs use continuous embeddings to connect the MLLM with the segmentation head — the language model produces some special token whose hidden state is then fed into a mask decoder. This continuous interface makes it impossible to directly apply standard RL algorithms (such as GRPO, PPO, or DPO) that operate on discrete token sequences and rely on verifiable text outputs for reward computation. The authors state this explicitly:

"Current state-of-the-art pixel-wise MLLMs [73, 85, 101] cannot directly and concisely apply reinforcement learning (RL) to mask generation tasks since they use continuous embeddings to connect the MLLM with the segmentation head."

This is a critical bottleneck because RL-based post-training has become the dominant paradigm for pushing LLMs beyond their supervised fine-tuning ceiling — as seen in reasoning models like DeepSeek-R1 and OpenAI o1/o3. If pixel-wise MLLMs cannot participate in this RL-driven improvement cycle, they risk being left behind as the rest of the field advances. Some works have attempted workarounds by predicting bounding boxes or points as intermediate representations, then applying RL to those coordinates, and finally converting boxes/points to masks using SAM — but this adds substantial system complexity and loses the precision advantages of direct mask generation.

3. Co-training requirements create engineering and scaling debt. The specialized modules added for mask understanding and generation typically must be co-trained with the MLLM backbone. This means that adding pixel-wise capabilities to a new base model (say, upgrading from QwenVL-2 to QwenVL-3) requires re-engineering the integration, retraining the specialized heads, and managing multiple training objectives simultaneously. Moreover, different training losses for different tasks (segmentation loss for mask generation, cross-entropy for text generation) introduce substantial complexity for scaling training with VQA and pure text data. A VQA batch might only need the language modeling head, while a segmentation batch needs the mask decoder — the forward pipelines differ, the loss scales differ, and the optimization dynamics are harder to balance.

"Specially designed modules added for mask understanding and generation capabilities typically require co-training with the MLLM [53, 68, 73, 85, 93, 101]. In addition, the different training losses and forward pipelines introduce substantial complexity for scaling training with VQA and pure text data."

4. Token-inefficient mask representations limit practical deployment. Several recent works have explored treating masks as text or images to avoid the architectural complexity of specialized decoders:

• Text-based approaches represent masks as run-length encoding (RLE), polygon coordinates, or similar serializations
• Image-based approaches treat each mask as a separate image that the model autoregressively generates

However, both approaches share a fatal flaw: a single mask requires dozens or even hundreds of tokens. In RLE encoding, a complex region boundary can require hundreds of coordinate tokens. In image-based approaches, generating a mask image at reasonable resolution (say, 256×256) requires thousands of tokens autoregressively. Beyond the obvious inefficiency — a single mask consuming more tokens than an entire paragraph of text — this token bloat has architectural consequences: (a) it wastes context window budget that could be used for the actual conversation or reasoning; (b) it slows inference dramatically since each mask token must be autoregressively generated; (c) it makes interleaved text-mask generation (where multiple masks appear throughout a caption) prohibitively expensive.

The authors frame this as "an essential question":

"How can we non-intrusively endow base MLLMs (such as the QwenVL series) with pixel-wise capabilities, making the learning process as simple as VQA training — requiring only next-token prediction loss for supervised fine-tuning (SFT) and straightforward reinforcement learning (RL)?"

The word "non-intrusively" is doing significant work here. The goal is not just to achieve good performance on segmentation tasks — it's to achieve it through a mechanism that is architecturally invisible to the base MLLM, requiring no new modules, no new loss functions, and no changes to the training pipeline. If masks are just another language, then adding pixel-wise capabilities is as simple as adding new vocabulary tokens and training data — exactly how you'd add a new language to a multilingual model.

Why This Problem Matters Now

The paper's timing reflects a critical juncture in the MLLM field. Over 2024–2025, general-purpose MLLMs (GPT-4o, Gemini, QwenVL, InternVL) have achieved remarkable performance on image-level understanding — describing whole images, answering questions about visual content, reading text in screenshots. However, region-level understanding remains a differentiator between "good enough" and truly capable visual AI systems. Consider these concrete use cases:

• Medical imaging: A radiologist needs to ask "what is this specific region?" by circling an anomaly in a CT scan, not by verbally describing its location.
• Collaborative design: A user marks a region in a UI mockup and says "make this button blue" — the system must understand which region is being referenced and generate an edited output.
• Robotics: A robot must segment "the red cup to the left of the blue bowl" and then grasp it — requiring joint language grounding, spatial reasoning, and precise mask generation.
• Accessibility: Screen readers for visually impaired users need to describe specific regions of images on demand — "what's in the top-left corner?" — requiring region-captioning capabilities.

These applications all require bidirectional region understanding: going from language to mask (referring expression segmentation) and from mask to language (region captioning). They also require interleaved capabilities: the ability to process multiple masks in input, generate multiple masks in output, and reason across them. Current systems handle subsets of these requirements well, but the fragmented architecture makes building a unified system difficult.

Moreover, the field is moving toward scalable training recipes. The enormous success of LLMs stems partly from the simplicity of the training paradigm: next-token prediction on a massive text corpus, optionally followed by instruction tuning and RLHF. If pixel-wise capabilities require bespoke architectures and multi-objective training, they cannot ride this scaling curve — each new model generation requires re-engineering the pixel-wise components from scratch.

Where Prior Approaches Fall Short

The paper categorizes existing approaches into three families for both mask understanding and mask generation, and identifies specific failure modes in each. I'll walk through these systematically.

Prior Approaches to Mask Understanding (Mask as Input)

Approach 1: Visual prompts on images. Methods like Set-of-Mark [78] and VIP-LLaVA [5] directly overlay visual markers (circles, arrows, bounding boxes) on the input image to indicate the region of interest. The MLLM then processes this marked-up image normally.

Where this falls short: While intuitive and precise — the model literally sees the region — this approach alters the original image content. Overlays can occlude fine details, change the visual statistics that the model was trained on, and make it impossible to mark multiple overlapping regions simultaneously. For multi-round interactions where multiple regions accumulate, the image becomes progressively cluttered. More subtly, this approach couples the visual prompt design to the model's visual encoder — a prompt that works well for one MLLM architecture may not transfer to another.

Approach 2: Textual coordinates. Methods like QwenVL [2, 3, 67], InternVL [8], and others represent regions as textual bounding boxes or points in the instruction — e.g., "describe the region at (100, 200, 300, 400)."

Where this falls short: While maximally compatible with the natural language interface (no image modifications needed), this approach suffers from precision and ambiguity problems. The MLLM must learn to associate abstract coordinate numbers with specific image locations, which is fundamentally different from how humans refer to regions. A bounding box like "(100, 200, 300, 400)" provides a coarse rectangle that may include multiple objects or image regions, creating ambiguity about what exactly is being referred to. The authors note that this "poses significant challenges for MLLMs in precisely identifying which image regions the coordinates refer to" — the model may correctly reason about the general area but fail to isolate the exact intended region.

Approach 3: ROI feature extraction. Methods like Osprey [87], GLaMM [48], DAM [30], and GRASP [62] either encode the region mask into image features or extract ROI features from feature maps using the region mask. For example, they might use RoIAlign to pool features from the specific masked area and feed these into the LLM alongside text tokens.

Where this falls short: This is the most accurate approach but also the most architecturally invasive. These methods:

• Require specialized module designs — feature extractors, ROI poolers, region encoders — that must be integrated into the MLLM architecture
• Need integrated training — the region encoder typically cannot be trained independently and then plugged in; it must be co-trained with the LLM to ensure feature alignment
• Have complex pipelines — handling masked input requires routing through a different forward path than plain text or image input
• Suffer from limited generalizability and scalability — if you want to switch base models (e.g., from Vicuna to Llama to Qwen), you must redesign and retrain the integration layer

Prior Approaches to Mask Generation (Mask as Output)

Approach 1: Textually outputting coordinates or contours. Methods that output 2D points [16], bounding boxes [3, 16, 100], or polygonal contours [70, 92, 95] as text tokens.

Where this falls short: The fundamental issue is a mismatch between training objective and evaluation metric. These models are trained with textual cross-entropy loss (predicting the next coordinate token correctly), but evaluated with spatial metrics like IoU. A small textual error — predicting "234" instead of "235" for a coordinate — may be nearly imperceptible in cross-entropy terms but cause a large spatial displacement. Conversely, the model may learn to output coordinates that are linguistically valid (forming a sensible bounding box) but spatially inaccurate. This objective mismatch creates a ceiling on achievable performance: "there exists a critical gap between textual cross-entropy and continuous localization."

Approach 2: Special token + dedicated segmentation decoder. Methods like LISA [25], Sa2VA [85], OMG-LLaVA [93], HyperSeg [73], and PaDT [53] use a paradigm where the MLLM produces a special token whose hidden state is fed into a SAM-like segmentation decoder. The decoder then produces the 2D mask, and training involves both the language modeling loss and segmentation-specific losses (typically Dice loss and binary cross-entropy).

The paper describes this as the dominant SOTA approach: "aggregating segmentation information via special tokens [25, 53, 60, 73, 85, 93, 94] and decoding the tokens into 2D masks with a dedicated segmentation model [10, 24]."

Where this falls short: This approach has several critical limitations:

• Requires joint training of MLLM with an additional segmentation decoder — the decoder cannot be trained offline and plugged in; it must be co-trained with the LLM to align the continuous hidden states
• Relies on segmentation loss optimization — adding Dice loss and BCE alongside the cross-entropy objective complicates the training dynamics and may conflict with language learning objectives
• Continuous hidden states prevent RL — the special token's hidden state is a continuous vector, not a discrete token, so standard RL algorithms that work over token sequences cannot directly optimize mask quality through text-based rewards
• Architectural coupling — the decoder's design (SAM-based? Mask2Former-based? custom?) is baked into the MLLM architecture, making it hard to swap or upgrade independently

Approach 3: Treating masks as images for autoregressive generation. Methods like HiMTok [68], ARGenSeg [71], Text4Seg [26], and Alto [66] treat a region mask as a small image and autoregressively generate the mask pixels as image tokens (often using a VQ-VAE image tokenizer).

Where this falls short: While this approach can produce high-precision segmentation masks (since it's essentially modeling the full mask distribution autoregressively at the pixel/patch level), it has crippling efficiency problems:

• Enormous token budgets — a single 256×256 mask encoded through a VQ-VAE image tokenizer might require hundreds of tokens (e.g., a 16× downsampled latent grid of 16×16 = 256 tokens)
• Substantial computational costs — generating 100+ tokens per mask, especially in interleaved generation tasks with multiple masks, balloons inference time
• Mask reconstruction degradation with joint training — the paper explicitly notes (Appendix E, Fig. 7 and Section D) that some joint-training mask tokenizers suffer from mask reconstruction collapse — all masks degenerate into ellipses or blurry approximations because the LLM training objective interferes with the tokenizer's reconstruction fidelity

The paper's critique here is sharp: "other mask tokenizers [66, 68] require joint training with the MLLM, which ultimately leads to degraded mask reconstruction performance (with all masks reconstructed as ellipses)."

Reinforcement Learning in Pixel-Wise MLLMs: A Stalled Frontier

The paper also contextualizes the RL challenge specifically. While RL has become crucial for pushing LLM performance in math, coding, and reasoning (through GRPO, PPO, DPO, etc.), the pixel-wise MLLM subfield has been largely left behind:

• Mainstream VLM RL works [20, 29, 41, 61] achieve success through simple answer-matching rewards — checking if the model's textual answer matches the ground truth. This works because answers are text.
• Pixel-wise RL works [36, 37, 38, 52, 101] require converting the model's output into a mask first (via a decoder), then computing mask IoU or box IoU as the reward. This requires the intermediate step of mask de-tokenization, adding complexity.
• Some works [36, 37, 101] predict bounding boxes or points instead of masks, apply RL to those, and then use SAM to convert them to masks for reward computation — which "substantially increases the overall system complexity."

The authors' insight is that if masks were simply text tokens, then RL for mask generation would reduce to standard answer-matching — the same approach that has been so successful for text-based reasoning tasks. This is elegant: the domain gap between "mask RL" and "text RL" disappears entirely.

How This Paper Positions Itself

SAMTok positions itself not as another pixel-wise MLLM in the crowded field of LISA-derivatives and SAM-integrated architectures, but as a fundamental rethinking of the interface between masks and MLLMs. The paper's thesis, stated upfront, is:

"We propose SAMTok, a discrete mask tokenizer that tokenizes masks into textual special words (text tokens) and detokenizes these textual special words into masks, thereby transforming masks into a new language for MLLMs to learn from, similar to regular text data."

This framing is important because it changes the problem from "how do we build a better pixel-wise MLLM?" to "how do we build a better mask tokenizer?" — and then the MLLM problem becomes trivial. If SAMTok works well, then any base MLLM can acquire pixel-wise capabilities by simply adding vocabulary tokens and training on SAMTok-formatted data. No architecture changes. No new loss functions. No specialized decoders.

The paper explicitly distinguishes its approach along three dimensions that directly address the failures of prior work:

1. Decoupling tokenizer from MLLM. SAMTok is trained independently (on 209M masks) and then used as a frozen preprocessing/postprocessing tool. This means:

• The tokenizer's reconstruction quality doesn't degrade during MLLM training (unlike joint-training approaches)
• The tokenizer can be swapped or upgraded without retraining the MLLM
• Different MLLMs can share the same tokenizer (the paper demonstrates this with Qwen2.5-VL, Qwen3-VL, and PerceptionLM in Table 11)

2. Extreme compression: 2 tokens per mask. By using residual vector quantization (RQ) with a compact codebook (256×2 = 512 total codes, but only 2 tokens per mask), SAMTok achieves a representation that is orders of magnitude more compact than RLE-text (~10–100 tokens) or autoregressive image generation (~100+ tokens). This makes it practical for interleaved generation tasks with many masks and for long-context scenarios.

3. Discrete representation enabling text-only RL. Because mask tokens are literal text tokens in the LLM's vocabulary, RL can use a purely textual reward function — count how many predicted mask tokens match ground-truth mask tokens — without needing to de-tokenize masks, compute IoU, or involve any external segmentation model during RL training. The paper frames this as a first:

"To the best of our knowledge, this is the first successful attempt to optimize mask generation performance using text-only reward signals."

The paper also draws an implicit analogy to the NLP community's evolution from task-specific architectures to unified language models. In 2017–2018, NLP tasks (translation, summarization, QA, sentiment analysis) each had specialized architectures. The transformer + pretraining revolution showed that all these tasks could be unified under text-to-text frameworks (T5, GPT). SAMTok makes the analogous move for pixel-wise vision tasks: if masks can be losslessly converted to text tokens, then all these tasks (referring segmentation, region captioning, grounded conversation, interactive segmentation) become instances of text-to-text generation — and the entire MLLM toolkit (SFT, RLHF, GRPO, chain-of-thought) becomes available without modification.

This positioning is simultaneously ambitious and humble: ambitious in claiming that a 2-token representation can replace entire segmentation decoders and region encoders; humble in acknowledging that SAMTok is "just" a tokenizer — it doesn't try to also be the MLLM, the RL algorithm, or the training framework. This separation of concerns is the key architectural insight.

3. Technical Approach

3.1 Reader Orientation

The paper builds SAMTok, a frozen mask tokenizer that compresses any 2D region mask into exactly two discrete text tokens and can faithfully reconstruct the mask from those tokens alone, and then demonstrates that by treating these mask tokens as a new "language" — simply adding them to the MLLM's vocabulary and training with standard next-token prediction — any base MLLM can acquire pixel-wise understanding and generation capabilities without a single architectural modification, specialized loss function, or task-specific decoder. The problem it solves is the fragmentation of pixel-wise MLLM design: current approaches require separate mechanisms for mask input (region feature extractors) and mask output (segmentation decoders), cannot straightforwardly apply reinforcement learning because they use continuous embeddings, and couple the mask handling so tightly to the MLLM backbone that upgrading either component forces retraining both. SAMTok's shape of the solution is to move all the mask-specific complexity into a standalone tokenizer trained once on 209M masks, converting a messy 2D spatial representation problem into a clean text-to-text problem that the entire LLM ecosystem already knows how to handle.

3.2 Big-Picture Architecture (Diagram in Words)

The system has three major components, arranged in a pipeline that is trained independently at each stage:

1. SAMTok Tokenizer — a frozen encoder-quantizer-decoder based on SAM2 that converts any (image, mask) pair into two discrete codebook indices (the "mask words"), and conversely converts two discrete codebook indices back into a 2D mask. This is trained once on 209M masks and then never updated.

2. Base MLLM (e.g., QwenVL series) — an off-the-shelf multimodal LLM whose vocabulary is extended with 512 new special tokens (one per codebook entry plus start/end markers). No other architectural changes are made. The MLLM is never told about segmentation decoders, Dice loss, or ROI pooling; it only sees text tokens, some of which happen to represent masks.

3. Reinforcement Learning Module (GRPO) — a standard RL training loop that operates purely on text: it samples multiple responses from the MLLM, checks which ones contain the correct mask tokens (character-level string matching), assigns rewards based on precision/recall of mask tokens, and updates the model to increase the probability of correct mask-word sequences. No mask decoding occurs during RL training.

Information flows as follows: Training the tokenizer — a dataset of (image, mask) pairs flows into SAMTok's encoder, which produces a continuous embedding; this is residual-quantized into two discrete codes; the decoder reconstructs the mask from these codes; reconstruction + commitment losses update the quantizer and decoder. Training the MLLM — all mask data (grounding, captioning, VQA, interactive segmentation) is preprocessed through the frozen SAMTok encoder, converting every mask into two text tokens; these tokenized dialogs become standard vision-language training data; the MLLM is fine-tuned with pure next-token prediction loss, learning to output mask tokens just like any other vocabulary word. RL training — the MLLM generates multiple candidate responses for each query; the textual answer-matching reward counts how many predicted mask tokens match ground-truth mask tokens; the GRPO algorithm updates the model to favor higher-reward responses.

3.3 Roadmap for the Deep Dive

• First, the SAMTok encoder — how a 2D mask gets converted into a single continuous embedding vector using SAM2's prompt encoder and mask decoder, and why the architecture reuses SAM components this way.
• Second, the residual vector quantizer — how the continuous embedding is discretized into exactly two codebook indices, why residual quantization (RQ) is chosen over VQ or FSQ, and the critical tradeoff between codebook size, number of quantization steps, reconstruction fidelity, and the MLLM's ability to learn the token space.
• Third, the SAMTok decoder — how the two discrete codes are converted back into a 2D mask, the design choice to treat them as "sparse prompt embeddings" for a full SAM2 model, and why decoupling the decoder from MLLM training preserves reconstruction quality.
• Fourth, the SAMTok training procedure — the reconstruction loss (cross-entropy + Dice), the commitment loss, the 209M-mask training dataset, and the freezing strategy (which SAM2 parts are frozen vs. trained).
• Fifth, the unified mask-token interface for MLLMs — how mask tokens are integrated into the MLLM vocabulary, the special token formatting (<|mt_0000|> through <|mt_0511|> plus <|mt_start|> and <|mt_end|>), how mask understanding tasks (mask → text) and mask generation tasks (text → mask) become instances of next-token prediction, and how diverse tasks like GCG, GRES, REC, region captioning, and interactive segmentation are all unified under a single text-to-text formulation.
• Sixth, the reinforcement learning approach for mask generation — the textual answer-matching reward function (Equation 5), why discrete mask tokens enable RL where continuous embeddings cannot, the GRPO training setup, and the cold-start chain-of-thought data used to bootstrap RL training.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems and interface paper whose core idea is that masks can be represented as an extremely compact (two-token) discrete code that is learned once by a standalone tokenizer and then treated as indistinguishable from natural language by any downstream MLLM. The technical contribution is the tokenizer architecture and training, the integration protocol with MLLMs, and the demonstration that text-only RL rewards work for mask generation when masks are text tokens.

SAMTok Encoder: From 2D Mask to Continuous Embedding

The SAMTok encoder $f_{\text{enc}}$ takes an image $I$ and a binary region mask $M$ (specifying which pixels belong to the region of interest) and produces a single $d$-dimensional continuous embedding vector $z$ that compactly captures the spatial extent, shape, and appearance context of that region. The key architectural insight is that the encoder is instantiated using components from SAM (Segment Anything Model) — specifically SAM2 [49] — so that the rich visual understanding and segmentation priors learned by SAM during its massive pretraining are directly leveraged for the tokenization task.

The encoder's forward pass proceeds in three steps, mirroring SAM's own interactive segmentation pipeline but with a crucial difference in what prompts are provided:

Step 1: Image feature extraction. The SAM image backbone $f_{\text{img}}$ processes the input image $I$ to produce a spatial feature map. In SAM, this backbone is typically a Vision Transformer (ViT) that outputs features at some downsampled resolution (e.g., 16× or 32× smaller than the input image). These features encode rich semantic and geometric information about the image contents — where edges are, what textures exist, how objects relate spatially. The authors freeze the SAM image backbone during SAMTok training (Section B), meaning these features are computed exactly as they would be in the original SAM model, preserving the pretrained visual representations.

Step 2: Mask-to-dense-prompt encoding. The SAM prompt encoder $f_{\text{prm}}$ takes the binary mask $M$ and converts it into dense prompt embeddings that share the same spatial resolution as the image features from Step 1. In standard SAM, this prompt encoder can accept various types of spatial prompts: points, bounding boxes, or coarse masks. When given a mask, SAM typically downsamples it to match the feature resolution and applies a learned embedding to create a dense conditioning signal. The authors leverage this existing mechanism: the mask itself becomes the "prompt" that tells SAM which region to attend to.

Why is this clever? In standard interactive segmentation with SAM, the user provides sparse prompts (clicks, boxes) and SAM predicts a mask. Here, the role is inverted: the mask is the input, and SAM's prompt encoder is repurposed to encode the mask's spatial information into the feature space. This means the encoder inherits SAM's ability to understand what constitutes a coherent region — boundaries, connectivity, objectness — when encoding the mask.

Step 3: Feature fusion and mask embedding extraction. The dense prompt embeddings from Step 2 are added to the image features from Step 1 (element-wise addition at each spatial location). This enriched feature map — which now carries both image content and region-of-interest information — is fed into the SAM mask decoder $f_{\text{msk}}$.

Now comes the critical design choice. In a standard SAM forward pass, the mask decoder would interact with the enriched features and a set of learnable mask tokens (output tokens) via cross-attention and self-attention, ultimately producing a segmentation mask through its mask prediction head. However, the SAMTok encoder removes the final mask prediction head from $f_{\text{msk}}$. Instead, it only runs the mask decoder up to the point where it produces a mask embedding — the internal representation that would normally be fed into the prediction head. Specifically, the decoder has a pre-initialized mask embedding (a learnable query vector, analogous to object queries in DETR or mask tokens in Mask2Former) that interacts with the image features via cross-attention and self-attention layers. The output of this interaction is the $d$-dimensional continuous mask embedding $z$:

$z = f_{\text{enc}}(I, M) = f_{\text{msk}}(f_{\text{img}}(I), f_{\text{prm}}(M)) \in \mathbb{R}^d$

where $f_{\text{img}}(I)$ is the image feature map from the frozen SAM image backbone, $f_{\text{prm}}(M)$ is the dense prompt embedding derived from the mask, $f_{\text{msk}}(\cdot, \cdot)$ is the SAM mask decoder (minus prediction head) that fuses these two inputs through the pre-initialized mask embedding, and $z$ is the resulting $d$-dimensional continuous vector representing the entire region.

What this computes: For a given image and a highlighted region, the encoder produces a single fixed-size vector that captures everything SAM "knows" about what makes that region a coherent segmentable unit — its shape, its boundary, its relationship to surrounding image content, its semantic category (implicitly). This is not a box coordinate or a contour; it's a learned continuous representation that encodes the mask in SAM's internal feature space.

Why this form: Three reasons. (1) By reusing SAM's pretrained image backbone and prompt encoder, the encoder inherits strong visual priors and segmentation capabilities without needing to learn them from scratch — the image features already know where object boundaries are, and the prompt encoder already knows how to map spatial prompts to feature-space conditioning. (2) By using the mask decoder to produce an embedding rather than a mask, the encoder creates a bottleneck representation that is both compact (a single vector) and information-rich (it went through SAM's full attention-based refinement), which is exactly what a tokenizer needs — the continuous embedding must capture enough detail for reconstruction while being low-dimensional enough for efficient quantization. (3) By removing the mask prediction head, the encoder deliberately separates representation from reconstruction — the quantizer and decoder will handle converting this embedding to discrete codes and back to a mask, creating a modular design where each component can be optimized independently.

The dimensionality $d$ is not explicitly stated in the paper but is determined by the mask decoder's output embedding size in SAM2 (Large variant). In a typical SAM architecture, this is on the order of 256 dimensions — the same size as the transformer's hidden dimension.

Residual Vector Quantizer: From Continuous Embedding to Two Discrete Codes

The quantizer's job is to take the continuous mask embedding $z$ and convert it into a small number of discrete tokens that can be directly inserted into an MLLM's vocabulary. The core challenge is the rate–distortion tradeoff: fewer tokens and smaller codebooks make it easier for the MLLM to learn to predict the correct codes (the token sequence is short, the search space is small), but they also limit how accurately the mask can be reconstructed (the discrete representation has limited capacity). The authors choose residual quantization (RQ) [27] after experimentally comparing it against vector quantization (VQ) [58] and finite scalar quantization (FSQ) [44] in Table 12.

How residual quantization works. Standard VQ maps the continuous vector to its single nearest neighbor in a codebook — pick the codebook entry that minimizes Euclidean distance to $z$, and that one entry (represented as an integer index) is the discrete representation. The quality of this approximation depends on the codebook size: a larger codebook means more possible representatives, so less approximation error, but also a larger vocabulary for the MLLM to learn.

RQ generalizes this by allowing multiple quantization steps, each using the same or a different codebook, applied to the residual error from the previous step. The intuition: if the first codebook entry $e_1$ approximates $z$ but leaves some error $r_1 = z - e_1$, a second codebook can approximate that error $r_1$ with $e_2$, and the final discrete representation is the concatenation $[e_1, e_2]$. The reconstruction is then $e_1 + e_2$, which can represent a much richer set of vectors than either codebook alone because the sum of two codebook vectors spans a combinatorially larger space.

Formally, the paper defines the two-step RQ procedure as:

$\begin{cases} e_1 = \arg\min_{e \in \mathcal{C}} \|z - e\|_2^2 \\ r_1 = z - e_1 \\ e_2 = \arg\min_{e \in \mathcal{C}} \|r_1 - e\|_2^2 \\ q = [e_1, e_2] \end{cases}$

where $\mathcal{C}$ is the codebook (a learned set of embedding vectors, each of dimension $d$), $e_1$ is the codebook vector closest to $z$ (first-stage quantization), $r_1$ is the residual — what the first vector failed to capture, $e_2$ is the codebook vector closest to the residual $r_1$ (second-stage quantization), and $q = [e_1, e_2]$ is the final discrete representation consisting of two embedding vectors (or equivalently, two integer indices into the codebook).

What this computes: Starting from the continuous mask embedding $z$, the quantizer greedily finds a two-vector approximation: the first vector $e_1$ captures the dominant structure (the "coarse" aspects of the mask), and the second vector $e_2$ captures the fine details that $e_1$ missed. The operation produces two indices into the codebook $\mathcal{C}$, which become the two "mask words" that the MLLM will see.

Why this form: The paper's ablation in Table 12 provides the empirical justification. Using standard VQ with a small codebook (1024 entries) severely limits reconstruction accuracy (r-Acc = 0.50) because the codebook cannot represent the diversity of region masks appearing across all images — a mask representation that's off by one codebook entry can cause the reconstructed mask to shift dramatically. Increasing VQ codebook size to 65,536 improves reconstruction (r-Acc = 0.66) but creates an enormous vocabulary for the MLLM — the model must learn to select among 65,536 possible tokens at each position, which is a harder learning problem and increases the risk of the MLLM never learning to use the rare codes. FSQ with a large codebook achieves better codebook utilization (r-Acc = 0.69, g-Acc = 78.1), meaning more of the 65,536 entries are actually used during training.

RQ with two codebooks of 1024 entries each (denoted "1024×2") achieves better reconstruction accuracy than FSQ with a 64× larger total codebook (r-Acc = 0.70 vs. 0.69), while keeping each individual codebook small (1024 entries) and requiring only two tokens. This is the key insight: two small codebooks, applied sequentially to the residual, can represent a combinatorially large space ($1024 \times 1024$ possible combinations) while keeping the MLLM's vocabulary manageable (1024 + 1024 = 2048 tokens, or 256 + 256 = 512 tokens in the final default configuration). The MLLM only needs to learn to pick one of 256 options at each of two positions, rather than one of 65,536 options at a single position.

The paper further ablates codebook size and quantization steps in Table 13:

• 1024×4 (two codebooks of 1024, four quantization steps) improves reconstruction (r-Acc = 0.75) because more steps = less residual error, but hurts generation accuracy (g-Acc = 77.3 vs. 78.3 for 1024×2) because the exponentially larger search space ($1024^4$ possible combinations) makes it harder for the MLLM to learn to generate the correct token sequences. The authors explicitly note: "the exponentially expanded search space ($1024^4$ vs. $1024^2$) makes it more difficult for the MLLM to learn to generate mask tokens effectively."
• 256×4 (codebook 256, four steps): similar r-Acc to 1024×2 (0.72 vs. 0.70), but lower g-Acc (77.3 vs. 78.3).
• 256×2 (the default): r-Acc = 0.70, g-Acc = 77.6 — a sweet spot where reconstruction is adequate, the vocabulary is compact (256 + 256 = 512 special tokens), and the MLLM can learn the generation task effectively.
• 512×2: slight improvements (r-Acc = 0.71, g-Acc = 77.8), but the authors choose 256×2 as default for its compactness.

Codebook details. The default configuration uses two non-shared codebooks, each of size 256, meaning there are $256 \times 256 = 65,536$ possible mask representations total. The codebooks are randomly initialized and trained jointly with the encoder/decoder. The non-shared design means the first and second quantization steps can specialize — the first codebook might learn to encode coarse mask properties (location, approximate size, broad shape category), while the second codebook learns to encode residual details (fine boundary adjustments, protrusions, holes). This is more expressive than a shared codebook where both steps must use the same set of vectors.

A critical subtlety about the two codebooks. The MLLM vocabulary must accommodate both codebooks. The paper introduces mask tokens formatted as <|mt_0000|> through <|mt_0511|>, where <|mt_0000|> through <|mt_0255|> correspond to the first-level codebook (256 entries) and <|mt_0256|> through <|mt_0511|> correspond to the second-level codebook (256 entries). Additionally, two special tokens <|mt_start|> and <|mt_end|> mark the beginning and end of mask token sequences. This means a single mask is represented in the MLLM's input/output as three tokens: <|mt_start|><|mt_XXXX|><|mt_YYYY|><|mt_end|>, where XXXX selects from the first codebook and YYYY selects from the second. The start/end tokens serve two purposes: they disambiguate mask token sequences from regular text (the MLLM knows that anything between start and end tokens represents a mask), and they can be used to extract mask tokens during RL reward computation by simple string matching between <|mt_start|> and <|mt_end|>.

SAMTok Decoder: From Two Discrete Codes Back to a 2D Mask

The decoder $f_{\text{dec}}$ takes the original image $I$ and the two discrete mask embeddings $[e_1, e_2]$ (retrieved from the codebook using the two integer indices) and reconstructs the 2D binary mask $\hat{M}$. Like the encoder, the decoder is instantiated using a full SAM model (SAM 2.1 Large) — including the image backbone, prompt encoder, mask decoder, and critically, the mask prediction head that was removed from the encoder.

The decoder's forward pass proceeds as follows:

Step 1: Image feature extraction. Exactly as in the encoder, the SAM image backbone $f_{\text{img}}$ processes the original image $I$ to produce a spatial feature map. The paper uses the same frozen SAM image backbone weights for both encoding and decoding, meaning the image features are identical regardless of whether the image is being used for tokenization or de-tokenization. This consistency is important: the decoder sees the same visual context that the encoder saw, preventing any feature-space mismatch.

Step 2: Sparse prompt encoding from discrete codes. This is where the decoder diverges from the encoder. The encoder used the mask itself as a dense prompt (the mask was spatially downsampled and embedded to create a prompt at every spatial location). The decoder, however, only has the two discrete embedding vectors $e_1$ and $e_2$ — it does not have the original mask. The authors treat these two vectors as "special language prompt embeddings" and feed them into the SAM prompt encoder $f_{\text{prm}}$.

Inside the prompt encoder, the two discrete embeddings are directly summed to form a single sparse prompt embedding. This summation is the key operation:

$\text{prompt\_embedding} = e_1 + e_2$

Why summation rather than concatenation or separate processing? The authors don't elaborate, but the logic follows from the residual quantization structure: since the encoder's quantizer produced $e_1 + e_2$ as the reconstruction of $z$ (modulo quantization error), summing them in the decoder recovers (an approximation of) the original continuous mask embedding. This means the decoder's prompt encoder receives essentially the same information that the encoder's mask decoder produced — the continuous mask embedding — just delivered through a different mechanism (summed discrete codes rather than direct output of the mask decoder). This design choice preserves the information flow: encoder produces embedding → quantizer discretizes it → decoder reconstructs it by summing and treats it as a prompt.

A subtle point about the prompt embedding's sparsity: in the encoder, the dense prompt embedding was added to the image features at every spatial location (dense conditioning). In the decoder, the sparse prompt embedding is not added to image features directly. Instead, it's passed into the mask decoder as a separate input — analogous to how SAM processes a point prompt (which is also a single embedding vector rather than a dense spatial map). The mask decoder then handles the interaction between this sparse prompt and the dense image features through its attention mechanisms.

Step 3: Mask reconstruction through the mask decoder. The mask decoder $f_{\text{msk}}$ (this time including the mask prediction head) takes three inputs: the image features from Step 1, the sparse prompt embedding from Step 2, and a pre-initialized mask embedding (the same learnable query token used in the encoder, or a separate one — the paper is not explicit on whether they're shared).

The key computational dynamic inside the mask decoder is:

• Self-attention between the pre-initialized mask embedding and the prompt embedding: This allows the mask token (the query) to attend to the prompt (the sum of $e_1 + e_2$), effectively "reading" the information encoded in the discrete codes. This self-attention mechanism is how the mask decoder recovers the features of the continuous mask embedding $z$ — the summed discrete embeddings interact with the learnable query token, and through the learned attention weights, the query token reconstructs something functionally equivalent to the original $z$ that the encoder produced.

• Cross-attention between the refined query and the image features: The query token (now enriched with mask information from the prompt embedding) attends to the dense image features. This is the standard SAM mechanism: the query token "looks at" the image and decides which spatial locations are relevant given the prompt. Because the prompt embedding carries the decoded mask information, the query token attends to the specific image region that corresponds to the original mask.

• Mask prediction head: The final representation of the query token (after self-attention and cross-attention layers) is passed through a linear projection and upsampling to produce the binary mask $\hat{M}$. This is the standard SAM prediction head — typically a small MLP that outputs a logit per pixel, followed by a sigmoid activation to produce a probability map, then thresholded (usually at 0.5) to get the binary mask.

The full decoder operation is expressed as:

$\hat{M} = f_{\text{dec}}(I, [e_1, e_2]) = f_{\text{msk}}(f_{\text{img}}(I), f_{\text{prm}}([e_1, e_2]))$

where $f_{\text{prm}}([e_1, e_2])$ produces the sparse prompt embedding by summing the two discrete codebook vectors, $f_{\text{img}}(I)$ is the frozen image features, and $f_{\text{msk}}$ is the full SAM mask decoder including the prediction head that converts the refined query token into a binary mask $\hat{M}$.

What this computes: Given an original image and two codebook indices representing a mask, the decoder reconstructs that mask at the original image resolution. It does so by first recovering the continuous mask embedding (via summation of the two codebook vectors), then using SAM's standard interactive segmentation machinery — treating the recovered embedding as if it were a user-provided prompt — to delineate the specific region in the image.

Why this form: Three architectural motivations. (1) By treating the discrete codes as prompts for a full SAM model, the decoder leverages SAM's powerful segmentation capability and its pretrained understanding of object boundaries, semantic coherence, and visual context. The decoder doesn't need to learn "how to segment" from scratch — it inherits this from SAM and only needs to learn the mapping from discrete codes to a prompt embedding that triggers the correct segmentation. (2) By summing $e_1 + e_2$ and passing the result as a sparse prompt, the decoder creates a direct analog to the encoder: the encoder produced $z$ and quantized it to $e_1 + e_2$ (approximately), and the decoder starts from $e_1 + e_2$ and reconstructs. This symmetry ensures that no information is lost in the prompt encoding step (beyond quantization error). If the decoder used a more complex prompt construction (e.g., cross-attention between $e_1$ and $e_2$), the architecturally simpler summation approximation might still hold, but unnecessary complexity would be added. (3) Critically, the decoder is fully decoupled from the MLLM. Once trained, the SAMTok decoder is never updated during MLLM fine-tuning or RL. This means the mask reconstruction quality is invariant to whatever happens during MLLM training — a significant advantage over joint-training approaches where mask quality can degrade.

Training of SAMTok: Reconstruction and Commitment Losses

SAMTok is trained with a single objective: given an image and a mask, can the encoder-quantizer-decoder pipeline reconstruct the mask accurately? The training data consists of 209 million masks drawn from diverse open-source segmentation datasets [13-15, 24, 28, 46, 49, 77, 80, 81, 98], spanning indoor scenes, outdoor environments, website UIs, part-level annotations, object-level annotations, entity-level annotations, and semantic-level annotations. This diversity is crucial: the tokenizer must learn to represent masks across vastly different domains (from COCO objects to ADE20k semantic regions to Cityscapes street scenes to website buttons) so that it generalizes to any mask an MLLM might encounter during deployment.

The training objective follows the standard RQ-VAE formulation [27], combining a reconstruction loss that measures mask quality and a commitment loss that encourages the encoder to produce embeddings close to codebook entries:

$\begin{cases} \mathcal{L} = \mathcal{L}_{\text{recon}} + \lambda \mathcal{L}_{\text{commit}} \\ \mathcal{L}_{\text{recon}} = \mathcal{L}_{\text{CE}}(M, \hat{M}) + \mathcal{L}_{\text{DICE}}(M, \hat{M}) \\ \mathcal{L}_{\text{commit}} = \|z - \text{sg}(e_1)\|_2^2 + \|r_1 - \text{sg}(e_2)\|_2^2 \end{cases}$

where $\mathcal{L}_{\text{recon}}$ is the total reconstruction loss, $\mathcal{L}_{\text{CE}}$ is the pixel-wise binary cross-entropy between the ground-truth mask $M$ and the reconstructed mask $\hat{M}$, $\mathcal{L}_{\text{DICE}}$ is the Dice loss (1 minus the Dice coefficient, which measures overlap quality), $\mathcal{L}_{\text{commit}}$ is the commitment loss, $\text{sg}(\cdot)$ is the stop-gradient operator, $\lambda$ controls the tradeoff between reconstruction quality and codebook commitment, $z$ is the continuous mask embedding from the encoder, $e_1$ is the first-stage quantized vector, and $r_1 = z - e_1$ is the first-stage residual.

What this computes, term by term:

• Cross-entropy loss $\mathcal{L}_{\text{CE}}(M, \hat{M})$: For each pixel in the image, the decoder outputs a probability $\hat{p}_{ij}$ that pixel $(i, j)$ belongs to the mask. The cross-entropy loss penalizes deviations between this probability and the ground-truth binary label $M_{ij} \in \{0, 1\}$. This is the standard pixel-wise classification loss — it treats mask reconstruction as per-pixel binary classification, which works well for capturing fine spatial details but can be imbalanced (most pixels in an image are typically background).

• Dice loss $\mathcal{L}_{\text{DICE}}(M, \hat{M})$: The Dice coefficient is $2|M \cap \hat{M}| / (|M| + |\hat{M}|)$ — twice the intersection area divided by the sum of areas. The Dice loss is $1 - \text{Dice}$. Unlike cross-entropy, the Dice loss is region-based: it directly optimizes overlap quality and is less sensitive to class imbalance because it normalizes by region sizes. Small masks (which have few positive pixels, making cross-entropy dominated by easy background negatives) get appropriate weight under Dice loss. This is why most segmentation models use it alongside cross-entropy.

• Commitment loss $\mathcal{L}_{\text{commit}}$: This is the critical term for vector quantization training. The stop-gradient operator $\text{sg}(\cdot)$ means that when computing $\|z - \text{sg}(e_1)\|_2^2$, gradients flow only through $z$ (pulling the encoder output toward the codebook vector), not through $e_1$ (which is updated by a different mechanism — typically exponential moving averages or the straight-through estimator). Similarly, $\|r_1 - \text{sg}(e_2)\|_2^2$ encourages the first-stage residual to be close to the second codebook vector. Without this loss, the encoder could produce embeddings arbitrarily far from any codebook entry, making quantization essentially random. With it, the encoder learns to produce embeddings that are well-approximated by the codebook.

Why this form: The combination of cross-entropy and Dice loss for reconstruction is standard practice in segmentation because they have complementary properties: cross-entropy provides per-pixel gradient signal (good for learning fine boundaries), while Dice provides region-level signal (good for handling size variation and avoiding the "ignore small objects" bias). The commitment loss with the stop-gradient operator is the standard VQ-VAE training technique — it's the minimum mechanism needed to train discrete representations with gradient-based optimizers, because the $\arg\min$ operation in quantization is non-differentiable.

Training configuration (from Appendix B):

• Initialization: The encoder and decoder are initialized with SAM 2.1 Large pretrained weights; the codebooks are randomly initialized.
• Frozen components: The SAM image encoder $f_{\text{img}}$ and SAM prompt encoder $f_{\text{prm}}$ are frozen during SAMTok training; only the SAM decoder $f_{\text{msk}}$ (within both encoder and decoder) and the codebook are trainable. This preserves SAM's general visual understanding and prompt encoding capabilities while allowing the decoder to adapt to the mask reconstruction task.
• Optimizer: AdamW [40]
• Learning rate: $4 \times 10^{-5}$ with cosine decay schedule [39]
• Global batch size: 1024
• Framework: Xtuner [12]
• Hardware: NVIDIA A100 GPUs (80 GB)

The choice to freeze the image encoder and prompt encoder is pragmatic: SAM2's pretrained weights already encode rich visual features, and fine-tuning them on the mask reconstruction task might cause catastrophic forgetting of the general visual knowledge that makes the tokenizer robust across domains. Freezing also reduces memory requirements and training time, since only the decoder and codebook parameters receive gradients.

Unified Mask-Token Interface for MLLMs

This section describes how SAMTok, once trained and frozen, is integrated into an MLLM to enable mask understanding and generation through standard next-token prediction. The key conceptual move is to treat the mask tokens produced by SAMTok as a new language — analogous to adding a new natural language (like French or Japanese) to a multilingual model's vocabulary, but for a visual "language" of region masks.

Vocabulary extension. The base MLLM (QwenVL series by default) has its tokenizer's vocabulary extended with 514 new special tokens:

<|mt_0000|> through <|mt_0255|>   →  256 entries, corresponding to the first-level codebook
<|mt_0256|> through <|mt_0511|>   →  256 entries, corresponding to the second-level codebook
<|mt_start|>                       →  marks the beginning of a mask token sequence
<|mt_end|>                         →  marks the end of a mask token sequence

The embedding vectors for these new tokens are randomly initialized using the mean and variance statistics of the original token embeddings. This is an important initialization detail: rather than starting from zeros or a random Gaussian, the new embeddings are centered at the existing vocabulary's embedding distribution. This helps gradient flow during the early stages of fine-tuning — the mask token embeddings start in a similar regime to the text tokens they'll interact with, rather than being outliers that the attention mechanism would ignore.

No other architectural changes are made to the base MLLM. The visual encoder, the projection layer, the LLM backbone — all remain exactly as in the original QwenVL architecture. The only difference is that the vocabulary now includes tokens that happen to represent masks.

Preprocessing for mask understanding tasks. For any task where a mask is part of the input (region captioning, region VQA, interactive segmentation where previous masks inform the current query), SAMTok's encoder is used to tokenize each mask into two codebook indices. For example, if a region of interest in an image is represented by mask $M$, SAMTok encodes it into indices $(i_1, i_2)$, which are then formatted as the text string:

<|mt_start|><|mt_i1|><|mt_i2|><|mt_end|>

This string is inserted into the MLLM's text prompt at the position where the mask reference is needed. For a region captioning task, the prompt might look like:

"<image> Describe the region <|mt_start|><|mt_0011|><|mt_0347|><|mt_end|> in detail."

where <image> is the standard image placeholder token used by QwenVL. The MLLM then processes this as a normal multimodal input — the image goes through the visual encoder, the text tokens (including mask tokens) go through the LLM's embedding layer, and the model applies its standard cross-attention to fuse visual and textual information. The MLLM learns during fine-tuning that <|mt_0011|><|mt_0347|> refers to a specific region in the image — the one that SAMTok would reconstruct from those codes — because the training data consistently pairs these token sequences with captions that describe the corresponding regions.

Preprocessing for mask generation tasks. For any task where the MLLM must output masks (referring segmentation, grounded conversation generation, panoptic scene graph generation), the ground-truth masks in the training data are tokenized with SAMTok, and the target text sequence includes the mask tokens. For example, in referring expression segmentation (given a text query like "the red cup on the left," produce the mask for that cup), the training example would be:

Input:  "<image> Segment the red cup on the left."
Output: "<|mt_start|><|mt_0042|><|mt_0187|><|mt_end|>"

During training, the MLLM learns to output these token sequences through standard next-token prediction (cross-entropy loss on each output token, including the mask tokens). There is no segmentation loss, no Dice loss, no mask decoder training — the MLLM's training is indistinguishable from any other text generation task. The MLLM never "sees" the 2D mask; it only sees the two integer indices.

During inference for mask generation tasks, the MLLM autoregressively generates text, which may include mask token sequences. To visualize the result, these sequences are postprocessed: the <|mt_start|> and <|mt_end|> tokens are stripped, the two codebook indices are extracted, the corresponding embedding vectors $e_1$ and $e_2$ are retrieved from the codebook, and SAMTok's decoder converts them back into a 2D mask (using the original image as context). This postprocessing step is entirely external to the MLLM — the MLLM doesn't know or care that masks are being reconstructed.

Unified co-training across diverse tasks. The paper's crucial claim is that this preprocessing pipeline enables all mask-related tasks to be unified under a single text-to-text formulation. The authors explicitly describe their data preprocessing:

"After SAMTok tokenizes all region masks into words, all mask-related tasks can be preprocessed into purely textual corpora, including mask-to-text, text-to-mask, interleaved text–mask generation, and interactive tasks (where masks serve as both inputs and outputs)."

Concrete examples of how diverse tasks are reformulated:

• Region captioning (mask → text): Input contains the image and mask tokens referring to a specific region; target output is the textual description of that region. Both input mask tokens and output text are trained with next-token prediction (the mask tokens in the input provide context; only the output text tokens contribute to the loss).

• Referring expression segmentation (text → mask): Input is the image and a natural language query ("the red cup"); target output is the mask token sequence. The loss is computed only on the output mask tokens.

• Grounded conversation generation (interleaved text + masks): Input is the image and an instruction ("describe this scene and localize mentioned objects"); target output interleaves text phrases and mask token sequences (e.g., "A <|mt_start|><|mt_0042|><|mt_0187|><|mt_end|> red cup sits next to a <|mt_start|><|mt_0023|><|mt_0156|><|mt_end|> blue bowl."). Both text and mask token predictions contribute to the loss.

• Multi-round interactive segmentation (masks as both input and output): In round $t$, the input includes mask tokens for all previously segmented objects; the model must parse the new instruction (e.g., "now segment the object behind the first one"), reason about spatial relationships, and output new mask tokens. Both input mask tokens (providing context) and output mask tokens (the target) are handled by the same next-token prediction mechanism.

Because all tasks are now pure text generation, the MLLM can be fine-tuned on a mixture of all these tasks simultaneously with a single loss function — next-token prediction — without any task-specific loss weighting, gradient accumulation tricks, or architectural switches. The paper calls this "co-training under any multimodal training framework, all these tasks can be co-trained using the standard next-token prediction loss, without introducing any customized loss functions or architectural modifications." This is a significant simplification compared to prior work, where segmentation tasks required Dice loss + cross-entropy, VQA tasks required cross-entropy, and these had to be balanced with weighting hyperparameters.

Training configuration for MLLM fine-tuning (from Appendix B):

• Frozen components: The MLLM's image encoder is frozen; the projection layer (connecting visual features to the LLM's embedding space) and the LLM backbone are fine-tuned.
• Embedding initialization: The 514 new mask token embeddings are randomly initialized with mean/variance matching of existing token embeddings.
• Optimizer: AdamW [40]
• Learning rate: $2 \times 10^{-5}$ with cosine decay schedule [39]
• Global batch size: 256
• Supervised fine-tuning data: ~5M samples spanning mask generation (grounding, referring segmentation, GCG, instance segmentation, scene parsing), region understanding (region captioning, region VQA), and collaborative tasks (multi-round interactive segmentation). All samples are pre-tokenized with SAMTok before training.

Reinforcement Learning for Mask Generation

The discrete nature of SAMTok's mask tokens unlocks a capability that has been largely inaccessible to prior pixel-wise MLLMs: reinforcement learning with purely text-based reward signals. The authors describe this as both a technical contribution and a conceptual demonstration — proving that the entire RL toolkit developed for text-based reasoning (GRPO, PPO, DPO, etc.) transfers to mask generation when masks are represented as text.

Why continuous embeddings block RL. In approaches like LISA [25], Sa2VA [85], or OMG-LLaVA [93], the MLLM outputs a special token whose hidden state (a continuous vector) is passed to a segmentation decoder. To apply RL, one would need to: (1) decode this continuous vector into a mask (requiring a forward pass through the mask decoder), (2) compute a reward based on mask quality (e.g., IoU with ground truth), (3) backpropagate through the mask decoder into the MLLM. This is possible in principle but practically challenging: the reward function operates in continuous mask space, not discrete token space, so standard RL algorithms that compare text outputs (like answer-matching rewards) cannot be used. Some works [36, 37, 101] have attempted workarounds by predicting bounding boxes or points as text, applying RL to those coordinate sequences, and using SAM to convert them to masks for reward computation — but this adds substantial complexity and still doesn't enable direct mask-level RL.

SAMTok sidesteps this entirely: because mask tokens are literal text tokens in the MLLM's vocabulary, the model's output is a text string that can be compared directly against a ground-truth text string. The RL problem reduces to: did the model output the right "words" for the masks?

The textual answer-matching reward function. The paper proposes a simple but effective reward for mask generation tasks, computed entirely through character-level string matching on the MLLM's text output:

$R_{\text{mask}} = N_{\text{TP}} / \max(N_{\text{pred}}, N_{\text{gt}})$

where $N_{\text{TP}}$ is the number of predicted true positive masks (mask token sequences that appear in both the prediction and the ground truth, after deduplication), $N_{\text{pred}}$ is the total number of predicted mask token sequences without deduplication (to penalize repetitive predictions — if the model outputs the same mask token pair five times, it counts as five predictions, but only one can be a true positive), and $N_{\text{gt}}$ is the number of ground-truth mask token sequences.

What this computes: The reward function simultaneously measures precision (through the denominator $\max(N_{\text{pred}}, N_{\text{gt}})$ — if the model predicts more masks than the ground truth, the denominator grows, reducing the reward) and recall (through the numerator $N_{\text{TP}}$ — finding more correct masks increases the reward). The use of $\max(N_{\text{pred}}, N_{\text{gt}})$ rather than just $N_{\text{pred}}$ is a deliberate design choice: if the model predicts fewer masks than the ground truth but all are correct (e.g., $N_{\text{pred}} = 2$, $N_{\text{gt}} = 5$, $N_{\text{TP}} = 2$), the reward is $2/5 = 0.4$ rather than $2/2 = 1.0$, penalizing incomplete predictions. Conversely, if the model over-predicts (e.g., $N_{\text{pred}} = 10$, $N_{\text{gt}} = 3$, $N_{\text{TP}} = 3$), the reward is $3/10 = 0.3$, penalizing false positives.

Why this form: Standard answer-matching rewards for text RL (in math, coding, VQA) simply check if the ground-truth answer string appears anywhere in the model's output — a binary reward (1 for correct, 0 for incorrect). This binary approach doesn't work well for mask generation because tasks like GRES may require multiple masks, and partial correctness (finding 2 out of 3 referred objects) should receive some reward. The $N_{\text{TP}} / \max(N_{\text{pred}}, N_{\text{gt}})$ formulation provides a continuous, granular reward that distinguishes between "got all masks right," "got some masks right," and "got no masks right," enabling more nuanced optimization. The deduplication on the numerator but not the denominator is intentional: it ensures that finding the same correct mask multiple times doesn't inflate the reward (you only get credit for each unique correct mask once), but outputting the same mask multiple times is still penalized as wasteful generation.

The extraction procedure. During RL training, the reward is computed as follows:

1. The MLLM generates a rollout response (a text string) for a given prompt.
2. All substrings between <|mt_start|> and <|mt_end|> are extracted from the response — these are the predicted mask token sequences.
3. The same extraction is performed on the ground-truth answer string.
4. Duplicate predicted mask sequences are identified (for the denominator, which counts without deduplication).
5. Unique predicted mask sequences are compared against unique ground-truth mask sequences — a predicted sequence is a true positive if it exactly matches a ground-truth sequence.
6. The reward is computed using the formula above.

This entire process is purely text-based: no SAMTok decoder is invoked, no masks are reconstructed, no IoU is computed. The reward function operates entirely in the MLLM's token space. This means RL training runs at the same speed as text-only RL, without the overhead of mask decoding (which would require GPU memory for SAM2 and additional forward passes).

RL algorithm and setup. The paper uses Group Relative Policy Optimization (GRPO) [51], which has become the standard RL algorithm for post-training LLMs on reasoning tasks. GRPO works by:

1. For each prompt, sampling $K$ candidate responses from the current model (the "group").
2. Computing a reward for each candidate using the reward function.
3. Computing the relative advantage of each candidate within its group: how much better (or worse) its reward is compared to the group mean.
4. Updating the model to increase the probability of high-advantage responses and decrease the probability of low-advantage responses, using a clipped policy gradient objective similar to PPO.

The relative nature of GRPO (comparing within a group rather than to an absolute baseline) makes it more stable and sample-efficient than standard PPO, and it eliminates the need for a separately trained value function (critic network). This is important for mask generation tasks where the reward function might have high variance (some images are inherently harder than others) — by comparing candidates from the same prompt, GRPO normalizes away prompt-level difficulty.

Cold-start chain-of-thought data. To bootstrap RL training, the authors first generate 26,000 cold-start samples by prompting the Qwen3-VL-235B model (a much larger model than the Qwen2.5VL-3B/7B being fine-tuned) to simulate chain-of-thought (CoT) reasoning for mask generation tasks. This creates an initial distribution of high-quality reasoning traces that the smaller model can learn from during supervised fine-tuning, before RL further refines the generation.

For the actual RL training, the authors select:

• 8,000 challenging samples from the GRES dataset [34] for the thinking RL setting (where the model generates chain-of-thought reasoning before outputting mask tokens).
• 41,000 challenging samples from the GCG dataset [48] for the non-thinking RL setting (where the model directly outputs interleaved text and mask tokens).

The distinction between thinking (CoT) and non-thinking RL matters because GRES requires compositional reasoning about referring expressions (e.g., "the cup that is to the left of the red bowl and above the wooden table"), while GCG requires the model to simultaneously describe a scene and ground mentioned phrases — different reasoning demands.

RL configuration (from Appendix B):

• Framework: Easy-R1 [97] with GRPO [51] algorithm
• Learning rate: $1 \times 10^{-6}$
• Base model: The SFT-trained QwenVL-SAMTok model (after 5M-sample supervised fine-tuning)

Why this enables something new. The paper claims this is "the first successful attempt to optimize mask generation performance using text-only reward signals." The significance is not just the performance gains (which are substantial — +6.8% gIoU, +18.9% N-acc on GRES; +4.5% AP50, +6.6% Recall on GCG), but the demonstration that the entire LLM RL pipeline — GRPO, answer-matching rewards, chain-of-thought reasoning — transfers to pixel-wise tasks without modification when masks are text tokens. This has profound implications: future improvements in RL algorithms, reward modeling, and reasoning strategies developed for text tasks will directly benefit pixel-wise MLLMs that use SAMTok, without requiring separate research programs for "mask RL" vs. "text RL."

4. Key Insights and Innovations

Innovation 1: Masks Are a Language — The Decoupling of Spatial Representation from Linguistic Models

The paper's most fundamental insight is not a specific architecture or training recipe, but a conceptual reframing of masks as linguistic tokens rather than spatial entities. Before SAMTok, the dominant assumption in pixel-wise MLLMs was that masks require specialized spatial handling — feature-pooling modules for mask input, segmentation decoders for mask output — because their 2D continuous nature seemed fundamentally different from the 1D discrete token sequences that language models process. This assumption created a sharp architectural divide: text got handled by the LLM's standard mechanisms, while masks required custom engineering at every interface point.

SAMTok challenges this assumption directly by asking: what if the spatial complexity of masks can be fully absorbed into a standalone tokenizer, making the downstream MLLM's job indistinguishable from processing text? The answer, demonstrated empirically across six task categories and multiple model families (Tables 1–10), is that this separation of concerns works surprisingly well — a 3B-parameter MLLM with no segmentation-specific components achieves SOTA or near-SOTA on GRES (Table 3), GCG (Table 1), interactive segmentation (Table 2), and region captioning (Tables 7–9). The excellent performance is a diagnostic result: it shows that pixel-wise understanding does not require the MLLM to model spatial structure internally; it only requires the MLLM to learn a mapping between natural language queries and a pre-learned "mask vocabulary."

This insight places SAMTok in a direct lineage with the NLP community's abandonment of task-specific architectures in favor of unified text-to-text frameworks, but it makes a substantially different claim. The NLP unification (T5, GPT) argued that diverse language tasks share enough structure that a single architecture suffices. SAMTok argues something stronger: that visual tasks with fundamentally non-linguistic outputs (2D binary masks) can be lossily but effectively represented in a purely linguistic modality — a claim about cross-modal representation capacity, not just architectural unification. This is a genuine conceptual contribution: the finding that two discrete tokens can capture enough mask information for high-quality downstream task performance (evidenced by the 0.70 reconstruction IoU in Table 13 and the competitive SOTA results) implies that the information the MLLM needs about a region — its rough shape, location, and boundary properties — is far lower-dimensional than the raw pixel mask would suggest. The tokeizer absorbs the compression, and the MLLM operates in an information regime where masks are effectively "just another noun phrase" in the conversation.

The scalability implication is equally significant: since SAMTok is a frozen, independently trained component, the burden of making a new MLLM pixel-capable is reduced to vocabulary extension and data preparation. This transforms the engineering challenge from "design a new region encoder and decoder for each MLLM architecture" to "run SAMTok once on the training data and fine-tune any MLLM normally" — a change that the paper validates by showing SAMTok works across Qwen2.5-VL, Qwen3-VL, and PerceptionLM (Table 11) with minimal performance variation. This decoupling is not just convenient; it's a principled separation that allows the tokenizer and the MLLM to be optimized independently on different data and at different cadences — exactly the property that made tokenizers like BPE and SentencePiece foundational to the LLM ecosystem.

Innovation 2: The Discrete Representation as an RL Unlock — Text-Only Rewards for Spatial Tasks

The paper's second distinctive contribution is the diagnosis that continuous mask-to-decoder interfaces are the bottleneck preventing RL-based optimization of pixel-wise MLLMs, paired with the demonstration that discrete mask tokens eliminate this bottleneck entirely. This is not just an engineering convenience — it's a finding that changes what optimization techniques are available for pixel-wise tasks.

Prior work had established that RL can improve MLLMs on reasoning tasks through answer-matching rewards (checking whether the model's textual output matches a ground-truth answer string). However, pixel-wise MLLMs using continuous embeddings (LISA-style segmentation tokens, SAM-integrated decoders) could not participate in this paradigm because their outputs were not text strings — they were continuous vectors fed into downstream decoders. Computing a reward required either (a) decoding the mask, computing IoU, and backpropagating through the decoder (cumbersome and decoder-dependent), or (b) predicting intermediate text representations like bounding boxes, applying RL to those, and converting to masks post-hoc (complex and lossy). Neither approach provided the simplicity of "compare output string to ground-truth string" that made text RL successful.

SAMTok's discrete tokens change this landscape fundamentally. Because mask tokens are literal vocabulary entries, the MLLM's output for a mask generation task is a text string like <|mt_start|><|mt_0042|><|mt_0187|><|mt_end|>. The reward function (Equation 5) reduces to character-level substring matching: does the predicted string contain the same mask-token sequences as the ground truth? The computation requires no mask decoding, no IoU calculation, and no external models — it is purely text-based. This means that the entire RL pipeline developed for text reasoning (GRPO, PPO, DPO, chain-of-thought prompting, cold-start data generation) transfers without modification to mask generation tasks.

The results validate the significance of this unlock quantitatively. Applying GRPO with text-only rewards produces gains on GRES of +6.8% average gIoU and +18.9% average N-acc across all splits (Table 3), pushing the 3B model from 70.1% to 76.7% average gIoU — past the previous 8B-parameter SOTA methods. On GCG, the gains are +4.5% AP50, +2.0% mIoU, and +6.6% Recall on average (Table 1). These improvements come from purely text-based optimization, without any pixel-level supervision during RL. The paper explicitly notes that caption quality metrics decreased slightly because the reward function did not evaluate caption quality — confirming that the mask gains are from mask-specific optimization, not general model improvement.

This is more than a performance result; it's a proof of concept for a new class of RL-able vision tasks. If mask generation can be optimized through text rewards, then any spatially-grounded visual task that can be expressed through discrete token outputs — object detection as center + size tokens, keypoint prediction as coordinate tokens, depth estimation as quantized depth tokens — becomes amenable to the same text-RL paradigm. SAMTok demonstrates the principle concretely for segmentation, but the conceptual implication is broader: discrete tokenization of continuous visual outputs is the bridge between vision tasks and the rapidly advancing LLM RL ecosystem. The paper's finding that this works "surprisingly well" — that two tokens can carry enough information for RL to meaningfully improve mask quality — is the diagnostic result that should motivate similar discretization efforts for other visual outputs.

Innovation 3: Diagnostic Analysis of the Rate–Distortion–Learnability Tradeoff in Mask Tokenization

Beyond the system-level contributions, the paper provides a careful empirical characterization of a three-way tradeoff that governs any discrete mask tokenizer: reconstruction fidelity, codebook/code size, and the MLLM's ability to learn the token space. This is not a new theoretical concept — it's the classic rate–distortion tradeoff from information theory plus a learnability constraint — but the paper's systematic ablation (Tables 12 and 13) turns it from an abstract principle into a design diagnostic with concrete numbers.

The tradeoff has three interacting dimensions:

• Reconstruction fidelity (r-Acc): How accurately can the tokenizer reconstruct a mask from discrete codes? Higher fidelity means the MLLM's generated tokens correspond to masks that are spatially precise, but achieving it requires larger codebooks or more quantization steps.
• Search space size (implicitly, token count × codebook size): How many possible mask representations exist? A larger space (more steps, bigger codebooks) improves reconstruction but makes the MLLM's generation task harder — the model must select among more possible token sequences, and rare codes may never be learned.
• MLLM generation accuracy (g-Acc): How well does the MLLM learn to output the correct mask tokens for a given task? This is the downstream metric that matters, and it is not monotonic in reconstruction fidelity.

Table 13's comparison of 1024×2 vs. 1024×4 (same codebook size, different quantization steps) is particularly illuminating: four-step RQ improves reconstruction (r-Acc 0.75 vs. 0.70) because finer residual quantization captures mask details better, but degrades generation accuracy (g-Acc 77.3 vs. 78.3). The authors attribute this to the exponentially expanded search space — from 1024^2 ≈ 10^6 possible token pairs to 1024^4 ≈ 10^12 possible token quadruples — which makes it harder for the MLLM to learn the correct token sequences from limited data. The paper explicitly notes that "longer words incur much higher computational costs" for dense mask prediction settings. This is the learnability penalty: more expressive representations come with a learning burden that can outweigh the reconstruction benefit.

The default configuration (256×2, meaning two non-shared codebooks of 256 entries each, producing 256^2 = 65,536 possible masks) represents a pragmatic operating point on this tradeoff curve: reconstruction is adequate (r-Acc 0.70), the vocabulary is compact (512 special tokens total), and generation performance is competitive (g-Acc 77.6). Moving to 512×2 improves reconstruction slightly (r-Acc 0.71, g-Acc 77.8) but the authors choose the more compact 256×2 — demonstrating a preference for learnability and inference efficiency over marginal reconstruction gains.

This analysis also surfaces a subtle but important comparison with standard VQ in Table 12. A single large VQ codebook of 65,536 entries achieves r-Acc = 0.66 — worse than RQ's 0.70 with two codebooks totaling only 2,048 entries. The reason is that RQ's two-step structure produces combinatorial codebook utilization: each additional residual step adds a multiplicative factor to the representational capacity (256 × 256 = 65,536 combinations) while keeping individual codebook sizes small. This means the MLLM's learning problem — picking one of 256 options at each of two positions — stays manageable, even though the effective representational capacity equals that of a single massive codebook. Standard VQ asks the MLLM to pick one of 65,536 options at a single position, which empirically leads to worse codebook utilization (many entries never used) and worse downstream performance.

This diagnostic contribution is incremental rather than fundamental — it's an empirical characterization of a known tradeoff — but it is practically valuable because it provides concrete guidance for anyone building discrete tokenizers for visual outputs. The numbers in Tables 12 and 13 serve as reference points: if your mask reconstruction IoU drops below ~0.70, generation performance suffers; if your search space exceeds ~256^4 combinations, the MLLM struggles to learn; two residual steps with medium codebooks strikes the best balance for this model scale and data regime. These thresholds may not generalize to all settings, but the analytical framework — measuring reconstruction, search space, and generation accuracy together rather than optimizing any one in isolation — is the durable contribution.

Innovation 4: Verifier-Free RL as an Existence Proof — Discrete Tokens Remove the Need for External Reward Models

Embedded within Innovation 2 is a more specific and consequential claim that deserves separate treatment: SAMTok demonstrates that mask generation RL can work without an IoU verifier, a mask decoder in the loop, or any external model — using only string matching on the MLLM's text output. This is an existence proof that challenges a implicit assumption in the field.

The prior implicit assumption (visible in the design of Seg-Zero [36], VisionReasoner [37], Lens [101], and others) was that optimizing mask generation through RL requires some form of pixel-level reward signal — typically IoU between predicted and ground-truth masks, computed by decoding the model's output and comparing spatial overlap. This forced a specific architecture: the reward computation must include a mask decoder (to convert model outputs into masks) and a spatial comparison function (to compute IoU), making the RL training loop dependent on segmentation-specific components. Some works attempted to circumvent this by predicting bounding boxes or points (which are naturally textual — coordinate strings), applying RL to those, and then converting to masks with SAM. But this "two-hop" approach (MLLM → boxes → SAM → masks → IoU reward) introduces complexity, latency, and information loss at each hop.

SAMTok's text-only reward function (Equation 5) proves that this complexity is unnecessary. The reward N_TP / max(N_pred, N_gt) is computed by extracting mask-token substrings from the model's text output and checking for exact matches with ground-truth mask tokens — operations that require only string processing. The reward captures precision and recall of mask generation without ever reconstructing a single mask. The success of this approach — +18.9% N-acc on GRES (Table 3) and +6.6% Recall on GCG (Table 1) — validates that optimizing token-level accuracy indirectly optimizes spatial mask quality. This is non-obvious: one might reasonably worry that two masks with very different token representations could be spatially nearly identical (tokenizer aliasing), or that identical tokens could correspond to slightly different masks depending on image context, making token-matching rewards noisy or misaligned with spatial quality. The empirical results suggest that these concerns are either minor in practice or that the GRPO algorithm is robust to them.

The significance of this existence proof extends beyond segmentation. Many visual tasks produce structured continuous outputs — depth maps, surface normals, optical flow, keypoint heatmaps — that currently resist RL because computing rewards requires decoding these outputs and measuring spatial metrics. SAMTok's demonstration that token-level matching can serve as a proxy for spatial quality opens the door to similar discretization + text-RL approaches for these tasks. The key architectural requirement is a tokenizer that: (1) maps continuous visual outputs to discrete tokens with sufficient fidelity, and (2) produces tokens where exact match implies spatial similarity (or near-similarity) with high probability. SAMTok satisfies both, and the GRPO results provide the first concrete evidence that this approach can work at scale.

This is a fundamental insight with practical consequences, not just a performance result. It changes the design space for pixel-wise MLLMs: rather than asking "how do we incorporate pixel-level rewards into RL training?" (which forces decoder-integrated architectures), future designers can ask "how do we tokenize visual outputs so that token-matching rewards are well-aligned with spatial quality?" — a tokenizer design problem, not an RL architecture problem.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The primary evaluation is conducted on the MATH benchmark (Hendrycks et al., 2021), specifically the 500-question test split from Lightman et al. (2022). The training set contains 12,000 questions, and all supervised fine-tuning and PRM training use this training split. The choice of MATH is deliberate: mathematical reasoning problems require multi-step logical deduction rather than novel factual recall, making them ideal for testing whether test-time compute can improve performance when the model already possesses the necessary knowledge (Section 4). Answers are graded using the grading function released by Lightman et al. (2022) (Appendix G), which performs exact string matching after normalization.

• Base model(s). All experiments use PaLM 2-S* (Codey) (Anil et al., 2023) as the base LLM, described as "representative of the capabilities of many contemporary LLMs" (Section 4). The model achieves roughly 10–19% pass@1 on MATH depending on the prompt and sampling configuration — far from saturation, leaving substantial room for test-time compute to improve performance. For the FLOPs-matched comparison (Section 7), a second model with approximately 14× more parameters is used as the pretraining-scaled baseline, though the paper does not specify this model's exact name or architecture beyond it being a larger PaLM 2 variant.

• Metrics. The primary metric is MATH test accuracy (%) — the fraction of the 500 test questions for which the selected final answer matches the ground truth. For search experiments, accuracy is reported using best-of-N weighted selection (Section 5). For revision experiments, accuracy is reported using majority voting or verifier-based selection across the revision chain (Section 6). Difficulty-dependent results are reported as accuracy within each of five difficulty quintiles (Section 3.2). The paper also reports pass@1 for the revision model's per-step outputs (Figure 6, left) to show improvement over the revision chain.

• Baselines. The paper compares against several baselines:

  • Majority voting: selecting the most common final answer among N independently sampled solutions, with no learned verifier.
  • ORM best-of-N weighted: scoring N complete solutions with an outcome reward model (ORM) trained to predict final-answer correctness, then applying best-of-N weighted selection (following Li et al., 2023).
  • PRM best-of-N weighted: same as ORM best-of-N weighted but using the process reward model (PRM) that scores each intermediate step, with last-step aggregation.
  • Parallel sampling (for revision experiments): generating N independent solutions from the revision model and selecting the best via verifier or majority voting.

  For the FLOPs-matched comparison, the baseline is a ~14× larger model using greedy decoding with no additional test-time compute.

• Generation budget / compute accounting. All test-time compute budgets are measured in number of generations — one generation equals one complete sampled answer from the base LLM. For best-of-N and beam search, the budget equals the number of samples or beams N. For lookahead search with k lookahead steps, the cost is N × (k+1) to account for additional rollout computation (Section 5.3). Budgets are swept across powers of 2 from 2^0 to 2^9 (1 to 512 generations). For the FLOPs-matched comparison (Section 7), compute is measured in total FLOPs using standard approximations: pretraining FLOPs X = 6ND_pretrain and inference FLOPs Y = 2ND_inference, where N is parameters, D_pretrain is pretraining tokens, and D_inference is inference tokens. Three values of the ratio R = D_inference / D_pretrain are tested: 0.16 (R ≪ 1), 0.79 (R ≈ 1), and 22 (R ≫ 1).

• Cross-validation / statistical protocol. For compute-optimal strategy selection, the paper uses two-fold cross-validation within each difficulty bin on the 500-question test set (Section 3.2). The best-performing strategy (choice of search algorithm, beam width, sequential-to-parallel ratio) is selected on one fold and evaluated on the other, with results averaged. This prevents the circularity of selecting the best strategy and evaluating it on the same data. Difficulty bins are constructed using either oracle pass@1 (requiring ground-truth labels) or predicted difficulty using the PRM's average final-answer score across 2048 samples per question. Both binning methods are evaluated.

Main Quantitative Results

Search Against PRM Verifiers (Section 5, Figures 3–4)

Aggregate search algorithm comparison (Figure 3, left). Across all 500 MATH test questions, the paper compares best-of-N weighted, beam search (M = √N and M = 4), and lookahead search (k = 1 and k = 3) at generation budgets from 2 to 256:

• At low budgets (2–8 generations), beam search with M = 4 significantly outperforms best-of-N weighted. Quoting from Section 5.3, at 4 generations beam search achieves roughly 27% accuracy versus roughly 16% for best-of-N weighted — an 11-percentage-point gap.
• At high budgets (64–256 generations), beam search performance flattens and falls slightly below best-of-N weighted. Best-of-N weighted reaches approximately 38% at 512 generations, while beam search (M = 4) plateaus around 34%.
• Lookahead search (both k = 1 and k = 3) generally underperforms at the same generation budget due to its higher per-step cost reducing the effective number of beams explored. The 3-step lookahead variants converge to similar performance as other methods at very high budgets but never surpass them.
• Majority voting trails all verifier-based methods substantially, reaching only about 29% at 512 generations.

Difficulty-dependent behavior of search (Figure 3, right). The per-difficulty breakdown (comparing beam search M = 4 versus best-of-N weighted at four budget levels: 4, 16, 64, 256 generations) reveals the central pattern motivating compute-optimal allocation:

• Bin 1 (easiest): Beam search accuracy decreases from roughly 78% to 77% as budget goes from 4 to 256, while best-of-N weighted increases from 68% to 88%. This is the clearest evidence of PRM over-optimization — beam search finds solutions that score highly under the PRM but are incorrect.
• Bin 2: Beam search improves modestly (roughly 14% → 32%) but best-of-N weighted improves faster (roughly 14% → 60%), maintaining a clear advantage at high budgets.
• Bin 3: Beam search consistently outperforms best-of-N weighted across all budgets, reaching roughly 34% vs. 23% at 256 generations.
• Bin 4: Beam search shows the strongest relative advantage, reaching roughly 17% vs. 10% for best-of-N at 256 generations.
• Bin 5 (hardest): Both methods hover near 1–3% regardless of budget. No method makes meaningful progress.

Compute-optimal search (Figure 4). By selecting the best search strategy per difficulty bin at each budget level (using both oracle and predicted difficulty bins):

• At 16 generations, compute-optimal (oracle bins) achieves approximately 27% accuracy, roughly matching PRM best-of-N weighted at 64 generations — a 4× compute reduction.
• At 256 generations, compute-optimal oracle reaches approximately 39.5%, surpassing PRM best-of-N weighted at the same budget (roughly 37%).
• Compute-optimal with predicted difficulty bins tracks the oracle version closely, with the two curves "largely overlapping" (Section 5.3). The predicted version reaches approximately 37% at 256 generations — slightly below oracle but still above best-of-N weighted.
• Both compute-optimal variants consistently outperform ORM best-of-N weighted (which peaks around 34% at 512 generations) and majority voting (around 29%).

The key takeaway from the search experiments: beam search is superior on medium-hard problems (bins 3–4) but degrades on easy problems (bins 1–2) due to verifier over-optimization, and difficulty-conditioned strategy selection recovers most of the theoretical best performance at each budget.

Revision Model Results (Section 6, Figures 5–8)

Revision model pass@1 trajectory (Figure 6, left). Starting from approximately 18.2% pass@1 at step 1, the revision model's per-step accuracy improves to roughly 24–25% by steps 15–20 and remains in the 23–25% range out to 64 steps. The fact that pass@1 continues improving beyond the 4-step training horizon is evidence that the model learns a generalizable revision skill, not just memorizing the training distribution.

Sequential vs. parallel comparison (Figure 6, right). At 64 generations:

• Sequential + best-of-N weighted: approximately 41.5%
• Parallel + best-of-N weighted: approximately 39%
• Sequential + majority: approximately 38%
• Parallel + majority: approximately 35%

Sequential revisions outperform parallel independent sampling under both selection mechanisms, with the verifier-based gap (roughly 2.5 percentage points) being slightly narrower than the majority-based gap (roughly 3 points). This demonstrates that the revision model genuinely improves answer quality across the chain, not just providing more candidates for the verifier to select from.

Sequential-to-parallel ratio sweep (Figure 7, left). For a fixed total generation budget, varying the ratio of sequential revisions per chain to parallel chains reveals the optimal allocation structure:

• At 256 generations, the optimal ratio is around 2:1 to 8:1 sequential-to-parallel, achieving approximately 43–44% accuracy.
• Fully parallel (leftmost point) yields approximately 40%.
• Fully sequential (rightmost point) yields approximately 42%.
• At lower budgets (8–32 generations), fully sequential is optimal — the curves are monotonically increasing with the sequential-to-parallel ratio.

This non-monotonic pattern at high budgets shows that both sequential refinement (improving within a chain) and parallel diversity (exploring different approaches) contribute to performance, and the optimal tradeoff depends on budget.

Difficulty-dependent optimal ratio (Figure 7, right). At a fixed budget of 128 generations, breaking out the sequential-to-parallel ratio sweep by difficulty bin:

• Bin 1: Performance is essentially flat across all ratios, around 90–92%. Easy questions are insensitive to allocation strategy.
• Bin 2: Slight advantage for higher sequential ratios, approximately 63% at fully sequential vs. 58% at fully parallel.
• Bin 3: A clear optimal ratio emerges at moderate sequential-to-parallel values (around 2:1 to 8:1), reaching approximately 42% vs. 35% at the extremes.
• Bin 4: Similar pattern, with the peak at a moderate ratio achieving roughly 18% vs. 14% at fully parallel.
• Bin 5: All ratios produce roughly 2–3% accuracy. No allocation strategy helps.

The key pattern mirrors the search results: easy problems benefit from exploitation (sequential refinement of mostly-correct initial answers), hard problems benefit from exploration (parallel diversity to try different approaches), and the hardest problems remain unsolvable regardless of strategy.

Compute-optimal revisions (Figure 8). Selecting the optimal sequential-to-parallel ratio per difficulty bin:

• At 64 generations, compute-optimal oracle achieves approximately 40%, matching parallel best-of-N weighted at 256 generations — a 4× improvement in compute efficiency.
• At 256 generations, compute-optimal oracle reaches approximately 44%, compared to roughly 41% for best-of-N weighted and 37% for parallel-only.
• Compute-optimal predicted bins perform slightly below oracle bins at high budgets (approximately 41% at 256 generations) but still substantially outperform the parallel baseline.
• Notably, the parallel baseline appears to plateau around 36–37% at high budgets, while compute-optimal scaling continues to improve, suggesting that gains from adaptive allocation compound at higher budgets.

FLOPs-Matched Comparison: Test-Time vs. Pretraining Compute (Section 7, Figures 9 and 1)

The FLOPs-matched comparison asks: given a fixed total FLOPs budget, is it better to train a larger model or to keep the smaller model and spend the extra FLOPs on inference-time computation? The comparison uses PaLM 2-S* with compute-optimal test-time scaling versus a model with approximately 14× more parameters (greedy decoding, no extra test-time compute), at three different values of R = D_inference / D_pretrain: 0.16, 0.79, and 22.

Revisions (Figure 9, left; Figure 1, top-right bar chart). The results from the bar chart in Figure 1 (top-right), which aggregates the five difficulty bins into three groups (easy, medium, hard) for readability:

At R ≪ 1, test-time compute with the smaller model outperforms the 14× larger model across all difficulty levels — including hard problems where the absolute performance is low but the relative advantage exists. At R ≫ 1, test-time compute remains preferable only on easy questions. The negative numbers on hard questions at R ≫ 1 mean that in high-inference-volume regimes, the larger model's per-token cost dominates and pretraining is the better investment.

PRM search (Figure 9, right; Figure 1, bottom-right bar chart). The pattern is starker for PRM search:

PRM search shows weaker FLOPs-matched benefits than revisions, with substantial disadvantages on medium and hard questions even at moderate R values. On easy questions, test-time compute remains preferable across all R regimes, though the margin narrows from +19.1% at R ≪ 1 to +2.0% at R ≫ 1. The 0.0% and negative values on medium/hard questions indicate that the 14× larger model is a better use of compute than PRM search for these difficulty tiers unless inference volume is very low.

Figure 9 detail. The line plots show accuracy per difficulty bin as test-time compute scales, with the 14× larger model's greedy performance marked as stars at three x-axis positions corresponding to the three R values. Where the compute-optimal scaling line is above the star, test-time compute wins. On bin 1 (purple, topmost line for revisions), the scaling line is above all three stars — test-time compute dominates regardless of R. On bin 5 (blue, bottommost line), the line is below all three stars and essentially flat near 0–5% — no amount of test-time compute helps. The intermediate bins show R-dependent crossover points.

Ablation Studies and Robustness Checks

PRM step-wise aggregation strategy (Appendix E, Figure 13): Comparing "min" (minimum score across steps), "prod" (product of step-level correctness probabilities), and "last" (only the final step's PRM prediction) for aggregating per-step PRM scores into a single solution-level score, the paper finds that "last" performs best — approximately 37% at 256 samples versus roughly 35% for "min," 27% for "prod," and 34% for a separately trained ORM. This contradicts prior work (Lightman et al., 2023; Wang et al., 2023) which found "min" to be optimal. The authors hypothesize the discrepancy arises because their PRM is trained with soft Monte Carlo labels rather than binary correctness labels, changing how per-step scores distribute. An interesting consequence: using last-step prediction effectively makes the PRM behave like an ORM at aggregation time, yet the PRM still outperforms a separately trained ORM, suggesting that step-level PRM training provides beneficial representation learning even when intermediate predictions aren't directly used at aggregation time.

PRM vs. ORM comparison (Appendix F, Figure 14): At 2048 samples, PRM best-of-N weighted achieves approximately 40% accuracy versus roughly 35% for ORM best-of-N weighted and roughly 30% for majority voting. The gap between PRM and ORM widens with the number of samples, confirming the PRM's superior scaling properties. This gap, combined with the "last" aggregation finding, provides the key insight: the PRM's advantage does not come from using intermediate step information at selection time (since "last" ignores it), but from the better representations learned during PRM training compared to ORM training.

Revision model verifier choice (Appendix J, Figure 15a): The base-LM PRM (trained on PaLM 2-S* base model outputs) underperforms a revision-specific ORM when scoring revision model outputs. Sequential + base-LM PRM achieves roughly 40% at 64 generations versus sequential + revision ORM at roughly 42%. This confirms that distribution shift between the base model's outputs (used to train the PRM) and the revision model's outputs is a practical concern — the PRM's scores become less calibrated on revision outputs.

Revision history in verifier context (Appendix J, Figure 15b): Including previous revisions in the ORM's input context provides a small improvement over the no-history ablation (approximately 1–2 percentage points at 64 generations), but both variants substantially outperform the parallel baseline. This confirms that the sequential sampling benefit is not solely attributable to the verifier seeing more context from the revision chain — the revision model genuinely produces better candidates.

Oracle vs. predicted difficulty bins (Figures 4, 8, and Appendix C, Figures 11–12): Both oracle and predicted difficulty bins yield qualitatively similar trends across difficulty levels, confirming that the compute-optimal strategy works without ground-truth labels. Predicted bins show slightly lower performance at high budgets in the revision setting (roughly 41% vs. 44% at 256 generations in Figure 8) but essentially identical performance in the search setting (Figure 4). This is the critical robustness check — if the gains required oracle difficulty, the approach would be circular. The fact that predicted bins work nearly as well means the system is deployable without access to answers.

Majority voting for revisions (Appendix B, Figure 10): The sequential-to-parallel ratio trends observed with verifier-based selection are replicated with majority voting: easy questions are insensitive to ratio, hard questions show an optimal intermediate ratio, and fully sequential marginally outperforms fully parallel in aggregate. This demonstrates that the difficulty-dependent optimal allocation is a fundamental property of the revision process, not an artifact of the verifier's behavior — majority voting, which makes no learned quality assessments, shows the same qualitative patterns.

ReST^EM revision model optimization (Appendix K, Figure 16): Attempting to further optimize the revision model using ReST^EM (Singh et al., 2024), a reinforcement learning-based self-improvement method that uses on-policy data collection, substantially hurts performance. At 256 generations, fully sequential performance with the ReST^EM model drops to approximately 33.5% compared to roughly 38.5% at the optimal ratio for the base revision model. The authors hypothesize that on-policy data collection in ReST^EM exacerbates spurious correlations in revision data, causing the model to fail to learn the revision task properly. This is a notable negative result that highlights the sensitivity of revision training to the data generation procedure — offline data construction with edit-distance-based pairing (the paper's approach) works, while on-policy RL-based improvement backfires.

Critical Assessment

Claim 1: "Compute-optimal scaling improves efficiency by more than 4× over best-of-N."

This claim holds, but with important qualifications about what "4×" means and in what regime it applies. The 4× figure specifically refers to matching the accuracy of a best-of-N baseline with 4× fewer generations: for search, 16 generations of compute-optimal scaling matches 64 generations of best-of-N weighted (Figure 4); for revisions, 64 generations of compute-optimal matches 256 generations of best-of-N weighted (Figure 8). The evidence is consistent across both oracle and predicted difficulty settings, though the gap narrows somewhat with predicted bins at the highest budgets (Figure 8 shows predicted bins at ~41% vs. oracle at ~44% at 256 generations).

The qualification is that the 4× figure is computed without accounting for difficulty estimation cost. The paper's method for estimating difficulty — generating 2048 samples per question and scoring them with the PRM — is computationally more expensive than the largest test-time budgets being studied (256–512 generations). The authors acknowledge this in Section 3.2: "our experiments do not account for this cost largely for simplicity." In a realistic deployment, the total cost would be difficulty estimation + strategy execution, and the former could dominate the latter. The 4× figure should therefore be understood as an upper bound on achievable efficiency gain once a cheap difficulty estimator is available. The paper does not demonstrate that such an estimator exists or that the gains persist when using one, making this claim partially aspirational.

Additionally, the gains are largest in the low-to-moderate budget regime (16–64 generations). At very high budgets (256+), the advantage narrows, and in the revision case with predicted bins, compute-optimal scaling achieves only ~41% vs. ~37% for best-of-N weighted at 256 generations (Figure 8) — still an improvement but closer to 1.1× than 4×. The 4× claim is most accurate for "matching moderate-budget best-of-N with lower-budget compute-optimal," not for "always achieving 4× improvement at any budget level."

Claim 2: "Test-time compute with a smaller model can outperform a ~14× larger model."

This claim is rigorously supported but only under specific, clearly bounded conditions, and the paper is commendably transparent about these boundaries. The claim holds convincingly when:

• Problems are easy-to-medium difficulty (bins 1–3), where the base model already produces correct solutions at a non-trivial rate (~10%+ pass@1)
• The inference-to-pretraining token ratio R is low (R ≪ 1 or R ≈ 1), meaning total inference compute is small relative to pretraining compute
• The method is revisions (not PRM search), which shows stronger FLOPs-matched benefits

The claim weakens or reverses when:

• Problems are hard (bins 4–5), where test-time compute provides minimal benefit regardless of budget (Figure 9, bin 5 lines are flat near 0–5%)
• R is large (R ≫ 1), where the larger model's per-token inference cost dominates the total budget
• PRM search is used instead of revisions (Figure 9, right), where even medium-difficulty problems show −30.8% to −35.3% relative disadvantage at R ≈ 1 and R ≫ 1

The paper does an excellent job of characterizing where the substitution fails rather than overselling a universal advantage. The statement that "test-time compute amplifies existing capability but does not create it from nothing" (Section 7 takeaway box) is empirically well-grounded: bin 5 problems show essentially zero improvement under any method, confirming that if the base model cannot produce correct solutions at all, no amount of search or revision helps.

However, there is an important weakness in the baseline: the 14× larger model uses greedy decoding only — no majority voting, no best-of-N, no search, no revision. This is acknowledged in Section 7, where the paper notes they "fix training data and scale only model parameters when increasing pretraining compute, matching the approach of the LLaMA model series (Touvron et al., 2023)." The paper concedes that "compute-optimal pretraining would scale both data and parameters equally (Hoffmann et al., 2022)," and leaves that comparison to future work. Giving the larger model even a modest test-time compute budget (e.g., best-of-8 or best-of-16) would create a much stronger baseline and would likely narrow or reverse many of the claimed advantages. The current comparison is really "small model + optimized test-time compute vs. large model + zero test-time compute," not "test-time compute vs. pretraining." A fairer comparison would give both models the same total FLOPs budget split between pretraining and inference — but this would require optimizing the larger model's test-time strategy as well, which significantly expands the experimental scope.

Claim 3: "Efficacy of test-time compute strategies depends critically on prompt difficulty."

This is the most robustly supported claim in the paper and arguably its most important scientific contribution. The evidence is overwhelming and replicated across multiple independent experimental axes:

• In search (Figure 3, right), beam search degrades easy-problem performance at high budgets while improving medium-problem performance — qualitatively opposite effects of the same strategy at different difficulty levels.
• In revisions (Figure 7, right), the optimal sequential-to-parallel ratio shifts from "fully sequential is best" on easy problems to "balanced ratio is best" on medium-hard problems — a non-monotonic relationship between difficulty and optimal strategy.
• In FLOPs-matched analysis (Section 7), the substitution of test-time compute for pretraining compute works on easy problems but fails on hard problems — a binary qualitative difference.
• The patterns replicate across both search and revision mechanisms, across oracle and predicted difficulty bins, and across verifier-based and majority-based selection.

The robustness is noteworthy: a difficulty-dependent effect that appears in only one experiment could be a fluke; an effect that appears in five independent comparisons (search algorithms × difficulty, revision ratios × difficulty, FLOPs-matched × difficulty, revisions × majority voting × difficulty, PRM vs. ORM × sample count) is almost certainly real. This claim does not depend on the 4× efficiency figure, the FLOPs-matched comparison, or any specific method — it is a structural finding about test-time compute that should generalize to other models and tasks, though the specific difficulty thresholds and optimal strategies would need to be re-estimated.

Weaknesses and missing experiments:

• Single benchmark, single model family. All results are on MATH with PaLM 2-S*. The paper argues this model is "representative" but provides no evidence — no replication on Llama, GPT, Gemini, or Claude families; no testing on non-math reasoning benchmarks (code generation, logical reasoning, scientific QA). The difficulty-dependent patterns (over-optimization on easy problems, no improvement on hard problems) could be specific to math reasoning, where correctness is binary and verifiable, or to PaLM 2's particular error patterns. The paper would be substantially strengthened by at least one out-of-domain evaluation.

• Small test set for strategy selection. The 500-question test set split into five difficulty quintiles (~100 per bin), then further split by two-fold cross-validation, means the compute-optimal policy is selected based on ~50 questions per fold per bin. This is a very small sample for discrete strategy selection (choosing among 3–4 search algorithms or 7–8 sequential-to-parallel ratios). The paper does not report confidence intervals on the compute-optimal scaling curves, making it difficult to assess whether the observed gains are statistically reliable or could be noise from the small bin sizes. Given the coarse discretization (5 bins), a question at the boundary of two bins could have its strategy substantially affected by which side it falls on.

• No combination of search and revisions. The paper studies PRM tree-search and iterative revisions as independent mechanisms but never combines them. Section 8 acknowledges this as a limitation: "we did not experiment with PRM tree-search techniques in combination with revisions." This is a significant gap because the two mechanisms have complementary strengths: revisions improve candidate quality (better proposal distribution), while PRM search improves candidate selection (better verifier). A natural next experiment — using the revision model as the proposal distribution within beam search — is not run, meaning the current results represent a lower bound on what a fully integrated system could achieve.

• Unclear generalizability of the predicted difficulty estimator. The PRM-based difficulty estimator (averaging PRM final-answer scores over 2048 samples) works well for PaLM 2-S* on MATH (Figures 4, 8 show oracle and predicted bins largely overlapping). But the estimator's quality depends on how well the PRM's confidence correlates with actual correctness — a property that could vary substantially across models, tasks, and domains. If the PRM is poorly calibrated (e.g., overconfident on certain error types), the predicted bins would misclassify difficulty, and the compute-optimal policy would degrade. The paper does not analyze the calibration of the difficulty estimator or its robustness to distribution shift.

• No latency analysis. The paper measures compute in "generations" (total FLOPs proxy) but ignores wall-clock time. Sequential revisions are inherently serial — each revision depends on the previous one — while parallel best-of-N can be executed simultaneously with sufficient hardware. A strategy allocating 128 generations as 64 sequential × 2 parallel takes roughly 64× longer than 128 parallel samples in wall-clock time. For latency-sensitive applications, the sequential-heavy strategies favored by the compute-optimal policy on easy problems may be practically unusable regardless of their FLOPs efficiency. The paper never addresses this tradeoff, which is a significant oversight for a method that claims practical relevance.

• The 14× larger model may not be compute-optimally trained. As discussed above, the pretraining baseline scales parameters only (not data), following the LLaMA paradigm rather than Chinchilla-optimal training. A Chinchilla-optimal model trained with 14× more total FLOPs (scaling both parameters and data) would likely outperform the parameter-only-scaled model used here, making the test-time compute advantages smaller or non-existent. This is not a fatal flaw — the paper is transparent about it — but it means the FLOPs-matched results are a lower bound on pretraining's effectiveness and may overstate the case for test-time compute.

What would strengthen the paper:

1. A cheap difficulty estimator validated against the expensive 2048-sample method, with the full cost (estimation + execution) included in the budget accounting.
2. At least one out-of-domain evaluation (code generation, scientific QA) to test whether the difficulty-dependent patterns generalize.
3. A combined search + revisions experiment, even at small scale, to establish whether the two mechanisms are complementary or redundant.
4. Confidence intervals or bootstrap estimates on the compute-optimal scaling curves, particularly given the small per-bin sample sizes.
5. A latency-aware reanalysis of the sequential-to-parallel tradeoff, annotating which strategies are practical for latency-sensitive vs. throughput-oriented deployments.
6. A fairer FLOPs-matched baseline where the larger model also gets some test-time compute budget (e.g., best-of-8 or best-of-16), or a Chinchilla-optimally trained larger model.

These gaps do not undermine the paper's core contributions — the difficulty-dependent scaling patterns, the 4× efficiency gains (conditional on cheap difficulty estimation), and the diagnosis of verifier over-optimization as the primary bottleneck are all well-supported. But they do bound the strength of the claims and suggest that the practical deployment picture is more nuanced than the paper's narrative might suggest.

6. Limitations and Trade-offs

Difficulty Estimation Cost Is Excluded From the Headline Efficiency Numbers

The assumption or constraint. The entire compute-optimal framework depends on knowing each prompt's difficulty before deciding how to allocate the inference budget. The paper's method for estimating difficulty — generating 2048 samples per question and either checking ground-truth correctness (oracle) or averaging PRM final-answer scores (predicted) — is extraordinarily expensive. The authors acknowledge this in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

Generating 2048 samples per question consumes more compute than the largest test-time budgets studied in the paper (256–512 generations). This means that in any realistic deployment, the total cost would be 2048 + N generations per question, not N generations, and the difficulty estimation cost would dominate for all but the highest-budget scenarios.

The consequence. The paper's headline claim — "more than 4× better efficiency over a standard best-of-N baseline" (Section 1) — is computed after difficulty is known, without amortizing the cost of learning it. If difficulty estimation costs 2048 generations and the compute-optimal strategy uses 16 generations (matching best-of-N at 64, giving the claimed 4×), the total cost is actually 2048 + 16 = 2064 generations — far more than the 64 generations of best-of-N it was compared against. The 4× figure is therefore an upper bound on achievable efficiency that assumes a zero-cost difficulty oracle. In practice, until a cheap difficulty estimator exists, the approach as described is more expensive than the baseline it claims to beat, not less.

The problem is compounded by the fact that the predicted difficulty estimator (using PRM scores) is only validated as a relative ranking mechanism (it produces bins that correlate with oracle bins), not as a cost-effective deployment tool. Figure 4 and Figure 8 show that predicted bins perform similarly to oracle bins in terms of strategy selection quality, but both rely on the same 2048-sample generation step. The paper never tests whether difficulty can be estimated with, say, 8 or 16 samples instead of 2048, nor does it train a lightweight classifier that predicts difficulty from the question text alone — an approach the authors explicitly flag as future work (Section 8).

What evidence exists in the paper. The limitation is acknowledged in Section 3.2, but the magnitude of the cost is never quantified or compared to the test-time budgets being studied. The paper provides no experiment where difficulty estimation cost is included in the total budget. The 4× efficiency claim appears in Section 1 and is supported by Figures 4 and 8, both of which exclude difficulty estimation cost from the generation budget axis.

Mitigation status. The paper does not mitigate this limitation. It acknowledges the exploration-exploitation tradeoff (Section 3.2: "compute spent assessing difficulty versus compute spent solving the problem") and flags it as "a key avenue for future work," but provides no solution, no cheaper approximation, and no analysis of how the efficiency claims degrade as difficulty estimation cost increases. The authors suggest future work on "pretraining or finetuning models to directly predict difficulty of a question" (Section 8), but this model is not developed or evaluated.

Test-Time Compute Provides No Benefit on the Hardest Problems

The assumption or constraint. The paper's approach assumes that the base model can produce correct solutions at some non-trivial rate — that there exist correct answers in the proposal distribution to find (via search) or refine (via revisions). This assumption fails for problems that are fundamentally beyond the model's capabilities. The authors are explicit about this boundary in Section 7:

"test-time compute amplifies existing capability but does not create it from nothing"

On difficulty bin 5 (the hardest 20% of MATH problems by base model pass@1), the base model's pass@1 is near zero — it almost never produces a correct solution even with 2048 independent attempts.

The consequence. For the hardest problems, no amount of test-time compute helps. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all search methods and all budgets from 4 to 256 generations — the scaling curves are essentially flat. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy regardless of the sequential-to-parallel ratio at 128 generations. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is flat near 0–5%, and the ~14× larger model always outperforms test-time compute on these problems regardless of the ratio R. The paper reports a −52.9% relative disadvantage for PRM search on hard problems at R ≫ 1 (Figure 1, bottom-right bar chart), and even the revision approach shows a −37.2% disadvantage (Figure 1, top-right bar chart).

This is a fundamental capability bound, not an inefficiency that can be engineered around. If the model does not know how to solve a class of problems at all — if correct solutions are simply not in its output distribution — then search cannot find them and revisions cannot refine toward them. The practical implication is that test-time compute is not a substitute for pretraining on genuinely novel or out-of-distribution reasoning tasks. For problems in bin 5, there is no alternative to training a larger or more capable model.

What evidence exists in the paper. The bin 5 results appear in every difficulty-conditioned analysis in the paper: search algorithms (Figure 3, right), revision sequential-to-parallel ratios (Figure 7, right), compute-optimal scaling for both search and revisions (Figures 4 and 8, where the bin 5 curves are essentially flat and close to zero), and the FLOPs-matched comparison (Figure 9, where the bin 5 line is below all three R-value stars). The flatness across all methods and budgets is the key evidence that this is not a method-specific limitation but a fundamental capability ceiling.

Mitigation status. The paper is transparent about this limitation (Section 7 takeaway box) but does not attempt to solve it — and arguably should not, since solving it would require improving the base model's pretraining, which is outside the scope of a test-time compute paper. The limitation is inherent to the problem formulation: test-time compute optimizes the use of an existing model, and if the model is fundamentally incapable on a class of problems, no optimization can compensate. The paper's contribution here is characterizing the boundary rather than eliminating it.

The FLOPs-Matched Baseline Is Not Compute-Optimally Trained and Uses Only Greedy Decoding

The assumption or constraint. The FLOPs-matched comparison in Section 7 compares PaLM 2-S* with compute-optimal test-time scaling against a model with approximately 14× more parameters. The paper acknowledges two weaknesses in this baseline in Section 7:

"We fix training data and scale only model parameters when increasing pretraining compute, matching the approach of the LLaMA model series (Touvron et al., 2023)... We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

Additionally, the larger model uses only greedy decoding — no majority voting, no best-of-N, no search, and no revision. The comparison is therefore "small model + optimized test-time compute" versus "large model + zero test-time compute," not "test-time compute versus pretraining compute."

The consequence. Both weaknesses bias the comparison in favor of test-time compute. First, a Chinchilla-optimal model (Hoffmann et al., 2022) trained with 14× more total FLOPs — scaling both parameters and data equally — would likely outperform a parameter-only-scaled model. The paper's larger model may be undertrained relative to its parameter count, making it a weaker baseline than what compute-optimal pretraining would produce. The magnitude of this effect is unknown because the paper does not compare against a Chinchilla-optimal larger model.

Second, giving the larger model even a modest test-time compute budget would create a substantially stronger baseline. For example, best-of-8 with majority voting on the larger model might close or reverse many of the claimed advantages, especially on easy-to-medium problems where the paper's own results show that best-of-N already provides significant gains over greedy decoding (Figure 3, left). The paper's finding that test-time compute can outperform a 14× larger model on easy problems at R ≪ 1 (e.g., +27.8% relative improvement for revisions, from Figure 1) could shrink or disappear if the larger model were also allowed best-of-8 or best-of-16.

What evidence exists in the paper. The FLOPs-matched results are in Section 7, Figures 9 and 1. The paper explicitly states the baseline's limitations (parameter-only scaling, greedy decoding), but does not provide any alternative comparisons — no Chinchilla-optimal larger model, no larger model with test-time compute, no sensitivity analysis on how much of the advantage disappears if the larger model uses best-of-N. The relative improvements reported in Figure 1 (e.g., +27.8%, +19.1%) are therefore upper bounds on the advantage of test-time compute over pretraining, and the true advantage against a properly optimized baseline could be substantially smaller.

Mitigation status. The paper partially mitigates this limitation through transparency — it acknowledges both weaknesses explicitly and frames the results as a specific, parameter-only-scaling comparison rather than a universal claim. However, it does not provide the experiments that would bound the magnitude of the bias (e.g., "if the larger model used best-of-8, the +27.8% would reduce to +X%"). The promise to address compute-optimal pretraining scaling in future work is noted but not fulfilled. A practitioner reading the FLOPs-matched results should understand them as an existence proof that test-time compute can beat pretraining under specific (and somewhat favorable) conditions, not as a general prescription that it will beat pretraining in practice.

All Results Are on a Single Benchmark With a Single Model Family

The assumption or constraint. Every experiment in the paper uses the MATH benchmark (500 test questions) with PaLM 2-S* as the base model. The paper states that it "believe[s] this model is representative of the capabilities of many contemporary LLMs" (Section 4), but provides no empirical evidence for this claim — no replication on other model families (GPT, Llama, Gemini, Claude), no testing on non-math reasoning benchmarks (code generation, logical reasoning, scientific QA), and no cross-domain validation.

The consequence. The paper's central findings — the difficulty-dependent scaling patterns, the 4× efficiency gains from compute-optimal allocation, the specific strategies that work best at each difficulty level, and the FLOPs-matched tradeoff threshold — may be specific to mathematical reasoning with PaLM 2 models. Several aspects of the findings have plausible model-specific explanations:

• The PRM's over-optimization behavior (beam search degrading easy-problem performance at high budgets, Figure 3 right) depends on the specific errors and biases in PaLM 2-S*'s output distribution. A model with different calibration properties or different types of reasoning errors might show different over-optimization thresholds or might not over-optimize at all at the same budgets.
• The revision model's ability to improve through sequential refinement depends on the base model's in-context learning and self-correction capabilities, which vary substantially across model families (some models benefit more from self-reflection than others).
• The optimal sequential-to-parallel ratio (Figure 7) and the difficulty thresholds at which strategies switch (Section 5.3) are almost certainly model-specific — a stronger base model would compress the difficulty distribution, shifting bin boundaries and changing which strategies are optimal for which problems.
• MATH consists of competition-level problems with clean, verifiable answers. It is unclear whether the difficulty-dependent patterns generalize to tasks where correctness is graded on a continuum (e.g., summarization quality), requires human judgment (e.g., creative writing), or involves multi-modal reasoning (e.g., visual QA).

What evidence exists in the paper. The limitation is acknowledged implicitly (Section 4 mentions the model is "representative" without evidence) but is never tested. The paper provides no out-of-domain evaluation, no cross-model comparison, and no analysis of whether the difficulty-dependent patterns persist on non-math tasks. All quantitative results (Figures 3–9, Tables 1–3 in the main paper, all appendix figures) are on MATH with PaLM 2-S*.

Mitigation status. The paper does not mitigate this limitation. It does not even discuss the generalizability of the difficulty-dependent scaling patterns to other domains or model families. The computational cost of replicating the full experimental suite (training PRMs, revision models, sweeping strategies, and running FLOPs-matched comparisons) on multiple benchmarks with multiple models is acknowledged to be prohibitive, but the absence of even a small-scale cross-validation (e.g., testing whether the difficulty-bin assignments from PaLM 2-S* predict optimal strategies for a different model, or testing one search configuration on a code generation benchmark) leaves the generalizability claims entirely unsubstantiated. A practitioner deploying these methods with a different model or on a different task domain would need to replicate the full analysis pipeline to determine the appropriate strategies and difficulty thresholds — the paper's specific numbers may not transfer.

Sequential Revisions Are Fundamentally Serial and Create a Latency–Throughput Tradeoff the Paper Does Not Address

The assumption or constraint. The paper measures test-time compute in "generations" (total number of sampled solutions), which is a reasonable proxy for total FLOPs but ignores wall-clock latency. Sequential revisions are inherently serial: each revision step depends on the output of the previous step through autoregressive generation, and the revision chain cannot be parallelized. In contrast, parallel best-of-N can execute all N samples simultaneously given sufficient hardware. The paper does not discuss this tradeoff anywhere.

The consequence. The compute-optimal allocation strategies derived in the paper often favor sequential-heavy configurations. On easy problems (bins 1–2), the optimal strategy is fully sequential revisions (Figure 7, right) — a single chain of length N. At low total budgets (8–32 generations), fully sequential is optimal across all difficulty levels (Figure 7, left). Even at higher budgets, the optimal ratio often has a substantial sequential component (2:1 to 8:1 sequential-to-parallel).

In terms of wall-clock time, a strategy allocating 64 generations as 64 sequential revisions takes approximately 64× longer than 64 parallel samples (assuming sufficient GPU parallelism to run all samples simultaneously). For latency-sensitive applications — interactive assistants, real-time decision-making, any deployment where the user is waiting for a response — this makes sequential-heavy strategies practically unusable regardless of their FLOPs efficiency or accuracy advantages. A user experiencing a 30-second delay from 64 sequential revisions may prefer a 1-second response from 64 parallel samples, even if the latter is 2–3 percentage points less accurate.

The paper also does not discuss the hardware assumptions underlying its choice of "generations" as the cost metric. The FLOPs-matched comparison in Section 7 uses total inference FLOPs (Y = 2ND_inference), which treats all inference FLOPs as having equal cost. In practice, serial FLOPs are constrained by per-GPU throughput and memory bandwidth, while parallel FLOPs can be distributed across many GPUs. A deployment with 64 GPUs can run 64 parallel samples in roughly the time of one generation, making parallel sampling extremely cost-effective in latency-normalized terms. The paper's cost model does not capture this distinction.

What evidence exists in the paper. The paper provides no latency analysis, no wall-clock time measurements, and no discussion of the latency–throughput tradeoff. The generation budget N is the only cost metric used throughout Sections 5–7. The sequential-to-parallel ratio sweeps (Figure 7) are presented purely in terms of accuracy at a given N, with no annotation of which strategies are practical for latency-constrained deployments.

Mitigation status. The paper does not acknowledge, measure, or mitigate this limitation. This is a significant oversight for a method that claims practical relevance. A latency-aware reanalysis would likely shift the recommended strategies substantially — favoring parallel-heavy configurations for latency-sensitive applications and reserving sequential-heavy configurations for offline batch processing where throughput matters more than per-query latency. The absence of this analysis means a practitioner reading the paper could deploy the recommended compute-optimal strategies and find them unacceptably slow, despite being "optimal" in FLOPs terms.

The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate and the Training Recipe Is Fragile

The assumption or constraint. The revision model is trained only on sequences where all in-context answers are incorrect, followed by a correct answer (Section 6.1). This creates a fundamental asymmetry: the model learns to revise incorrect answers into correct ones, but never learns what to do when the current answer is already correct. At test time, when a revision chain generates a correct answer at step k, the model may "revise" it into an incorrect answer at step k+1 because its training never included examples of preserving correctness. The paper reports:

"approximately 38% of correct answers get converted back to incorrect ones using a naive approach" (Section 6.1)

The consequence. The revision chain is unstable: progress is not monotonic, and a correct answer produced at step 5 may be lost at step 6. To mitigate this, the paper uses within-chain selection — majority voting or verifier-based selection across all steps of the revision chain, not just the final step (Section 6.1). This means the system must generate the entire chain, then retrospectively select the best answer, which has two drawbacks: (1) it wastes compute on revisions that degrade correct answers, and (2) it requires storing and scoring all intermediate outputs rather than terminating early when a correct answer is found. A revision model that could reliably recognize correct answers and stop revising would be more efficient and would not need the within-chain selection mechanism.

Additionally, the fragility of the revision training recipe limits its transferability. The paper's specific training approach — offline data construction with edit-distance-based pairing of incorrect and correct solutions, training only on correct answer tokens — works for PaLM 2-S* on MATH, but the negative result with ReST^EM (Appendix K, Figure 16) shows that seemingly reasonable modifications to the training procedure can cause the model to completely fail at the revision task:

"the on-policy data collection in ReST^EM exacerbates spurious correlations in revision data, causing the model to fail to learn the revision task properly" (Appendix K)

At 256 generations, the ReST^EM-trained model's fully sequential performance drops to approximately 33.5% compared to roughly 38.5% at the optimal ratio for the base revision model — a degradation rather than an improvement from additional RL-based training. This suggests that the revision approach is sensitive to dataset construction and training methodology in ways that are not fully understood, and that practitioners cannot simply apply standard RL fine-tuning recipes and expect improvements.

What evidence exists in the paper. The 38% reversion rate is reported in Section 6.1. The within-chain selection mitigation is described in the same section. The ReST^EM negative result appears in Appendix K and Figure 16, showing that the ReST^EM-optimized model performs substantially worse than the base revision model, particularly in the fully sequential regime. The paper notes this as a cautionary finding but does not analyze the root cause in depth.

Mitigation status. The paper partially mitigates the reversion problem through within-chain selection (majority voting or verifier-based selection across the chain), which recovers correct answers even if they are later "un-revised." However, this is a post-hoc patch, not a solution — the model still wastes compute generating revisions that degrade correct answers, and the selection mechanism adds overhead. The paper does not explore more principled solutions, such as training the revision model on mixed sequences (some correct-to-correct, some incorrect-to-correct) to teach it when to stop, or adding an explicit "no revision needed" token. The ReST^EM failure is acknowledged but not investigated further, leaving open questions about what specifically causes the degradation and whether alternative RL-based optimization approaches could avoid it. A practitioner seeking to deploy revisions in a new domain would need to carefully validate the training recipe and may encounter similar fragility.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper does not propose a better segmentation model or a more sophisticated pixel-wise MLLM architecture. Instead, it makes a more fundamental move: it eliminates the need for specialized segmentation architectures in MLLMs entirely, by demonstrating that a frozen, standalone mask tokenizer — producing just two discrete text tokens per mask — is sufficient to give any base MLLM pixel-wise capabilities through standard language modeling training. This is not an incremental improvement to the LISA/Sa2VA lineage of segmentation-decoder-integrated MLLMs; it is a reframing that renders that entire architectural category optional. If SAMTok's approach generalizes, the field's default answer to "how do we add pixel-wise capabilities to our new MLLM?" shifts from "design a region encoder and mask decoder, co-train them with the LLM, balance segmentation and language losses" to "download the SAMTok tokenizer, preprocess your training data to convert masks to two tokens, and fine-tune your MLLM normally." The magnitude of this simplification is substantial: it reduces the engineering challenge from custom multi-objective architecture design to vocabulary extension plus data formatting.

The paper reconciles a tension that has been implicit in the pixel-wise MLLM literature since LISA [25] popularized the SAM-decoder-integration paradigm. On one side, decoder-integrated approaches achieve strong performance (LISA, Sa2VA, OMG-LLaVA, HyperSeg, PaDT) but lock the MLLM into a specific decoder architecture, require co-training, and block RL. On the other side, text-output approaches (VisionLLM, Ferret, bounding-box methods) preserve architectural simplicity but sacrifice spatial precision and cannot straightforwardly apply RL either (continuous embeddings block it; coordinate-based approaches need SAM as a post-hoc bridge). SAMTok resolves this tension by showing that the spatial precision of decoder-integrated methods can be decoupled from the MLLM architecture and absorbed into a frozen tokenizer, while the architectural simplicity of text-output approaches can be preserved at the MLLM level. The resolution is clean: the tokenizer handles all spatial complexity; the MLLM operates purely in text space. Prior work treated this as a fundamental tradeoff (spatial precision requires architectural complexity); SAMTok shows it is a separable concern.

More subtly, this paper identifies and resolves a bottleneck that had not been clearly articulated before: continuous mask representations are the root cause preventing RL-based optimization of pixel-wise MLLMs. Prior RL-for-vision works like Seg-Zero [36], VisionReasoner [37], and Lens [101] all had to route reward computation through de-tokenization steps (decoding masks from continuous embeddings, computing IoU) or through intermediate representations (predicting boxes, converting to masks with SAM). The field implicitly accepted this complexity as necessary — that mask generation RL required pixel-level reward signals. SAMTok's demonstration that a purely token-matching reward N_TP / max(N_pred, N_gt) — computed through character-level string comparison, with no mask decoding whatsoever — can produce substantial gains (+6.8% gIoU, +18.9% N-acc on GRES; +4.5% AP50, +6.6% Recall on GCG in Table 3 and Table 1) is an existence proof that challenges this implicit assumption. It shows that the bottleneck was not a fundamental property of mask generation, but an artifact of continuous representations.

This finding redirects research attention in a specific way: rather than designing better RL algorithms that can handle continuous visual outputs (a challenging optimization problem), the field should focus on building better discrete tokenizers for visual outputs — a representation learning problem that is both more tractable and more composable with the rapidly advancing LLM RL ecosystem. Every improvement in text-based RL (better reward modeling, more stable GRPO variants, chain-of-thought reasoning strategies) will directly benefit pixel-wise MLLMs that use discrete mask tokens, without requiring any modification to the pixel-wise pipeline. This is a force multiplier: the RL community's progress is automatically inherited.

The paper also establishes verifier-free RL for mask generation as a viable paradigm. In the language domain, answer-matching rewards work because correctness is textually verifiable — the model's output is a string that can be compared to a ground-truth string. SAMTok shows that the same paradigm extends to masks when masks are text tokens: correctness is "did the model output the right two tokens?" rather than "does the decoded mask have high IoU?" This is not a minor implementation detail; it changes what experiments are possible. A researcher who wants to improve mask generation through RL no longer needs to implement mask decoding in the training loop, manage the memory overhead of running SAM2 during RL, or worry about the non-differentiability of IoU computation. They can use the same RL infrastructure they use for math or code reasoning, with a different reward function.

Finally, the paper provides the first systematic characterization of the rate–distortion–learnability tradeoff for discrete mask tokenization (Tables 12 and 13). The specific numbers — 256×2 as the sweet spot, 1024×4 degrading generation despite better reconstruction, RQ outperforming VQ and FSQ — give the field concrete reference points for future tokenizer design. The analytical framework (measure reconstruction IoU, generation accuracy, and search space size together; optimize generation accuracy, not reconstruction fidelity alone) is the durable contribution that should guide any future work on discrete visual tokenizers.

Certain research directions become less attractive after this work. The paper's demonstration that a 3B model with SAMTok can outperform 7–8B models with custom segmentation decoders on GRES and GCG (Tables 3 and 1) weakens the case for investing in increasingly sophisticated decoder-integrated pixel-wise MLLMs — if a simple tokenizer + standard fine-tuning achieves SOTA, the marginal benefit of custom decoders is shrinking. Similarly, approaches that represent masks as dozens or hundreds of tokens (RLE text, polygon coordinates, autoregressive mask image generation) face a steep uphill argument: SAMTok achieves competitive or superior performance with 2 tokens per mask, making the efficiency cost of longer representations hard to justify except in applications where extreme spatial precision (beyond SAMTok's ~0.70 reconstruction IoU) is required.

Follow-Up Research This Work Enables

Characterizing the mask reconstruction IoU threshold below which MLLM generation performance degrades. The paper establishes that r-Acc = 0.70 is adequate (g-Acc = 77.6, Table 13) and r-Acc = 0.50 is not (g-Acc = 63.3, Table 12), but does not map the continuous relationship between reconstruction quality and downstream task performance. A study that varies reconstruction IoU systematically — by changing codebook size, quantization steps, or adding controlled noise to codebook vectors — and measures g-Acc on GRES and GCG at each point would produce a minimum-viable-reconstruction curve. This is practically important: it tells tokenizer designers how much reconstruction quality they must preserve, and conversely, how much compression they can tolerate before downstream performance degrades. The experiment requires training multiple SAMTok variants (e.g., 128×2, 256×2, 512×2, 256×3, 128×3) and running the full SFT + evaluation pipeline for each — a significant but well-defined engineering effort.

Training a lightweight difficulty estimator for compute-optimal strategy allocation that costs <10 generations. The paper's headline 4× efficiency gain is computed without accounting for difficulty estimation cost (2048 samples per question). A natural follow-up would close this gap: train a small classifier (e.g., a linear probe on the base LLM's last hidden state, or a lightweight model that takes only the question text as input) to predict the difficulty quintile, using the PRM-based difficulty estimates as pseudo-labels. Evaluate whether this cheap estimator preserves the compute-optimal scaling gains when the difficulty estimation cost is included in the total budget. The key metric is: does cost(estimation) + cost(compute-optimal execution) still beat cost(best-of-N at matched accuracy)? This would convert the paper's aspirational 4× figure into a deployment-ready claim. The experiment requires generating difficulty labels for the 500 test questions (already done), training a classifier, and recomputing the compute-optimal scaling curves with estimation cost amortized.

Combining PRM tree search with the revision model as the proposal distribution, and measuring whether the gains are additive. The paper studies search and revisions independently but acknowledges they were never combined (Section 8). The natural experiment: use the revision model (not the base model) to generate candidates within beam search. At each step of the beam, instead of sampling from the base LLM, condition the revision model on the partial solution history (previous incorrect attempts from that beam) and generate the next step. Evaluate on MATH difficulty bins 3–4 (medium problems), where both search and revisions individually show benefits but neither is dominant. The central question is whether the two mechanisms provide complementary gains (search finds diverse solution strategies, revisions refine within each strategy) or redundant gains (revisions already cover the exploration that search provides). A non-result — additive gains are small — would still be informative, suggesting that the current compute-optimal policy (which selects between search and revisions) is near-optimal. A positive result — e.g., combined performance exceeds the best of either alone at matched budget — would motivate a unified compute-optimal policy over both axes simultaneously.

Testing whether text-only RL rewards are well-aligned with spatial mask quality by correlating token-matching rewards with IoU-based rewards on a held-out set. The paper's RL results (Tables 1 and 3) demonstrate that optimizing for token-matching rewards improves IoU-based metrics, but they do not directly measure the reward alignment: for a large set of model outputs, what is the correlation (e.g., Spearman ρ) between N_TP / max(N_pred, N_gt) and the average IoU of predicted masks? If the correlation is high (ρ > 0.8), then token-matching rewards are a reliable proxy and the approach generalizes. If the correlation is moderate (0.5–0.7), then there are regimes where token-matching rewards are misleading — e.g., where the model outputs the right tokens but the decoded mask is spatially off due to tokenizer reconstruction errors — and future work should investigate reward calibration or hybrid rewards. This experiment requires generating diverse model outputs (by varying temperature, model checkpoint, or RL training stage), decoding each output's masks, computing both token-matching and IoU rewards, and measuring the correlation. It is a diagnostic, not a method, but it would substantially strengthen the paper's central RL claim.

Stress-testing SAMTok on out-of-domain mask distributions to identify failure modes of the 2-token compression. The paper trains SAMTok on 209M masks spanning diverse datasets (indoor, outdoor, UI, part-level, object-level), but all evaluation tasks (GRES, GCG, RefCOCO, MR-RefCOCO, etc.) use natural images from similar distributions. A stress test would evaluate SAMTok's reconstruction fidelity on deliberately challenging masks: extremely thin structures (cracks, wires, text strokes), highly non-convex shapes with many holes (Swiss cheese patterns), heavily occluded objects where the mask is fragmented, and abstract/synthetic shapes that differ from the natural-object training distribution. The question is whether the 2-token representation has enough capacity for these edge cases, or whether reconstruction IoU drops below a usable threshold. If reconstruction degrades sharply on certain mask types, it would motivate hierarchical or variable-length tokenization (e.g., allocating more tokens to complex masks). This experiment requires constructing or curating a test set of edge-case masks and measuring per-category reconstruction IoU.

Replicating the full experimental suite (search scaling, revision scaling, compute-optimal allocation, FLOPs-matched comparison) with a non-PaLM base model on a non-math task to test generalizability. The paper's findings rely entirely on PaLM 2-S* and MATH. A replication with, say, Llama-3 on a code generation benchmark (HumanEval or MBPP, where correctness is verified by unit tests) would answer three critical questions: (1) Do the difficulty-dependent scaling patterns (beam search degrading easy problems, revisions helping easy problems, no method helping the hardest problems) replicate? (2) Do the specific difficulty thresholds and optimal strategy choices transfer, or must they be re-estimated per model/task? (3) Is the FLOPs-matched substitution of test-time compute for pretraining compute also viable in code domains, where the error modes are different (syntax errors vs. logic errors vs. runtime errors)? A failure to replicate — e.g., beam search never outperforms best-of-N on code, or the 4× efficiency gain shrinks to 1.2× — would bound the claims and motivate a conditional theory of when compute-optimal test-time scaling works. This is a large-scale experiment (requiring PRM training, revision model training, strategy sweeps, and FLOPs-matched baselines for a new model/task pair) but is the minimum required to establish the paper's claims as general principles rather than PaLM-on-MATH artifacts.

Practical Applications and Downstream Use Cases

Cost-efficient data labeling for segmentation datasets via LLM-based mask generation with text-only RL verification. The paper's demonstration that mask generation can be optimized through purely text-based rewards (no IoU computation, no mask decoding during RL) enables a lightweight data labeling pipeline: an MLLM with SAMTok generates candidate masks for referring expressions, a text-only reward function scores the outputs, and GRPO iteratively improves the model's mask generation without requiring ground-truth masks. This is relevant for domains where pixel-level annotations are expensive (medical imaging, satellite imagery, industrial inspection) but textual descriptions of regions are available or can be cheaply generated. The concrete benefit: the paper shows that GRPO with text-only rewards on the 3B model improves GRES gIoU from 70.1% to 76.7% (Table 3) and GCG Recall from 47.0% to 53.5% (Table 1) — gains that translate directly to fewer human corrections needed in a human-in-the-loop labeling pipeline. The 2-token representation means the LLM's output is compact enough for rapid iteration.

On-device pixel-wise MLLMs with frozen SAMTok for privacy-sensitive applications. SAMTok's decoupling of the tokenizer from the MLLM means that the heavy mask processing (encoding and decoding through SAM2) can run on-device or on a trusted server, while the MLLM itself can run in the cloud or on a separate device, seeing only the 2-token mask representations — not the image pixels and not the decoded masks. This is relevant for applications like medical image analysis (where a doctor marks a region on a scan and asks "what is this anomaly?"), personal photo organization (segmenting and captioning private images), or enterprise document processing (extracting regions from confidential documents). The concrete benefit: the MLLM never sees the full-resolution image or mask, only the 2-token codes, providing a privacy boundary at the tokenizer level. The paper's results with PerceptionLM (Table 11) show that the same SAMTok works across different MLLM backbones, enabling flexible deployment architectures.

Batch inference for grounded image captioning at scale with interleaved mask generation. The GCG benchmark requires describing an image while simultaneously generating masks for mentioned phrases — producing outputs like "A [<mask_1>] red cup sits next to a [<mask_2>] blue bowl." SAMTok makes this practical at scale because each mask costs only 2 tokens, compared to dozens or hundreds for RLE/polygon/image-based alternatives. In a production system that must caption and ground millions of images (e-commerce catalogs, stock photo libraries, social media content moderation), the token savings compound: generating 5 masks per caption with SAMTok adds 15 tokens (including start/end markers), while an RLE approach might add 50–200+ tokens per mask. At typical LLM inference costs ($1–10 per million output tokens), this difference is meaningful. The concrete numbers: Qwen2.5VL-SAMTok-3B achieves 36.8% AP50 on GCG with two tokens per mask (Table 1), while prior methods with heavier mask representations (LISA at 25.2%, GLaMM at 30.8%) underperform at similar model sizes. For high-throughput batch captioning, the efficiency gain from compact mask representation directly reduces cost.

Interactive segmentation as a chatbot feature without specialized backends. Multi-round interactive segmentation (Table 2, MR-RefCOCO and MR-PACO) requires the model to understand previously segmented regions (input masks), reason about spatial relationships, and generate new masks — a task that traditionally required coordinating a segmentation backend with an LLM. SAMTok enables this entirely within the LLM's text interface: previous masks are tokenized into the input prompt, the LLM reasons about them in text space, and outputs new mask tokens. This means an interactive segmentation feature can be built as a standard chatbot with a custom vocabulary — no separate segmentation server, no API orchestration, no passing decoded masks between services. The concrete numbers: Qwen2.5VL-SAMTok-3B achieves 83.3% cIoU on MR-RefCOCO (average across rounds 2–6), outperforming the 7B SegLLM expert model (79.7%) by 3.6 percentage points (Table 2). For a product team building a visual assistant, this means pixel-wise capabilities can be added in days (vocabulary extension + fine-tuning on SAMTok-formatted data) rather than months (designing and integrating a segmentation backend).

When to Prefer This Method

SAMTok is explicitly positioned against two alternatives: (1) decoder-integrated pixel-wise MLLMs (LISA, Sa2VA, OMG-LLaVA, HyperSeg) that use continuous special tokens + dedicated segmentation decoders, and (2) text-output approaches that represent masks as coordinates, polygons, RLE, or autoregressively generated mask images. The paper also implicitly contrasts with approaches that use bounding boxes or points as intermediate representations for RL.

Prefer SAMTok + standard MLLM fine-tuning when:

• You need to add pixel-wise capabilities to an existing base MLLM and want to avoid architectural modifications — SAMTok requires only vocabulary extension and data reformatting (Section 2.2).
• You want to apply reinforcement learning (GRPO, PPO, DPO) to mask generation tasks, since discrete tokens enable text-only reward functions (Equation 5, with demonstrated gains of +6.8% gIoU on GRES and +4.5% AP50 on GCG from RL, Tables 3 and 1).
• You need to handle tasks with multiple masks in input and output (interactive segmentation, grounded conversation generation) where token efficiency matters — SAMTok uses 2 tokens per mask versus dozens/hundreds for coordinate/polygon/RLE/image-based alternatives (Section 1, fourth challenge).
• You anticipate switching base MLLMs (e.g., upgrading from Qwen2.5-VL to Qwen3-VL to a future model), since SAMTok is frozen and decoupled — the paper shows it works across Qwen2.5-VL, Qwen3-VL, and PerceptionLM with minimal performance variation (Table 11).
• Inference token budgets are tight (limited context windows, cost-sensitive deployments), where the 2-token representation's efficiency advantage over longer mask encodings directly translates to lower cost and latency.

Prefer decoder-integrated approaches (LISA, Sa2VA, etc.) when:

• Mask reconstruction fidelity above SAMTok's ~0.70 IoU is critical for the application — e.g., medical image segmentation where boundary precision at the pixel level matters for clinical decisions, or fine-grained part segmentation where thin structures are important. Decoder-integrated methods can achieve higher mask IoU by directly optimizing segmentation losses on the MLLM's outputs, without the quantization bottleneck.
• The application requires only single-mask generation (not interleaved text-mask generation or interactive segmentation), and the additional engineering complexity of a specialized decoder is acceptable — in this narrow setting, the architectural cost of a decoder is less burdensome relative to the benefit.
• The base MLLM cannot be fine-tuned (e.g., API-only access to a proprietary model), and the mask decoder can be trained as a lightweight head on frozen LLM features — SAMTok requires fine-tuning the MLLM to learn the mask token vocabulary, which is infeasible without model weight access.

Prefer coordinate/polygon/RLE text-output approaches when:

• The masks are extremely simple (axis-aligned bounding boxes, convex polygons with few vertices) and the precision loss from coordinate quantization is acceptable — in this case, the overhead of training SAMTok is unnecessary.
• You cannot afford the 209M-mask pretraining cost for SAMTok and the downstream tasks involve only coarse localization (object detection, not instance segmentation), where bounding box coordinates suffice.

The key tradeoff SAMTok makes explicit: mask reconstruction fidelity (~0.70 IoU, from Table 13) versus architectural simplicity (no decoder, no segmentation losses, RL-compatible discrete tokens). For most applications, the paper demonstrates that 0.70 reconstruction fidelity is adequate for SOTA or near-SOTA downstream task performance (Tables 1–10) — but applications requiring near-pixel-perfect mask boundaries may find this fidelity ceiling too low.