OLLM: Options-based Large Language Models

📄 arXiv: 2604.19087v1 📥 PDF

作者: Shashank Sharma, Janina Hoffmann, Vinay Namboodiri

分类: cs.AI

发布日期: 2026-04-21


💡 一句话要点

OLLM:基于选项的大语言模型,提升数学推理任务的可控性和效率。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 选项模型 数学推理 潜在空间 强化学习

📋 核心要点

  1. 现有LLM依赖温度或采样启发式方法来生成多样性,缺乏对生成过程的显式控制。
  2. OLLM通过引入离散潜在变量索引的选项集,显式地建模下一个token的多种可能性,增强了可控性。
  3. 实验表明,OLLM在数学推理任务上显著优于LoRA微调的基线模型,且对齐效果更好。

📝 摘要(中文)

本文提出了一种名为Options LLM (OLLM) 的通用方法,它使用一组学习到的选项来替代标准LLM的单个下一个token预测,这些选项由一个离散潜在变量索引。OLLM通过一个小的潜在空间参数化多个合理的下一个token选项,下游策略可以选择或搜索这些选项,从而显式地模拟了变异,而无需依赖温度或采样启发式方法来诱导多样性。在架构上,OLLM是一个轻量级的“插件”,在输出头之前插入一个编码器和一个解码器,允许几乎任何预训练的LLM以最小的额外参数进行转换。作者将OLLM应用于一个在OpenMathReasoning上训练的17亿参数backbone(只有1.56%的参数可训练),并在OmniMath上进行评估。SOTA LoRA调整的基线模型的最终答案正确率峰值为51%,而OLLM的选项集在最佳潜在选择下允许达到约70%。然后,作者在潜在空间中训练一个紧凑的策略,该策略发出潜在变量来控制生成。在低维选项空间中操作使得奖励优化更有效,并显著减少了常见的未对齐问题(例如,语言切换或退化的推理),因为该策略被限制在SFT期间学习的选项中。重要的是,这种对齐来自模型结构,而不是额外的KL或手工制作的对齐损失。结果表明,optionized的下一个token建模增强了数学推理中的可控性、鲁棒性和效率,并强调了潜在空间策略学习是LLM中强化学习的一个有希望的方向。

🔬 方法详解

问题定义:现有的大语言模型在生成文本时,通常采用单一的token预测方式,并通过调整温度系数或采样策略来引入多样性。这种方式缺乏对生成过程的直接控制,难以保证生成结果的质量和一致性。尤其是在需要精确推理的任务中,例如数学问题求解,这种不确定性会导致错误累积,影响最终结果的准确性。

核心思路:OLLM的核心思路是将传统的单一token预测替换为基于选项集的预测。具体来说,模型不再直接预测下一个token,而是预测一个离散的潜在变量,该变量对应于一个预先学习好的token选项集合。通过选择不同的潜在变量,模型可以生成不同的token,从而实现对生成过程的更精细控制。这种设计使得模型能够显式地建模多种可能的输出,并根据下游任务的需求选择最合适的选项。

技术框架:OLLM的整体架构是在预训练的LLM基础上添加一个轻量级的“插件”,该插件包含一个编码器和一个解码器。编码器将LLM的隐藏状态映射到潜在空间,解码器则将潜在变量映射回token空间,生成对应的token选项。在训练阶段,模型学习潜在变量与token选项之间的对应关系。在推理阶段,下游策略可以选择或搜索潜在空间,从而控制生成过程。整个框架可以分为以下几个阶段:1) LLM backbone生成隐藏状态;2) 编码器将隐藏状态映射到潜在空间;3) 从潜在空间中选择一个潜在变量;4) 解码器将潜在变量映射回token选项;5) 根据下游策略选择最终的token。

关键创新:OLLM最重要的技术创新在于引入了基于选项集的token预测方式。与传统的单一token预测相比,OLLM能够显式地建模多种可能的输出,并允许下游策略对生成过程进行更精细的控制。此外,OLLM的轻量级插件设计使得它可以很容易地应用于各种预训练的LLM,而无需进行大量的参数调整。这种设计降低了模型的训练成本,并提高了模型的泛化能力。

关键设计:OLLM的关键设计包括:1) 离散潜在空间的维度:作者通过实验确定了合适的潜在空间维度,以平衡模型的可控性和表达能力。2) 编码器和解码器的网络结构:作者采用了简单的线性层作为编码器和解码器,以减少额外的参数量。3) 损失函数:作者采用了交叉熵损失函数来训练模型,目标是最大化正确token选项的概率。4) 潜在空间策略学习:作者训练了一个紧凑的策略网络,用于在潜在空间中选择合适的潜在变量,以优化下游任务的奖励。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在OmniMath数据集上的实验结果表明,OLLM在数学推理任务上显著优于LoRA微调的基线模型。SOTA LoRA调整的基线模型的最终答案正确率峰值为51%,而OLLM的选项集在最佳潜在选择下允许达到约70%。此外,OLLM在训练过程中不需要额外的KL或手工制作的对齐损失,即可实现更好的对齐效果。

🎯 应用场景

OLLM具有广泛的应用前景,尤其是在需要高精度和可控性的任务中,例如数学推理、代码生成、对话系统等。通过显式地建模多种可能的输出,OLLM可以提高生成结果的质量和一致性,并减少错误累积。此外,OLLM的潜在空间策略学习方法为LLM的强化学习提供了一个新的方向,有望在未来实现更智能和可控的文本生成。

📄 摘要(原文)

We introduce Options LLM (OLLM), a simple, general method that replaces the single next-token prediction of standard LLMs with a \textit{set 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 $\sim 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.