Output-Space Search: Targeting LLM Generations in a Frozen Encoder-Defined Output Space

AI summary

Overview

Research area: Natural Language Processing — controllable text and code generation, reinforcement learning from sequence-level objectives, and black-box optimization over LLM outputs.

Technical level: Advanced. The paper assumes familiarity with autoregressive decoding, encoder embeddings, PCA and Varimax rotation, group-relative policy optimization (GRPO/GSPO), QLoRA fine-tuning, and Bayesian optimization with Gaussian-process surrogates.

One-sentence scope: The paper proposes and evaluates "Output-Space Search" (OS-Search), a method that trains an LLM to generate outputs whose coordinates land near a requested point in a frozen 3-dimensional encoder-defined output space, so that generation becomes a search over endpoints rather than over token paths.

What This Paper Is About

LLM workflows often need either many different plausible completions or completions that score well under an external evaluator. Current practice — repeated token-space sampling, prompt chaining, or rejection/reranking — creates sequential dependence between samples and growing context, and offers no stable low-dimensional interface that an outer loop can optimize directly. The paper asks whether an autoregressive LLM can expose a single "state-like" target point that selects where an output should land, so that an external loop can sweep or optimize that target while decoding remains standard.

Key Contributions

1. The OS-Search interface itself: a fixed, external, encoder-defined output space Z (with d_z = 3) plus a target-conditioned controller (z*-conditioned policy) trained with sequence-level reinforcement learning to approximately hit requested coordinates and to self-report its own realised coordinates (ẑ_S).

2. Story-domain result (anchored Z_text): an anchored space whose first axis z_1 yields a stable reference axis, validated by a monotonic relationship between the requested target z_1* and an automated templatedness proxy (EQ-Bench Slop-Score). Sweeping a small grid in Z_text produces large embarrassingly-parallel multi-branch diversity gains.

3. Code-domain result (unanchored Z_code): a space that supports black-box optimization — Bayesian optimization over z* improves a withheld executable objective (CA++) under matched valid-program budgets while preserving validity.

4. A capability contrast (Table 1) between OS-Search and direct trajectory/program search: OS-Search searches low-dimensional continuous state-like targets (rather than high-dimensional discrete token paths or program edits), uses the evaluator only in the outer loop where it can be withheld, learns an objective-agnostic actuator, and produces conditionally independent samples per z* that can be generated in parallel.

Main Findings

• Target tracking improves precision with best-of-K. Success–tolerance curves (Fig. 3) show that best-of-5 improves precision over one-shot sampling, with the gap between early checkpoints (first 5% of the RL run) and the best checkpoint largely reflecting learning from RL. Malformed or invalid outputs count as failures.

• Retrieval grounding is essential for code control. Removing exemplars or providing mismatched exemplars collapses 3D control (best-of-5 ≤ 0.006) and drops the fraction of requests yielding any valid sample (ValidReq = 35.3%). Omitting the explicit REQUESTED TARGET line has little effect (Bo1 of 0.658 versus 0.664 with the default prompt), suggesting most code conditioning is carried by exemplar selection rather than numeric parsing.

• Calibration is tighter in code than in stories. Calibration plots with targets spanning quartile q_0.01 to above q_0.99 show code tracking within ±0.05, while stories are looser (Fig. 8, App. E).

• Story diversity is substantially higher than prompt chaining. Under a fixed 15-branch budget (B = 15) on WritingPrompts, OS-Search's Z-grid sweep (2 × 3 × 3) reaches LLMScore 57.15 with Degen 0.09, versus the strongest sampling baseline (chained top-p sweep, best T = 0.6, p = 0.95, k = 20) at LLMScore 18.45 with Degen 0.07. The diversity ratio is ρ_div(OS-Search) = 57.15 / 18.45 ≈ 3.10, described in the abstract as 3.1× higher LLM-scored diversity than prompt-chaining.

