Selection as a Model Shape

Abstract

We impart intelligence to models by training them with next token prediction. Not all problems are 'next token shape', which makes NTP a costly way to deploy this intelligence. We propose a new shape selection that selects output directly from context in a single pass. E.g. given a maze, it selects the path that solves it, in a single pass. Or given a large document, it selects paragraphs relevant for a search result. Using a custom head, we were able to convert Qwen 3.8 27B into this shape, and solve problems in single pass that otherwise would have taken thousands of tokens. A single forward pass has limited intelligence, but we were surprised by the number of tasks that can be done even with that.

Introduction

Next-token prediction is the standard training objective for language models, and autoregressive generation is their usual inference interface. Many applications, however, need an output that refers to the input itself. For example, a document system may need the exact clause that governs a contract, while a game controller may need the cells that form a shortest path through a maze. In each case, the input defines the available positions. Identifying the right subset can still require substantial reasoning.

We study selection as an output interface for these problems. Given text or an image and one or more questions, SEL predicts scores or spans over positions in the source. The application can map the result back to character offsets or image coordinates. We adapted Qwen 27B to perform this computation in one backbone pass and trained it on 26 text and image tasks. The release includes the training and inference code, LoRA adapter, selection head, and training dataset. GitHub · Model weights · Training dataset

Model depth, inference budget, and training data

Three aspects of a language model affect the computation available to solve a problem.

1. Model depth. A deeper model can apply more transformations to each position during a forward pass. Once the architecture is fixed, the number of layers traversed per pass is fixed. The total compute still varies with input length.
2. Tokens generated at inference. Each additional decoding step can use the tokens produced so far as an external scratchpad. This supports intermediate calculations, revisiting earlier choices, and decomposition into subproblems. The resulting sequential computation scales with the generation budget and adds latency. The context window and serving budget constrain how far it can scale.
3. Training data. Pretraining and adaptation determine which computations the model can carry out within a fixed inference budget. A model that has learned to recognize a useful pattern or apply a procedure can often do so in fewer steps. For SEL, supervision also specifies the output granularity, such as a supporting sentence, a single letter, or a complete path.

Systems such as Jev use a single backbone pass to score a set of options provided with the request. Selection uses a related inference budget with a different output space. An application can specify a board and ask for a route without enumerating candidate routes. The model instead predicts which cells belong to the path. This is analogous to segmentation, where the output refers to positions in the input.

Generating a sequence of moves can serve two purposes. Intermediate tokens can help the model determine the route, and output tokens can serialize a route that it has already determined. A source-aligned readout removes the need for that serialization. Our experiments ask which selections a pretrained backbone can learn to infer from a single set of hidden states.

Architecture

One backbone pass

The Qwen 27B backbone interleaves 16 full-attention layers with 48 Gated DeltaNet layers. We make the full-attention layers bidirectional and retain the left-to-right recurrence of the DeltaNet layers. With a source-first prompt, the modified layers allow source positions to incorporate information from a question that appears later in the sequence. We train rank-16 LoRA adapters on the language backbone along with the selection head. The base weights, vision encoder, and token embeddings remain frozen.

Encoding the source and questions

We encode the source, any environment rules, and the numbered questions in a single sequence using Qwen’s multimodal processor. The readout uses the final hidden states at selectable source positions and at the last token of each question. We retain the mapping from source tokens to character offsets or visual patches. The output head reads these states directly, bypassing the language model’s vocabulary projection.

Selection heads

Dense selection

The dense readout assigns a logit to each source position for each question. It projects source and question vectors separately from 5,120 to 256 dimensions. For projected vectors s and q, the scorer applies a learned linear readout to GELU(s + q + (s ⊙ q) / √256), where ⊙ denotes elementwise multiplication. The scores are computed in parallel and thresholded after a sigmoid. This readout is appropriate when the target is a mask, for example the union of supporting sentences or the cells of a specified maze path.

Pointer selection

A dense mask can combine positions from incompatible answers. In Tetris, two equally good placements may occupy different sets of cells. A mask that combines those cells need not correspond to a legal placement. To represent a single coherent answer, we use a recurrent pointer over source positions. It selects a span start or STOP and, after selecting a start, chooses an inclusive end. Additional spans are selected in source order and converted back to a mask.

Linked questions

Some requests contain dependent selections. For example, a request might ask which confirmed estimate to move and then which folder belongs to the corresponding client. The correct folder depends on the estimate selected in the first step. Both questions are encoded with the source in the same backbone pass.

At each step, the 256-dimensional GRU receives the projected question vector, the vector at the preceding pointer position, and an embedding for the action type. When the pointer proceeds to the next question, it resets the preceding-position input to the learned STOP vector and retains the recurrent state. The pointer’s sequential computation therefore operates on the states from a single backbone pass. A linear router on the pooled question states chooses the dense readout, a pointer for one answer, or a linked pointer for dependent questions.

