Technical Summary: Lightweight Latent Reasoning for Narrative Tasks

Problem Statement

Large Language Models (LLMs) have improved performance on complex tasks by generating long chains of thought (reasoning traces) via Reinforcement Learning (RL). However, this approach incurs high computational costs during both training and inference, particularly for narrative tasks involving thousands of tokens. While "latent reasoning" methods—where models generate continuous embeddings instead of discrete tokens to skip steps—have been proposed to improve efficiency, existing approaches face significant limitations in narrative domains:

1. Data Scarcity: Narrative tasks (e.g., story generation, plot hole detection) lack large, high-quality datasets of reasoning traces required to train latent heads effectively.
2. Adaptability: Narrative inputs are highly diverse, making fixed latent reasoning strategies difficult to adapt.
3. Performance Gaps: Existing latent reasoning methods often rely on pre-existing reasoning traces or fail to match the performance of standard non-latent RL training, especially in non-verifiable domains like story planning.

Methodology: LiteReason

The authors propose LiteReason, a latent reasoning framework designed specifically for narrative understanding and generation. It integrates a lightweight Reasoning Projector with standard RL training, allowing the model to interleave discrete token sampling with continuous latent reasoning.

Core Architecture

• Reasoning Projector: A small Multi-Layer Perceptron (MLP) stack added to the base LLM. It takes the model's last hidden state and predicts a continuous token embedding vector.
• Interleaved Modes: The model operates in two modes:
  • Discrete Mode: Standard token sampling via the LM head.
  • Latent Mode: Triggered by a special <implicit_thought> tag (generated by the model). In this mode, the Reasoning Projector predicts continuous embeddings which are fed back into the model, effectively "skipping" multiple reasoning steps without generating explicit text tokens. The tag includes a number indicating the complexity (number of latent steps) to perform.

Training Pipeline

LiteReason training involves three stages:

1. Data Collection: The authors construct an initial dataset via rejection sampling. They sample multiple reasoning traces for each prompt, filtering for those with positive rewards (Accuracy for plot hole detection; Contrastive Improvement for chapter prediction).
2. Supervised Fine-Tuning (SFT): The Reasoning Projector is trained to predict useful embeddings for "skipped" steps. A proportion of reasoning sentences in the traces are replaced with <implicit_thought> tags. The projector learns to bridge the gap between the remaining explicit reasoning steps. The base LLM is frozen during this stage.
3. Reinforcement Learning (RL): The full model is fine-tuned using Group Relative Policy Optimization (GRPO).
   • The policy model learns to decide when to activate the projector (i.e., when to use latent reasoning).
   • Gradients are computed only for discrete token sampling steps; the projector is not directly updated via policy gradients.
   • Periodic SFT: To prevent the projector from becoming outdated as the model's reasoning capabilities evolve, a small SFT step is performed after every RL epoch using the newly generated traces.

Reward Design for Narrative Tasks

• Flawed Fictions (Verifiable): Uses a binary reward (1 for correct plot hole detection, 0 otherwise).
• Next Chapter Prediction (Non-verifiable): Extends the existing VR-CLI (Verifiable Rewards via Completion Likelihood Improvement) objective. The authors introduce a Contrastive Improvement reward, which penalizes generic advice by subtracting the likelihood improvement of a randomly selected chapter from a different book. This encourages diverse, story-specific reasoning.

Key Contributions

1. Novel Framework: Introduction of LiteReason, the first method to interleave latent reasoning with discrete token sampling during RL training for narrative tasks, allowing the model to dynamically choose between latent and discrete reasoning without pre-existing reasoning datasets.
2. Efficiency-Performance Tradeoff: The method achieves performance gains comparable to non-latent RL training while drastically reducing token generation.
3. Task Adaptation: Successful application of latent reasoning to long-context narrative tasks (Flawed Fictions and Next Chapter Prediction), which differ significantly from the math/logic benchmarks where latent reasoning was previously dominant.

Experimental Results

The authors evaluated LiteReason on Flawed Fictions (plot hole detection) and Next Chapter Prediction (NCP) using Qwen 2.5-7B-Instruct.

• Performance vs. Baselines:
  • LiteReason significantly outperforms all training-free (Mixture of Inputs, Soft Thinking) and trained latent reasoning baselines (COCONUT, CoLaR).
  • On Flawed Fictions, LiteReason achieves 87.42% accuracy, reaching 96% of the performance gain of standard non-latent RL (88.71%), while generating 91.5% fewer tokens (34 tokens vs. 114 tokens).
  • On NCP, LiteReason achieves a 613% improvement in Contrastive Improvement over the baseline, reaching 69% of the non-latent RL gain, while generating 77.5% fewer tokens.
• Training Efficiency: LiteReason requires approximately 50% fewer tokens during RL training compared to standard RL, reducing computational costs significantly.
• Inference Speed: LiteReason models generate answers approximately 3x faster than RL-trained baselines in practice.
• Generalization: Unlike other latent methods (COCONUT, CoLaR) which suffered performance collapses on out-of-domain tasks (e.g., GSM-Hard math), LiteReason retained general model capabilities and showed robustness across different model sizes (tested on Qwen 3 4B and Gemma 3 1B).

Significance and Claims

The paper claims that LiteReason successfully guides RL training to a more efficient part of the performance-computation tradeoff curve. By allowing the model to learn when to compress reasoning into latent space, LiteReason avoids the rigid constraints of previous methods that rely on fixed latent stages or high-quality pre-existing reasoning datasets.

The authors position LiteReason as a practical solution for resource-constrained narrative tasks, demonstrating that continuous latent reasoning can be effectively combined with RL to achieve near-optimal performance with significantly reduced computational overhead. They note that this is the first application of continuous latent reasoning to narrative domains and the first to interleave latent and discrete reasoning dynamically during RL.