• Overlap metrics fall accordingly. OS-Search records Self-BLEU 8.01, ROUGE-L 23.27, METEOR 28.40, and Sent-Sim 48.17 — lower than every path-based baseline, where lower is more diverse. Chained greedy (T = 0) records 78.90 / 83.26 / 85.41 / 92.36; chained temperature sweep (best T = 0.9, p = 0.95, k = 20) 58.45 / 66.20 / 70.31 / 87.89; chained top-p sweep (best T = 0.6, p = 0.95, k = 20) 44.74 / 54.37 / 59.83 / 84.25; chained top-k sweep (best T = 0.6, p = 0.95, k = 40) 63.96 / 71.05 / 74.82 / 89.38.

• The anchored story axis is validated as a diagnostic. For final-10% 1D-control rollouts (S = {1}, N ≈ 7,000), Slop-Score correlates with the requested target z_1* at Pearson r ≈ 0.57 (p < 10⁻³), with roughly a 20–30 point upward shift over the target range and substantial within-setting variability. Correlating Slop-Score with realised coordinates (z_1, z_2, z_3) on the same samples shows only z_1 has a substantial association (r ≈ 0.56).

• Bayesian optimization over z* improves a withheld objective. Warm-started BO (GP-EI) improves best-so-far CA++ score up to 0.395 under a matched N_ok budget, compared to 0.371 for the 188-program library used to fit Z_code and above path-based sampling. A larger precomputed pool of random-z* controller rollouts (N_ok = 13,752) reached a best of 0.381 and covered a broader region of CA behavioral space, including high-score pockets beyond the library frontier.

• Updating the surrogate at realised coordinates beats updating at the request. In the surrogate-update ablation, updating at the realised coordinates z(x_best) of the best valid program per query outperforms updating at the requested target z*.

Methodology in Plain English

The approach separates three pieces and keeps them modular.

1. Representation. For each domain, outputs are mapped to fixed 3D coordinates. A frozen encoder embeds the task content (for code, only the extracted make_seed() function with comments stripped), and a frozen linear projection — fitted once as PCA followed by an orthogonal Varimax rotation within the PCA subspace — maps that embedding to a 3-dimensional coordinate z(x) = Uᵀ(E(x) − μ). Story coordinates are fitted on a corpus of public English story datasets (including the STORIES corpus), filtered to roughly 800–25,000 characters, and the first axis is anchored toward Qwen3-generated "default-style" stories relative to the corpus mean. Code coordinates are fitted on a task-specific library of N = 188 valid make_seed() implementations. The encoder, mean, and projection are then frozen so RL only changes the policy.

2. Actuation. A controller is prompted with the base prompt, retrieved exemplars from a nearest-neighbour library in Z (two near z* and one near −z*, used to ground otherwise-arbitrary numeric coordinates), and the numeric request z*_S. This design keeps inference as standard autoregressive sampling — no decoding-time guidance. The model fine-tuned is Qwen3-1.7B thinking with QLoRA adapters using group-relative policy optimization with GSPO-style sequence-level corrections, trained on 14 hand-written prompt templates with a curriculum over constrained-axis subsets and target magnitudes. A hard format/domain gate multiplies the reward: valid structured completions receive R = 3 + 3·R_dist(z(x), z*, S) + 1.5·R_hon(ẑ, z(x), S), rewarding targeting of the realised coordinate and calibration of the self-report; invalid outputs receive zero. For code, the gate requires AST sanitization and execution to a valid 16×16 board, and near-duplicates of retrieved exemplars receive zero reward to prevent copying.

3. Use. Once trained, the controller is a black-box actuator F: (p, z*) → x. For stories, the authors sweep a fixed target grid in Z_text for a fixed prompt, generating conditionally independent branches in parallel. For code, an outer loop (Bayesian optimization with a Gaussian-process surrogate and expected improvement) proposes z*, samples candidates, selects the best valid program, and updates the surrogate at the realised coordinate (z(x_best), f(x_best)) rather than at the request, because reachable support in Z_code is non-uniform. The CA++ benchmark objective f(x) ∈ [0, 1] is never used as an RL reward, so the controller stays objective-agnostic.

