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

TL;DR

OS-Search introduces an encoder-defined, fixed output space $Z$ and an outer-loop strategy to search targets $z^\u00a0$ for LLM generations, enabling parallel, endpoint-focused control without altering decoding. A retrieval-grounded $z^\u00a0$-conditioned controller is trained with sequence-level RL to produce outputs $x$ whose coordinates $z(x)\in Z$ align with $z^\u00a0$, while preserving autoregressive decoding. In storytelling, anchoring $Z_{\text{text}}$ yields a stable axis $z_1$ and, when sweeping $Z$, achieves substantially higher diversity than path-based baselines; in coding, Bayesian optimization over $z^\u00a0$ improves a withheld CA++ objective under matched budgets while preserving validity. The results demonstrate target-tracking accuracy in $Z$, calibration of $z(x)$ to targets, and notable performance gains, offering a reusable outer-loop interface for controllable generation across domains with a fixed, interpretable coordinate space.

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 higher LLM-scored diversity than prompt-chaining. On code, Bayesian optimization over Z (code) improves an objective withheld from the controller under matched inference budgets while preserving validity.

Figures (13)

• Figure 1: Path- vs state-like control (schematic 2D slice of $\mathbf{Z}$). (a) Standard decoding: a fixed prompt yields token-space paths whose endpoints concentrate in one region of the output space $Z$. (b) $z^\ast$-conditioned decoding: the same base prompt plus a requested target $z^\ast$ (grounded by retrieved exemplars) shifts endpoints toward $z^\ast$ in $Z$ (up to targeting error), while the decoder remains autoregressive. (c) External search: an optimizer moves $z^\ast$ across $Z$ to explore regions and maximize a score $f(x)$. Because branches are conditionally independent given the instantiated prompt, multi-branch generation is parallel.
• Figure 2: Prompt schema for OS-Search. The prompt includes retrieved exemplars with stored realised coordinates $z(x_j)$ and the numeric request $z^\ast_S$. The model outputs a structured completion $y$ containing task content $x$ and a self-report $\hat{z}_S$. We compute realised coordinates by embedding only '<text'> to obtain $z(x)$.
• Figure 3: Success--tolerance curves for 3D control in $Z$. Dashed curves show one-shot sampling, and solid curves show best-of-5 at the same request. Malformed/invalid outputs count as failures. Best checkpoint (blue) against first 5% of the RL run (gray).
• Figure 4: Relationship between requested $z_1^\ast$ and Slop-Score. Slop-Score as a function of the requested target on the first axis, $z_1^\ast$, for completions from late-training rollouts (final 10% of GRPO updates) generated under 1D control ($S=\{1\}$) ($N \approx 7{,}000$). Points are coloured by local density. The solid line shows a LOWESS smooth, the dotted line the global linear regression fit ($r \approx 0.57$), and the shaded band an approximate 80% prediction interval around the LOWESS trend.
• Figure 5: Programs in CA++ behavioral space (code domain). Each point is a valid make_seed() program scored by CA++. Axes are PCs of the CA++ trajectory-statistic feature vectors of the plotted programs (so proximity indicates similar CA dynamics), not $Z_{\text{code}}$. Hexbins show random-$z^\ast$ controller samples coloured by mean score per bin. Red outlines mark bins reached by the $N{=}188$ program-library baseline, and stars mark bins whose mean exceeds the best baseline score (0.371). Insets show representative seeds generated by corresponding make_seed() implementations. The printed $(z_1,z_2,z_3)$ values are the requested targets $z^\ast$ for those programs.