SSR

Selection-Based Structured Reasoning

Toward Efficient Multimodal Search Agents

Feiyu Gavin Zhu1,2Xiaoyu Zhu1,‡,*Jiqi Yang1Rui Yang1Arnab Kumar Mondal1Yancheng Wang1Xinke Deng1Jean Oh2Reid Simmons2Joerg Liebelt1Xiang Kong1Zhongyu Jiang1,*

1 Apple2 Carnegie Mellon University

feiyuz@cs.cmu.edu

‡ Project Lead   ·   * Corresponding Author

Animated comparison: freeform agents generate reasoning token by token, while SSR scores reusable reasoning candidates in parallel, selects a candidate, and generates the next action.
Reasoning by selection. Score reusable reasoning in parallel, then generate a context-specific action.
>90%lower reasoning latency per turn
28–54%lower model latency per question
61.37%average success · 4B · 7 benchmarks

Abstract

Multimodal agents commonly generate freeform reasoning before each action. For small models, limited model capacity can result in lengthy reasoning that provides little useful guidance for action generation while incurring substantial inference cost. We introduce Selection-based Structured Reasoning (SSR), a framework that reformulates reasoning as selection among pre-specified, reusable natural-language candidates. At each turn, the model selects a candidate based on its likelihood given the current context, without an auxiliary task head. Teacher-forced prefilling computes token likelihoods concurrently within and across candidates using a shared context KV cache. Across seven multimodal search benchmarks, 2B and 4B models, and multiple reinforcement learning objectives and supervised fine-tuning, SSR achieves competitive task performance while reducing per-turn reasoning latency by over 90% and total per-question model inference latency by 28–54%.

Selection-Based Structured Reasoning

A small library of reusable reasoning. A context-specific action at every turn.

Search questions differ in their entities, images, and answers, but often share the same information needs: identify an unfamiliar object, look up a missing fact, inspect an image detail, or answer from the available evidence. SSR represents these recurring decisions as reusable natural-language reasoning candidates.

At each turn, the language model scores every candidate given its current interaction history. The selected reasoning is inserted into the context to guide action generation. The action remains specific to the question: the model generates the search query, crop coordinates, or final answer.

SSR overview: the shared interaction history is prefilled, reusable reasoning candidates are scored in parallel, and a selected candidate guides action generation.
Parallel reasoning decoding scores all candidate tokens through teacher-forced prefill, with a shared history KV cache.

1. Score in parallel

Teacher-force the full candidate texts. Compute token likelihoods concurrently within each candidate and across candidates.

2. Select the reasoning

Use length-normalized likelihoods to sample a candidate, or choose the highest-scoring candidate at inference.

3. Generate and act

Append the selected text, generate an action, and execute it. Add the observation to the history and repeat until an answer or stopping condition.

Parallel reasoning decoding

For a candidate $r_i$ with $K_i$ tokens, SSR uses its length-normalized conditional log-likelihood:

$$S_\theta(r_i\mid h_t)=\frac{1}{K_i^\alpha}\sum_{j=1}^{K_i}\log p_\theta(r_{i,j}\mid h_t,r_{i,<j})$$

We use $\alpha=1$ by default. For selection temperature $\tau>0$, the scores define a categorical policy:

$$q_\theta(r_i\mid h_t)=\frac{\exp(S_\theta(r_i\mid h_t)/\tau)}{\sum_{k=1}^N\exp(S_\theta(r_k\mid h_t)/\tau)}$$

The model then generates the action from $p_\theta(\cdot\mid h_t,r_i)$. Scoring and action generation share the same parameters and system prompt; no auxiliary classification head is needed.

Pre-specified candidate tokens remove the sequential dependency of generating a freeform reasoning trace. With fixed model dimensions and idealized parallel reductions, the reasoning computation has depth

$$\begin{aligned}\mathcal D_{\mathrm{freeform}}&=O\!\left(L\log(H+L)\right)\\\mathcal D_{\mathrm{SSR}}&=O\!\left(\log(H+K_{\max})+\log N\right)\end{aligned}$$

Here, $H$ is the history length, $L$ the generated freeform reasoning length, $N$ the candidate count, and $K_{\max}$ the longest candidate. This characterizes parallel dependency depth; measured latency is reported below.

One reasoning library, seven benchmarks

The experiments use six manually written candidates shared across all benchmarks. Each describes a general information need and the corresponding next step.

Hover over an entry or tap it to read the full reasoning candidate.

Joint training for selection and action

SSR supports both supervised fine-tuning and reinforcement learning. SFT combines a categorical loss for selecting the target reasoning with the token-level action loss:

