BOW: Bottlenecked Next Word Exploration

2025-06-16

AI Summary

Traditional large language models (LLMs) rely on next-word prediction (NWP) training, which ensures surface-level fluency but hinders robust reasoning. To address this, we propose a “reasoning-path–word-prediction” decoupled architecture: a policy model explicitly generates interpretable reasoning paths, while a frozen discriminative model predicts token distributions solely conditioned on those paths—thereby externalizing and modularizing the reasoning process. We introduce a novel reward function based on next-word recovery capability and jointly optimize both modules via GRPO-based reinforcement learning. Experiments demonstrate substantial improvements across diverse reasoning benchmarks—including general-purpose reasoning (GSM8K, MMLU) and next-word prediction tasks (LAMBADA)—establishing our approach as a scalable, high-performance alternative to conventional NWP paradigms.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Large Multimodal Models (LMMs)
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Large language models (LLMs) are typically trained via next-word prediction (NWP), which provides strong surface-level fluency but often lacks support for robust reasoning. We propose BOttlenecked next Word exploration (BOW), a novel RL framework that rethinks NWP by introducing a reasoning bottleneck where a policy model first generates a reasoning path rather than predicting the next token directly, after which a frozen judge model predicts the next token distribution based solely on this reasoning path. We train the policy model using GRPO with rewards that quantify how effectively the reasoning path facilitates next-word recovery. Compared with other continual pretraining baselines, we show that BOW improves both the general and next-word reasoning capabilities of the base model, evaluated on various benchmarks. Our findings show that BOW can serve as an effective and scalable alternative to vanilla NWP.

Problem

• Improves reasoning in next-word prediction models
• Introduces bottlenecked reasoning path generation
• Enhances general and next-word reasoning capabilities

Innovation

• Introduces reasoning bottleneck before token prediction
• Uses GRPO for policy model training
• Combines policy and frozen judge models