# OLLM: Options-based Large Language Models

arXiv: https://arxiv.org/abs/2604.19087 · HTML: https://mlsanity.com/paper/2604.19087/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 11 of its 12 references that are in the map.
- First posted on arXiv 21 Apr 2026.

## Builds on (most cited first, 11 of 11)

- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Mitigating the Alignment Tax of RLHF](https://mlsanity.com/paper/2309.06256/) (2023)
- [AIMO-2 Winning Solution: Building State-of-the-Art Mathematical Reasoning Models with OpenMathReasoning dataset](https://mlsanity.com/paper/2504.16891/) (2025)
- [Critic-Guided Decoding for Controlled Text Generation](https://mlsanity.com/paper/2212.10938/) (2022)
- [A Distributional Lens for Multi-Aspect Controllable Text Generation](https://mlsanity.com/paper/2210.02889/) (2022)
- [Composable Text Controls in Latent Space with ODEs](https://mlsanity.com/paper/2208.00638/) (2022)
- [Controllable Text Generation via Probability Density Estimation in the Latent Space](https://mlsanity.com/paper/2212.08307/) (2022)
- [MacLaSa: Multi-Aspect Controllable Text Generation via Efficient Sampling from Compact Latent Space](https://mlsanity.com/paper/2305.12785/) (2023)
- [Embedding-Aligned Language Models](https://mlsanity.com/paper/2406.00024/) (2024)
- [JAM: Controllable and Responsible Text Generation via Causal Reasoning and Latent Vector Manipulation](https://mlsanity.com/paper/2502.20684/) (2025)

## Abstract (arXiv)

We introduce Options LLM (OLLM), a simple, general method that replaces the single next-token prediction of standard LLMs with a textitset of learned options for the next token, indexed by a discrete latent variable. Instead of relying on temperature or sampling heuristics to induce diversity, OLLM models variation explicitly: a small latent space parametrizes multiple plausible next-token options which can be selected or searched by a downstream policy. Architecturally, OLLM is a lightweight "plug-in" that inserts two layers: an encoder and a decoder, before the output head, allowing almost any pretrained LLM to be converted with minimal additional parameters. We apply OLLM to a 1.7B-parameter backbone (only 1.56% of parameters trainable) trained on OpenMathReasoning and evaluated on OmniMath. The SOTA LoRA-adapted baselines peak at 51% final answer correctness, while OLLM's option set allows up to ~ 70% under optimal latent selection. We then train a compact policy in the latent space that emits latents to control generation. Operating in a low-dimensional option space makes reward optimization far more sample-efficient and substantially reduces common misalignments (e.g., language switching or degenerate reasoning), as the policy is constrained to options learned during SFT. Crucially, this alignment arises from model structure rather than additional KL or handcrafted alignment losses. Our results demonstrate that optionized next-token modeling enhances controllability, robustness, and efficiency in math reasoning, and highlight latent-space policy learning as a promising direction for reinforcement learning in LLMs.