$$\mathcal L_{\mathrm{SFT}}=-\sum_t\left[\lambda^{\mathrm{SFT}}_{\mathrm{reason}}\log q_\theta(z_t\mid h_t)+\frac{1}{|a_t|}\sum_k\log p_\theta(a_{t,k}\mid h_t,z_t,a_{t,<k})\right]$$

For GRPO-style training, the inserted reasoning contributes one categorical importance ratio, while generated action tokens retain their usual token-level ratios:

$$\rho_t^{\mathrm{reason}}=\frac{q_\theta(z_t\mid h_t)}{q_{\mathrm{old}}(z_t\mid h_t)}\qquad \rho_{t,k}^{\mathrm{action}}=\frac{p_\theta(a_{t,k}\mid h_t,z_t,a_{t,<k})}{p_{\mathrm{old}}(a_{t,k}\mid h_t,z_t,a_{t,<k})}$$

Group-normalized rewards provide the advantages. The clipped objective weights reasoning selection by $\gamma^{\mathrm{RL}}_{\mathrm{reason}}$ and applies reference-policy KL regularization to generated action tokens. SSR can be trained directly with RL without an SFT warm start, and is also compatible with GSPO and SAPO. Algorithm 2 specifies the grouped GRPO update with detached competitor scores.

Competitive Performance, Lower Latency

We evaluate Qwen3-VL 2B and 4B agents on seven multimodal search benchmarks with a shared harness for reverse-image search, text search, and image cropping. Training uses VisualProbe, DeepEyes, and FVQA-train, with outcome and format rewards for RL.

Multimodal search performance

SSR reaches 61.37% average success at 4B, comparable to the strongest reported trained 4B baseline in the paper (61.25%). At 2B, SSR reaches 51.26%, close to the zero-shot Qwen3-VL-8B agent’s 51.33%.

ModelSizeMMSearchHR-MMSearchFVQA-testSimpleVQALiveVQAMAT-SearchInfoSeekAverage
MMSearch-R1*4B53.8025.5067.2068.5050.6078.0060.0057.66
SenseNova-MARS*4B44.0524.6757.1965.5543.0678.6748.4051.65
TAPO + GRPO4B54.9733.1165.5668.3154.0582.0058.8559.55
TAPO + GSPO4B60.8231.4866.6170.9856.1483.3359.4061.25
TAPO + SAPO4B56.7331.4867.1771.2755.3681.3359.9560.47
SSR (ours)4B61.4030.1670.2271.0857.1978.6760.9061.37
SSR (ours)2B51.1817.0661.1661.8351.5462.5953.4451.26

Trained 2B/4B agents. All models were trained for one epoch. Success rate (%), averaged across seven benchmarks. * Reproduced results. The highlighted rows are SSR.

Consistent efficiency across training objectives

SSR reduces mean reasoning latency by over 90% across all tested objectives. It improves average success over freeform GRPO at both model sizes, matches SFT, and shows modest decreases under GSPO and SAPO.

Training objectiveSizeSuccess (%) ↑Reasoning Latency / Turn
(s) ↓
Model Latency / Question
(s) ↓
Throughput (tokens/s) ↑
GRPO · freeform4B58.650.8965.45672.7
GRPO · SSR4B61.37 (+4.6%)0.061 (-93.2%)2.514 (-53.9%)3029.9 (+4070.2%)
GSPO · freeform4B60.450.9395.19571.9
GSPO · SSR4B58.60 (-3.1%)0.063 (-93.3%)3.706 (-28.7%)2526.2 (+3414.6%)
SAPO · freeform4B61.261.1136.14569.4
SAPO · SSR4B60.46 (-1.3%)0.060 (-94.6%)2.862 (-53.4%)3215.3 (+4535.0%)
SFT · freeform4B54.550.7705.30870.9
SFT · SSR4B54.57 (+0.0%)0.062 (-91.9%)3.318 (-37.5%)2567.0 (+3520.5%)
GRPO · freeform2B49.980.9024.15387.2
GRPO · SSR2B51.26 (+2.6%)0.053 (-94.1%)2.125 (-48.8%)3302.6 (+3687.9%)

Parenthetical percentages show relative changes from the paired freeform model. Success is measured on the full evaluation sets. Latency profiling uses 100 sampled questions per benchmark on an NVIDIA H100 with SGLang. Model latency includes prefill, reasoning, and action generation across the full trajectory; tool execution and judge time are excluded. Effective throughput counts all scored candidate tokens for SSR and generated reasoning tokens for freeform.

Faster and more predictable reasoning