Why This Matters

Impact on research. The paper reframes controllable generation as endpoint selection in a fixed external coordinate system rather than as path-level modification (temperature, top-p, logit bias, activation steering, decoding-time guidance). It offers a concrete alternative positioning to trajectory and program search methods such as MCTS, evolutionary program search, FunSearch-style approaches, and test-time RL: the evaluator is confined to the outer loop and can be entirely withheld from the controller. It also contributes a negative-leaning diagnostic — exemplar retrieval, not numeric coordinate parsing, appears to carry most of the conditioning in the code domain — and argues that outer-loop optimizers should model actuator noise and update on realised rather than requested coordinates.

Real-world applications.

• Multi-branch creative drafting: generating genuinely divergent story or script branches from one prompt in parallel, without prompt chaining or accumulating context.
• Inverse design with executable objectives: searching over a low-dimensional design space to produce valid programs, configurations, or seeds that an external simulator or benchmark scores.
• Black-box optimization with withheld objectives: settings where the scoring function is proprietary, expensive, or must not leak into the generator's training signal.
• Diversity during RL: replacing "turn up the temperature" with structured target sweeps that span distinct solution or reasoning modes for group-based RL objectives.

Industry relevance. The interface is deliberately modular and cheap at inference — training absorbs the cost, and generation stays standard sampling, so branches are embarrassingly parallel. That makes it attractive for serving systems that already batch requests, and for pipelines where a frozen, reusable coordinate knob is preferable to per-task prompt engineering. The paper also flags a dual-use concern directly: a low-dimensional reusable coordinate interface plus outer-loop search can lower the friction of systematically searching for undesirable stylistic or rhetorical outputs, so deployments are advised to add content filtering, exemplar-library guardrails, monitoring for adversarial outer-loop objectives, and conservative bounds on the target search region.

Future Directions

• Better spaces and stronger actuators. The authors state that larger backbones and richer exemplar sets may improve prompt-only use of targets, and list better spaces, stronger actuators, and outer loops beyond Bayesian optimization as next steps.
• End-to-end text optimization. The story evaluation currently uses probes, overlap metrics, and an LLM judge rather than a task-level objective; the paper explicitly does not demonstrate an end-to-end outer-loop text task analogous to the code-domain optimization.
• Portability and revalidation of Z. Because the geometry and semantics of the space depend on the frozen encoder, the fitting corpus, and embedding/preprocessing choices, moving to a new domain requires rebuilding and revalidating the space — an open question is how to do this efficiently.
• Reachability and approximate control. Requests are requests: realised coordinates occupy a smaller, non-uniform subset of Z, some targets may be unattainable or attainable only with low probability, and best-of-K improves targeting at higher inference cost. How to model actuator noise and reachable support well remains open.

Target Audience

Researchers and advanced practitioners in LLM controllability, reinforcement learning for generation, and automated search or discovery systems. It is most useful to readers comfortable with encoder embeddings, dimensionality reduction, policy-gradient fine-tuning, and Bayesian optimization — and to engineers designing serving or design-space-exploration pipelines who want a reusable, low-dimensional control interface over model outputs rather than another decoding-time heuristic. Readers primarily interested in human-interpretable semantics of control axes will find the paper's own caution relevant: PCA plus Varimax does not guarantee human-interpretable factors, and the anchored story axis is described as pragmatic rather than canonical.

Authors’ abstract

We introduce Output-Space Search (OS-Search), which turns LLM generation into endpoint search. An outer loop selects a target z* in a frozen encoder-defined 3D output space Z, and a retrieval-grounded policy trained with sequence-level RL generates outputs whose coordinates land near z* under standard autoregressive decoding. This enables parallel sweeps and black-box optimization in Z without path-dependent token/program search. On stories, sweeping Z (text) yields 3.1x…