Apple × CMU: 10× Faster Agent Reasoning by Turning Generation into Multiple Choice

In typical search-agent workflows following the ReAct paradigm, a model generates free-form reasoning before each action — analyzing observations, evaluating evidence, and deciding the next step. This thinking process is central to ReAct but incurs substantial inference overhead, especially for small on-device models (2B–4B parameters) that often produce lengthy yet unhelpful reasoning traces.

Researchers from Apple and Carnegie Mellon University observed that despite diverse entities, queries, and observations across tasks, the high-level reasoning logic that determines the next action recurs repeatedly. For example, “Where is an AI conference held?” and “When was LEGO founded?” share the same reasoning pattern: entity known, target attribute missing → search with entity name plus target attribute. This insight suggests the reasoning space need not cover all natural language; instead, a compact set of reusable reasoning options can be predefined.

Method: Selection-Based Structured Reasoning (SSR)

2.1 Reasoning Library and Selection-Based Reasoning

SSR defines a reasoning library of six options:

Identify entity in image via image-to-image search when uncertain

Query text search for specific facts

Crop or zoom image region to inspect details

Answer based on collected evidence

Read answer directly from image when visible

Answer directly from common knowledge

At each interaction step, the agent selects one option from the library as the current reasoning, then generates the concrete action (search query, crop coordinates, etc.) conditioned on that selection. This decouples reusable rationale (“why”) from instance-specific execution details (“how”).

2.2 Model as Its Own Classifier

Selection is performed by the LLM itself without extra classification heads. For each reasoning option, the model computes the average token probability under the current context (teacher forcing), then applies softmax to obtain a categorical distribution over the library. Sampling from this distribution completes the reasoning choice. The same model parameters and system prompt are shared for scoring and action generation, leveraging the base model’s semantic priors as a context-aware selector.

2.3 Parallel Reasoning Decoding

Because all options are predefined, scoring is fully parallelizable. At the token level, the entire option sequence is prefilled in a single batched forward pass via teacher forcing, computing all token likelihoods simultaneously. At the option level, all options share the same context KV cache and are scored concurrently. The model performs zero autoregressive token generation for reasoning.

A concrete illustration: given a photo of a University of Edinburgh building and the task “Find the exhibition named ‘People Experiencing Profound Change,’” a 4B model using free-form thinking generated 156 tokens in 1821.92 ms, misidentified the building as the National Museum of Scotland, and pursued a wrong search direction. SSR selected “identify entity via image search” in 32.66 ms and directly triggered the correct image search action.

2.4 Training Integration

SSR plugs into existing SFT and RL pipelines. A key advantage: because reasoning is constrained to the library, RL can be applied directly without an SFT warm-up phase that many prior works require to teach a specific thinking format. This eliminates SFT data construction and training time while making model behavior more predictable and controllable.

Experiments: Comparable Accuracy, 10× Speedup

Evaluated on Qwen3-VL 2B/4B with H100 + SGLang across seven multimodal search benchmarks (MMSearch, HR-MMSearch, FVQA-test, InfoSeek, SimpleVQA, LiveVQA, MAT-Search):

Accuracy on par: 4B SSR (GRPO) averaged 61.37% success rate, matching same-size SOTA; 2B SSR reached 51.26%, comparable to zero-shot 8B models. Across GRPO, GSPO, SAPO, and SFT, SSR matched free-form thinking success rates.

Large latency reduction: Per-step reasoning latency dropped >90%; total per-question model latency fell 28–54%. Versus baselines (Chain-of-Draft, Sketch-of-Thought, efficiency-reward RL, Probe & Prefill), SSR achieved highest success rate with p95 reasoning latency of only 0.07 s (baselines 0.8–1.9 s) and greater stability.

Ablation Findings

Semantics matter: Forcing the model to select only an index (e.g., “[1]”) instead of full option text reduced average success rate by 3.9% across all seven benchmarks, confirming that natural-language scoring uses the model’s semantic understanding for context matching, not merely index classification.

Library size yields monotonic gains: Expanding from 1 generic option to 6 options raised average success rate from 37.24% to 61.37%. The jump from 1 to 2 options (distinguishing “answer” vs “call tool”) alone contributed a 17.27-point leap, showing finer-grained options better guide action generation.

Selection is adaptive: The same checkpoint exhibited different option distributions per benchmark: “answer from image” dominated on factoid-heavy SimpleVQA, while “inspect image details” prevailed on high-resolution HR-MMSearch.

vs. autoregressive generation: Training a model to autoregressively generate library options required SFT to enforce format; RL with only prompt and format rewards failed to prevent hallucinated reasoning. SSR mechanically guarantees reasoning stays within the predefined library.

Conclusion: Redefining the Reasoning Space

The authors argue the contribution goes beyond shortening chain-of-thought. It raises a foundational question: must reasoning representation rely on unconstrained natural-language generation? For open-ended math, theorem proving, or code design, free-form thinking’s expressiveness is essential. But in multi-step search and tool-use settings, each step typically faces only a few high-level decisions: is current information sufficient? identify entity? continue search? zoom region? The truly input-dependent parts — search queries, tool arguments, crop coordinates, final answers — remain free-generation; the strategic reasoning layer need not be regenerated in the full language space every time.

SSR thus reorganizes the reasoning space: from a fully open language space to a compact, reusable, semantically explicit policy space. The model does not stop thinking; it replaces a portion of token-by-token generation with selection and inference in a structured reasoning space.