On the profiling subset, SSR achieves the highest success rate among the compared methods. Mean reasoning latency is 61 ms and the 95th percentile is 71 ms. Mean model latency is 2.514 seconds per question, about 28% below Sketch-of-Thought, the fastest baseline on this metric.

MethodSubset success (%) ↑Reasoning Latency / Turn (s) ↓Model Latency / Question (s) ↓Throughput (tokens/s) ↑
Meanp95Meanp95
MMSearch-R1-4B51.591.2351.8925.6017.50273.8
SenseNova-MARS-4B49.570.9601.5335.28910.39772.5
Chain-of-Draft54.060.4160.7983.6576.61669.9
Sketch-of-Thought56.070.5431.0543.4877.04871.7
Efficiency-reward RL56.280.7561.2734.6388.67568.9
Probe & Prefill55.490.8901.5645.21911.36965.8
SSR (ours)58.360.0610.0712.5144.6203029.9

All success rates in this table refer to the profiling subset, rather than the full evaluation sets. Latencies are in seconds; p95 is the 95th percentile.

What Makes Selection Work?

Score the meaning of the whole candidate

Scoring full reasoning text achieves 61.37% average success, compared with 58.98% when scoring only a candidate index. Both variants insert the same selected text before action generation; the difference is whether candidate selection uses the reasoning tokens themselves.

Autoregressive generation can reproduce library entries after an SFT warm start, but retains sequential reasoning generation. Without SFT, the autoregressive variant produces freeform reasoning despite the prompt and format reward. SSR enforces the reasoning library through selection.

Reasoning methodMMSearchHR-MMSearchFVQA-testSimpleVQALiveVQAMAT-SearchInfoSeekAverage
Entire reasoning61.4030.1670.2271.0857.1978.6760.9061.37
Index only60.8222.3067.5669.9955.5677.3359.3058.98
Autoregressive · with SFT61.4029.8469.2269.1056.9772.0061.5560.01
Autoregressive · without SFT60.2333.1167.2868.4154.6779.3358.6060.23

A small library supports different information needs

Average success rises from 37.24% with a single generic candidate to 61.37% with six candidates. Moving from one generic candidate to two choices—answer or use a tool—already yields a 17.27 percentage point gain.

Library size ablation: average success increases from 37.24 percent with one candidate to 61.37 percent with six.
More specific reasoning choices improve average success.
Reasoning selection distributions vary across the seven benchmarks, using the same checkpoint and reasoning library.
The same checkpoint uses different reasoning-selection distributions across benchmarks.

SimpleVQA most often selects “Answer from image,” while HR-MMSearch has the highest selection rate for “Examine image detail.” The generated action type matches the selected candidate’s expected action type in 99.9% of turns. These measurements show operational consistency and different selection patterns across tasks.

Pseudocode

Inference with SGLang, and grouped GRPO training.

Inference details

The running context $h_t$ contains the complete interaction before turn $t+1$. Each candidate is a nonempty token sequence including its reasoning delimiters; inputs include the required chat delimiters and assistant prefixes.

Action sampling parameters $\eta$ specify the token limit (e.g., max_new_tokens=8192), temperature, and stopping conditions, with one completion per input. The selection temperature $\tau$ is independent of these action settings.

While $h_t$ stores tokens, the persistent SGLang engine owns the KV cache of computed model states. KV tensors are not endpoint arguments or return values. New observations enter the cache when the next model call processes them. For multimodal inputs, the adapter supplies matching image_data per request row and tracks expanded token positions; those payloads are omitted below.

SGLang accepts variable-length candidate sequences without caller padding. Setting logprob_start_len to $L_t-1$ leaves the final $K_i$ input-token log-probabilities aligned with candidate $r_i$. Candidates use causal attention and do not attend to one another.

Algorithm 1 · SSR inference using SGLang’s generate interface

Require: Tokenized system prompt s and initial question q; reasoning library \mathcal R=\{r_1,\ldots,r_N\}; persistent SGLang engine \pi_\theta with prefix caching; harness \mathcal H; normalization exponent \alpha\in[0,1]; selection temperature \tau\geq0; action sampling parameters \eta; turn limit T_{\max}\geq1

Output: (h,\delta): final token context and stopping status

Conventions: \Vert concatenates tokens; response indexing is zero-based; [-K:] takes the last K entries. Each generate call reuses and updates the engine-owned KV cache \mathcal C_{\mathrm{KV}}.

