Technical Summary: Compressed Latent Reasoning (CoLaR)

Problem Statement

Large Language Models (LLMs) achieve superior performance in mathematical reasoning through Chain-of-Thought (CoT) prompting, which generates explicit, token-level reasoning chains. However, these lengthy sequences are computationally expensive and inefficient, creating significant server loads and hindering scalability, particularly under high-concurrency conditions. While existing efficiency methods attempt to skip redundant tokens or use compact notations, they remain constrained by sparse, token-based representations. Furthermore, recent latent-based reasoning approaches (e.g., Coconut, CODI) often rely on fixed-length reasoning chains and deterministic latent generation, limiting their adaptability and exploration capabilities.

Methodology: Compressed Latent Reasoning (CoLaR)

CoLaR is a novel framework designed to dynamically compress reasoning processes into a dense latent space while preserving the ability to explore diverse reasoning paths. The method operates through a two-stage training approach:

1. Supervised Fine-Tuning (SFT) with Auxiliary Objectives

CoLaR extends beyond standard next-token prediction by introducing an auxiliary next compressed embedding prediction objective.

• Dynamic Compression: During training, a compression factor $c$ is randomly sampled from a predefined range $[1, c_{max}]$. The embeddings of $c$ consecutive reasoning tokens are merged into a single compressed embedding using an Embedding Compress module. This module scales the sum of embeddings by $1/\sqrt{c}$ to preserve the original distribution properties, avoiding the distortion caused by simple mean pooling.
• Dual-Head Architecture: The model utilizes two heads:
  • Latent Head: A probabilistic head (a two-headed MLP) that predicts the mean ($\mu$) and standard deviation ($\sigma$) of the next compressed embedding distribution. This enables auto-regressive prediction of dense latent variables rather than deterministic values.
  • Language Head: Predicts the final answer and, crucially, provides dense supervision by predicting the original reasoning tokens corresponding to the compressed embeddings. This approximates a multi-label classification task to ensure the model learns the semantics of the compressed groups.
• Loss Function: The training loss combines the language modeling loss (for token prediction) and a latent loss. The latent loss utilizes a soft-MSE formulation that combines Mean Squared Error with an entropy regularization term. This encourages the model to maintain exploration capability (diverse latents) while fitting the data, addressing the under-fitting issues observed with standard Negative Log-Likelihood (NLL) loss on simpler datasets.

2. Reinforcement Learning (RL) for Exploration and Exploitation

To move beyond mimicking the CoT teacher and unlock the potential of latent reasoning, CoLaR employs Group Relative Policy Optimization (GRPO).

• Probabilistic Sampling: Because the Latent Head outputs a probability distribution, the model can sample diverse latent reasoning chains for the same input.
• Reward Mechanism: For each question, a group of outputs (latent chains + answers) is generated. Rewards are assigned based on answer correctness (1 for correct, 0 for incorrect).
• Balancing Strategy: The GRPO objective calculates group-normalized rewards. Crucially, the reward is averaged over the number of tokens/latents in the chain. This design inherently encourages the model to exploit shorter, correct reasoning paths (as the reward per token is higher) while still exploring diverse paths to find correct solutions.

Key Contributions

1. Dynamic-Speed Latent Reasoning: CoLaR introduces a framework that auto-regressively predicts latent variables encapsulating the semantics of multiple word tokens. This allows for dynamic adjustment of reasoning speed at inference time simply by prompting a desired compression factor.
2. Probabilistic Latent Head & RL Integration: Unlike prior deterministic latent methods, CoLaR employs a probabilistic Latent Head. Combined with RL, this enables the model to explore diverse reasoning pathways and exploit more compact solutions, significantly improving efficiency without sacrificing accuracy.
3. Novel Training Objectives: The introduction of the next compressed embedding prediction task and the soft-MSE latent loss allows the model to effectively learn dense representations and balance exploration-exploitation trade-offs.

Experimental Results

Extensive evaluations were conducted on four grade-school math datasets (GSM8k, GSM8k-hard, SVAMP, MultiArith) and the more challenging MATH dataset.

• Performance vs. Latent Baselines: On grade-school datasets, CoLaR achieved a 14.1% higher accuracy than state-of-the-art latent-based baselines (e.g., Coconut) at comparable compression ratios.
• Efficiency vs. Explicit CoT: Compared to explicit CoT, CoLaR reduced reasoning chain length by 53.3% with only a 4.8% performance degradation.
• RL Enhancements on Challenging Tasks: On the MATH dataset, applying RL to CoLaR resulted in a 5.36% accuracy gain over the non-RL version while reducing the reasoning chain length by 82.8%.
• Scaling: Experiments scaling from 1B to 8B parameters confirmed that CoLaR's efficiency and accuracy gains persist with larger model capacities. CoLaR-8B-RL surpassed its CoT teacher on the GPQA benchmark (37.5% vs 35.7%) while shortening the reasoning chain by 69%.
• Generalization: The model demonstrated robust generalization to unseen compression factors and out-of-domain datasets (e.g., MultiArith, GPQA).

Significance

The paper posits that CoLaR addresses the critical bottleneck of computational cost in LLM reasoning by shifting from sparse token-level processing to dense latent-space reasoning. Its significance lies in:

• Dynamic Adaptability: The ability to adjust reasoning speed (compression factor) at inference time offers a flexible trade-off between latency and accuracy, suitable for varying real-world application constraints.
• Efficiency without Sacrifice: By leveraging probabilistic latent reasoning and RL, CoLaR achieves substantial reductions in reasoning chain length (up to 82.8%) while maintaining or even improving accuracy on complex tasks, challenging the notion that efficiency must come at the cost of performance.
• Exploration-Exploitation in Latent Space: The work demonstrates that introducing stochasticity and reinforcement learning into latent reasoning processes allows models to discover more efficient reasoning paths that deterministic, fixed-length approaches miss.

The authors acknowledge limitations, including current performance approximating (rather than consistently surpassing) explicit CoT on some benchmarks and challenges in generalizing to non-integer compression factors due to discrete tokenization constraints. However, the framework establishes a promising direction for efficient, scalable, and adaptive LLM reasoning.