Training the selector

We trained a 4B and a 27B selector for one epoch on the same 26-task mixture. The training split contains 31,244 examples and 11.81 million input tokens, with an input limit of 2,048 tokens per example. The tasks include evidence selection, numerical lookup, screen interaction, and games. Annotations specify complete acceptable selections in source coordinates. When a task has multiple valid answers, we retain each answer as a separate target, including all questions in a dependent request.

For dense examples, we minimize binary cross-entropy over selectable positions against the lowest-loss complete target. For pointer examples, we use the negative log marginal likelihood of the valid span trajectories, normalized by their mean length, and add the dense loss with weight 0.25. We add a routing cross-entropy term with weight 0.10 in both cases. We use AdamW with four examples per optimizer step and a peak learning rate of 3 × 10⁻⁵. The schedule warms up over the first 5% of training source positions and decays by cosine to one-quarter of the peak.

The 27B run trains 116.7 million LoRA parameters and 3.17 million head parameters. Its frozen language weights are loaded in NF4 with BF16 compute. The 4B run uses a BF16 backbone.

Held-out results

Both selectors are evaluated on the same 2,605 held-out examples. For dense predictions, we use fixed thresholds of 0.40 for text and 0.30 for images. An example counts as exact if the predictions for all its questions match one complete reference alternative. For macro F1, we choose the complete alternative with the highest mean question F1 for each example, then average the resulting scores over questions. The comparison holds the data mixture and training schedule fixed. It reflects the two trained systems, whose backbones and weight precision differ.

The largest improvements among these examples occur on full-path selection and numerical lookup. On full image paths, patch F1 increases from 63.1% to 84.2%, alongside the increase from zero to 13 exact paths. Exact match penalizes any additional or missing square. Next-move accuracy is measured on supplied states. A rollout would additionally test whether errors accumulate over states reached by the model’s own predictions.

The Tetris example in the figure is the only exact held-out placement from SEL 27B. Its patch F1 of 59.0% indicates partial overlap with valid footprints, while legality and rows cleared would require evaluation at the level of complete moves. The geolocation figure shows the highest-scoring patch. Its score is below the standard threshold, and the estimate is approximately 1,100 km from the labeled location. A distance-based metric would provide more resolution than patch exact match for this task. Panorama by Yuka_Washio. Mapillary · CC BY-SA 4.0

Selection and generation on the same examples

We also measured latency on individual examples using SEL 27B and a causal Qwen 27B INT4 baseline with thinking enabled. The baseline received a serialization format appropriate to each task, including move letters for mazes, verbatim sentences for evidence selection, and row–column pairs for Tetris. Both systems ran warm on one RTX A6000. SEL timings are medians of five requests. They include loading prepared inputs for the held-out examples and returning scores to the CPU. The long-text request starts from raw text and produces an input of 1,888 tokens. Qwen timings are single vLLM requests with prefix caching disabled. Generated token counts include reasoning, final answers, and control tokens.

The arithmetic example is a matched-output comparison. Both systems returned M, the letter paired with 73, with a measured latency of 274 ms for SEL and 6.83 s for Qwen. For the other examples, timing should be read alongside the outcome. In an earlier maze run, Qwen generated 2,781 tokens and returned an incorrect route. In the timed run it exhausted the generation budget before returning a route. These measurements compare a selection-trained model and an unadapted causal baseline using different inference runtimes. They do not isolate the contribution of the readout or estimate a speedup at matched accuracy across the dataset.

Evidence selection after retrieval

To examine the downstream effect of selection, we indexed 233 essays by Paul Graham, comprising about 697,000 Qwen tokens. We asked why founders should do things that do not scale and how they should find their first users. Both conditions use the same eight BM25 results and the same answer model. The baseline passes all retrieved chunks to the answer model. The selection condition scores the sentences in those chunks and passes only the highest-ranked evidence. Answer-stage latencies are medians of three warm requests. The totals sum the measured stages.

Selection reduced the answer model’s input by 95.7%, and answer generation took 3.19 s instead of 7.72 s. The index lookup took 0.16 ms in either condition. Serial selection over the eight chunks added 7.35 s, yielding a total of 10.53 s. Both answers included manual user recruitment. The answer produced from selected context also recommended an untargeted launch, illustrating that reducing the input can change the substance of the answer. Do Things that Don’t Scale

The usefulness of source-aligned predictions depends on how the application evaluates them. A game can check the legality of a proposed move and measure performance over rollouts. A search system can evaluate whether the selected evidence supports an accurate answer, alongside the latency and context costs of producing it. These application-level evaluations would help establish when the savings from selection justify the cost of running the selector.

Citation

@misc{spock2026selection,
  author = {Gupta, Apurva and {Codex}},
  title = {Selection as a Model Shape},
  year = {2026},
  url = {https://spock.one/blog/selection-is-a-model-shape}
}