Validity: A final answer must be parseable and nonempty; a tool request must name an allowed tool and satisfy its argument schema, according to \mathcal H.

  1. h_0\gets s\Vert q; K_i\gets|r_i| for i=1,\ldots,N \triangleright context tokens; candidate lengths
  2. g_0\gets\{\texttt{"max\_new\_tokens"}:0\} \triangleright prefill-only sampling parameters
  3. \begin{aligned} \pi_\theta.\texttt{generate}(&\texttt{input\_ids}=s,\\ &\texttt{sampling\_params}=g_0,\\ &\texttt{return\_logprob}=\texttt{False}) \qquad \triangleright\text{ warm the system-prompt cache} \end{aligned}
  4. for t=0,\ldots,T_{\max}-1 do
  5. L_t\gets|h_t|; X\gets[h_t\Vert r_i]_{i=1}^N \triangleright N variable-length token sequences
  6. \begin{aligned} Y\gets\pi_\theta.\texttt{generate}(&\texttt{input\_ids}=X,\\ &\texttt{sampling\_params}=g_0,\\ &\texttt{return\_logprob}=\texttt{True},\\ &\texttt{logprob\_start\_len}=L_t-1) \qquad \triangleright\text{ list of }N\text{ responses} \end{aligned}
  7. for i=1,\ldots,N do
  8. D_i\gets Y[i-1].\texttt{input\_token\_logprob}[-K_i:]
  9. Verify [d[1]\text{ for }d\in D_i]=r_i \triangleright record: (log-probability, token ID, text)
  10. \ell_i\gets[d[0]\text{ for }d\in D_i]; verify all entries are finite \triangleright |\ell_i|=K_i
  11. S_{t,i}\gets K_i^{-\alpha}\sum_{j=0}^{K_i-1}\ell_i[j] \triangleright candidate scores S_t\in\mathbb R^N
  12. end for
  13. if \tau=0 then
  14. i_t\gets\min\arg\max_i S_{t,i}
  15. else
  16. i_t\sim\operatorname{Categorical}(\operatorname{softmax}(S_t/\tau))
  17. end if
  18. z_t\gets r_{i_t}; u_t\gets h_t\Vert z_t \triangleright append selected reasoning to current context
  19. \begin{aligned} A\gets\pi_\theta.\texttt{generate}(&\texttt{input\_ids}=u_t,\\ &\texttt{sampling\_params}=\eta,\\ &\texttt{return\_logprob}=\texttt{False}) \end{aligned}
  20. a_t\gets A.\texttt{output\_ids}; h_t^+\gets u_t\Vert a_t \triangleright action tokens; extended context
  21. if a_t is a valid final answer then
  22. return (h_t^+,\texttt{answer})
  23. else if a_t is not a valid tool request then
  24. return (h_t^+,\texttt{invalid-output})
  25. end if
  26. (o_t,\delta_t)\gets\mathcal H.\texttt{execute}(a_t) \triangleright observation tokens; continue or error
  27. h_{t+1}\gets h_t^+\Vert o_t \triangleright observation includes the next assistant prefix
  28. if \delta_t\neq\texttt{continue} then
  29. return (h_{t+1},\delta_t)
  30. end if
  31. end for
  32. return (h_{T_{\max}},\texttt{turn-limit})
Grouped GRPO training details

Each iteration collects a fresh minibatch of rollout groups and takes one optimizer step. Training uses clipping threshold $\epsilon>0$, reasoning-selection weight $\gamma^{\mathrm{RL}}_{\mathrm{reason}}\geq0$, reference-policy KL weight $\beta\geq0$, and reward-standard-deviation floor $\delta>0$.

During rollout logging, set return_logprob=True on Algorithm 1’s action-generation request. Logged behavior probabilities and current action likelihoods must use the same sampling distribution. Each trajectory retains its turns, final context, observations, media references, and stopping status, allowing $f_{\mathrm{reward}}$ to score unsuccessful and truncated trajectories.

Algorithm 2 · Grouped GRPO with detached competitor scores

Require: Parameters \theta; frozen reference p_{\mathrm{ref}}; system prompt s; question distribution \mathcal D; library \mathcal R; harness \mathcal H; reward function f_{\mathrm{reward}}; questions per minibatch m\geq1; group size n\geq2; rollout settings \alpha, \tau>0, \eta, T_{\max}; training hyperparameters \epsilon, \gamma^{\mathrm{RL}}_{\mathrm{reason}}, \beta, \delta; optimizer \mathcal O

Output: Trained parameters \theta

Conventions: \Vert concatenates tokens; arrays are zero-based and candidate indices one-based. Rollout records are detached. N=|\mathcal R|; \operatorname{LSE}(u)=\log\sum_i e^{u[i]}.

