OLLM: Options-based Large Language Models

1. Problem Statement

Traditional Large Language Models (LLMs) operate on a single next-token prediction paradigm. To generate diverse outputs, they rely on stochastic sampling heuristics (e.g., temperature, top-k, top-p). This approach has several limitations:

• Implicit Diversity: Variation is induced via noise rather than being explicitly modeled, making it difficult to control or optimize.
• Competition: Plausible alternative tokens compete within a single probability simplex (softmax), often causing the model to overfit to a dominant mode or fail to capture genuine multimodality in reasoning.
• Inefficient RL: Reinforcement Learning (RL) fine-tuning over the full vocabulary is sample-inefficient and prone to "reward hacking" or degenerate behaviors (e.g., language switching, incoherent reasoning) because the action space is too large and unstructured.
• Ambiguity vs. Determinism: Empirical analysis shows that while many token positions are deterministic (e.g., completing a math formula), a significant portion (~58%) are highly ambiguous, admitting multiple plausible continuations. Standard LLMs struggle to explicitly represent these alternatives without collapsing into a single mode.

2. Methodology: Options LLM (OLLM)

The authors propose OLLM, a lightweight architectural modification that replaces the single next-token distribution with a set of learned options indexed by a discrete latent variable.

Architecture

OLLM acts as a "plug-in" module inserted before the output head (lm_head) of any pretrained LLM. It consists of two new layers:

1. Encoder ($E_\psi$):
   • Input: The contextual hidden state ($h_t$) and the embedding of the ground-truth next token ($e(y_{t+1})$).
   • Function: Maps these inputs to a probability distribution over a small, discrete latent space $Z = \{1, \dots, K\}$.
   • Goal: During training, it learns to assign the correct "option" (latent) to the target token based on context.
2. Decoder ($D_\omega$):
   • Input: The selected latent variable $z$ and the hidden state $h_t$.
   • Function: Generates a bias vector $\Delta h_t$ that is added to the hidden state before the vocabulary projection.
   • Effect: This biases the logits such that different latent options correspond to different logit spaces, effectively disentangling competing continuations.

Training Objectives

The model is trained using a combination of three losses:

1. Cross-Entropy Loss ($L_{CE}$): Standard token prediction loss, ensuring the correct token is predicted given the latent $z$.
2. KL Divergence Regularization ($L_{KL}$): Encourages the encoder's distribution to remain close to a uniform prior to prevent "latent collapse" (where all tokens map to a single option). An adaptive scaling factor is used to maintain a target KL value (0.2).
3. Policy Loss ($L_{policy}$): A lightweight policy network $\pi_\phi$ is trained via Supervised Fine-Tuning (SFT) to mimic the encoder. It predicts the latent $z$ directly from the hidden state $h_t$ (without access to the ground-truth token).

Inference

During inference, the Encoder is replaced by the Policy. The policy selects a latent $z$ from the hidden state, and the Decoder biases the generation accordingly. This reduces the downstream optimization action space from the full vocabulary ($V$) to the small latent set ($K$).

3. Key Contributions

• Explicit Option Modeling: OLLM explicitly models variation by learning a compact latent space that indexes multiple plausible next-token alternatives, rather than relying on sampling noise.
• Plug-and-Play Efficiency: The architecture requires only two additional layers (Encoder/Decoder). In experiments, only 1.56% of the total parameters (on a 1.7B backbone) were trainable, making it highly parameter-efficient.
• Disentangled Reasoning: By separating competing continuations into different latent options, the model avoids the "softmax competition" problem, allowing it to preserve multiple valid reasoning paths simultaneously.
• Sample-Efficient RL Foundation: The low-dimensional latent space serves as a constrained action space for future RL, preventing the policy from exploring invalid or degenerate behaviors outside the distribution learned during SFT.
• No External Labels: Unlike prior controllable generation methods, OLLM does not require external labels (e.g., sentiment, topic); the options are learned purely from the data structure.

4. Experimental Results

• Dataset: Trained on OpenMathReasoning and evaluated on OmniMath (a universal Olympiad-level benchmark).
• Baseline: State-of-the-art LoRA-adapted baselines achieved a peak 51% final answer correctness.
• OLLM Performance:
  • Optimal Selection: Under optimal latent selection (oracle), OLLM achieved ~70% final answer correctness and ~83% per-token accuracy.
  • Policy Performance: Even with a simple SFT-trained policy operating in the latent space, OLLM significantly outperformed baselines, demonstrating that the latent space successfully captured the necessary variations.
• Qualitative Analysis:
  • In deterministic cases (e.g., math symbols), the model collapsed to a single consistent option.
  • In ambiguous cases (e.g., starting a reasoning step), the model successfully learned distinct options corresponding to different logical branches.

5. Significance and Future Work

• Improved Controllability: OLLM demonstrates that constraining generation to a learned option space enhances robustness and reduces common misalignments (like language switching) without needing complex KL penalties or handcrafted alignment losses.
• Interpretability: Different latent options can correspond to distinct reasoning trajectories, offering a new level of interpretability for LLM decision-making.
• Future Directions: The authors propose using this latent space for Reinforcement Learning (RL). Because the action space is small and grounded in SFT data, RL optimization is expected to be far more sample-efficient and stable than traditional RLHF over the full vocabulary. Future work will extend this to code generation, scientific QA, and multi-step planning.

Conclusion:
OLLM represents a paradigm shift from "sampling diversity" to "modeling diversity." By explicitly learning a discrete latent space of options, it enables more efficient, controllable, and robust reasoning in LLMs, particularly in complex domains like mathematics.