Surrogate: \ell_\epsilon(\rho,A)=\min\{\rho A,\operatorname{clip}(\rho,1-\epsilon,1+\epsilon)A\}.

  1. while the training budget remains do
  2. \theta_{\mathrm{old}}\gets\operatorname{stopgrad}(\operatorname{copy}(\theta))
  3. Synchronize \pi_{\mathrm{old}} to \theta_{\mathrm{old}}; invalidate its weight-dependent caches
  4. Sample m questions X\sim\mathcal D^m; \mathcal M\gets[\,] \triangleright current rollout minibatch
  5. for each question x\in X do
  6. for g=1,\ldots,n do
  7. Collect d_g with Algorithm 1 using \pi_{\mathrm{old}} and question x, with turn logging
  8. R_g\gets f_{\mathrm{reward}}(x,d_g) \triangleright reward, including non-answer stopping cases
  9. end for
  10. \mu\gets n^{-1}\sum_g R_g; \sigma\gets\sqrt{n^{-1}\sum_g(R_g-\mu)^2}
  11. for g=1,\ldots,n do
  12. d_g.\texttt{adv}\gets\operatorname{stopgrad}((R_g-\mu)/\max(\sigma,\delta)); append d_g to \mathcal M
  13. end for
  14. end for
  15. J\gets0 \triangleright accumulated surrogate; |\mathcal M|=mn trajectories
  16. for each trajectory d\in\mathcal M do
  17. for each turn v\in d.\texttt{turns} do
  18. (h,z,a,i)\gets(v.\texttt{h},v.\texttt{z},v.\texttt{a},v.\texttt{index}); K\gets|z|; w\gets z\Vert a
  19. b\gets[\log p_\theta(w[k]\mid h\Vert w[:k])]_{k=0}^{|w|-1} \triangleright teacher forcing; gradients enabled
  20. s_{\mathrm{sel}}\gets K^{-\alpha}\sum_{j=0}^{K-1}b[j] \triangleright current selected-candidate score
  21. u\gets\operatorname{stopgrad}(v.\texttt{scores}/\tau); u[i-1]\gets s_{\mathrm{sel}}/\tau \triangleright N selector logits
  22. \widetilde b^{\mathrm{reason}}\gets u[i-1]-\operatorname{LSE}(u)
  23. \rho^{\mathrm{reason}}\gets\exp(\widetilde b^{\mathrm{reason}}-v.\texttt{logq}) \triangleright one scalar ratio
  24. \rho^{\mathrm{action}}\gets\exp(b[K:]-v.\texttt{logp}) \triangleright |a| token ratios, elementwise
  25. J\gets J+\gamma^{\mathrm{RL}}_{\mathrm{reason}}\ell_\epsilon(\rho^{\mathrm{reason}},d.\texttt{adv})+\sum_k\ell_\epsilon(\rho^{\mathrm{action}}[k],d.\texttt{adv})
  26. end for
  27. end for
  28. \widehat{\mathcal K}\gets reference-policy KL estimate on generated action positions in \mathcal M
  29. \mathcal L\gets-J/|\mathcal M|+\beta\widehat{\mathcal K} \triangleright scalar loss; trajectory average
  30. \mathcal O.\texttt{zero\_grad}(); \mathcal L.\texttt{backward}(); \mathcal O.\texttt{step}() \triangleright one update from \mathcal M
  31. end while
  32. return \theta
Turn records and differentiable updates

Each turn stores the context h, selected reasoning z, action a, candidate index, and rollout candidate scores. It also stores the selected-candidate log-probability

$$\displaystyle v.\texttt{logq}=\frac{S^{\mathrm{old}}_{t,i_t}}{\tau}-\operatorname{LSE}\!\left(\frac{S_t^{\mathrm{old}}}{\tau}\right)$$

and the action-token log-probabilities logp from output_token_logprobs, in token order. Verify their token IDs against the action and retain exactly $|a_t|$ finite values.

Recompute likelihoods using the differentiable model $p_\theta$; SGLang rollout outputs are detached. Only the selected candidate’s score receives gradients, through both the softmax numerator and denominator. Teacher-forced turns may be batched with causal attention and true-length masks. Action losses and reference-policy KL apply only to generated action positions; prompts, inserted reasoning, and observations remain visible as context. Exclude padding from attention and loss.

Reasoning in a Reusable Space

Small search agents may not need to reason over the full language space at every step. By organizing recurring reasoning into a compact library and selecting through parallel scoring, SSR achieves competitive multimodal search performance with substantially lower model inference latency.