Table of Contents
Fetching ...

Think Silently, Think Fast: Dynamic Latent Compression of LLM Reasoning Chains

Wenhui Tan, Jiaze Li, Jianzhong Ju, Zhenbo Luo, Jian Luan, Ruihua Song

TL;DR

This paper tackles the inefficiency of token-level chain-of-thought reasoning in LLMs by introducing Compressed Latent Reasoning (CoLaR), which dynamically compresses reasoning into latent space and learns via a two-stage process: supervised fine-tuning with a next compressed embedding objective and reinforcement learning with GRPO to encourage diverse yet compact latent reasoning paths. CoLaR enables dense latent, silent reasoning and allows inference-time control of reasoning speed through a tunable compression factor, achieving strong accuracy at comparable compression to latent baselines and substantial reductions in reasoning length. Experiments across GSM8k, SVAMP, MultiArith, and the challenging MATH dataset show significant gains over latent baselines and notable scalability with larger base models, while providing interpretability insights into latent chains. While limitations remain in surpassing explicit CoT on all tasks and handling non-integer compression factors, CoLaR demonstrates a promising direction for efficient, flexible reasoning in LLMs with potential practical impact on real-world deployments.

Abstract

Large Language Models (LLMs) achieve superior performance through Chain-of-Thought (CoT) reasoning, but these token-level reasoning chains are computationally expensive and inefficient. In this paper, we introduce Compressed Latent Reasoning (CoLaR), a novel framework that dynamically compresses reasoning processes in latent space through a two-stage training approach. First, during supervised fine-tuning, CoLaR extends beyond next-token prediction by incorporating an auxiliary next compressed embedding prediction objective. This process merges embeddings of consecutive tokens using a compression factor randomly sampled from a predefined range, and trains a specialized latent head to predict distributions of subsequent compressed embeddings. Second, we enhance CoLaR through reinforcement learning (RL) that leverages the latent head's non-deterministic nature to explore diverse reasoning paths and exploit more compact ones. This approach enables CoLaR to: i) perform reasoning at a dense latent level (i.e., silently), substantially reducing reasoning chain length, and ii) dynamically adjust reasoning speed at inference time by simply prompting the desired compression factor. Extensive experiments across four mathematical reasoning datasets demonstrate that CoLaR achieves 14.1% higher accuracy than latent-based baseline methods at comparable compression ratios, and reduces reasoning chain length by 53.3% with only 4.8% performance degradation compared to explicit CoT method. Moreover, when applied to more challenging mathematical reasoning tasks, our RL-enhanced CoLaR demonstrates performance gains of up to 5.4% while dramatically reducing latent reasoning chain length by 82.8%. The code and models will be released upon acceptance.

Think Silently, Think Fast: Dynamic Latent Compression of LLM Reasoning Chains

TL;DR

This paper tackles the inefficiency of token-level chain-of-thought reasoning in LLMs by introducing Compressed Latent Reasoning (CoLaR), which dynamically compresses reasoning into latent space and learns via a two-stage process: supervised fine-tuning with a next compressed embedding objective and reinforcement learning with GRPO to encourage diverse yet compact latent reasoning paths. CoLaR enables dense latent, silent reasoning and allows inference-time control of reasoning speed through a tunable compression factor, achieving strong accuracy at comparable compression to latent baselines and substantial reductions in reasoning length. Experiments across GSM8k, SVAMP, MultiArith, and the challenging MATH dataset show significant gains over latent baselines and notable scalability with larger base models, while providing interpretability insights into latent chains. While limitations remain in surpassing explicit CoT on all tasks and handling non-integer compression factors, CoLaR demonstrates a promising direction for efficient, flexible reasoning in LLMs with potential practical impact on real-world deployments.

Abstract

Large Language Models (LLMs) achieve superior performance through Chain-of-Thought (CoT) reasoning, but these token-level reasoning chains are computationally expensive and inefficient. In this paper, we introduce Compressed Latent Reasoning (CoLaR), a novel framework that dynamically compresses reasoning processes in latent space through a two-stage training approach. First, during supervised fine-tuning, CoLaR extends beyond next-token prediction by incorporating an auxiliary next compressed embedding prediction objective. This process merges embeddings of consecutive tokens using a compression factor randomly sampled from a predefined range, and trains a specialized latent head to predict distributions of subsequent compressed embeddings. Second, we enhance CoLaR through reinforcement learning (RL) that leverages the latent head's non-deterministic nature to explore diverse reasoning paths and exploit more compact ones. This approach enables CoLaR to: i) perform reasoning at a dense latent level (i.e., silently), substantially reducing reasoning chain length, and ii) dynamically adjust reasoning speed at inference time by simply prompting the desired compression factor. Extensive experiments across four mathematical reasoning datasets demonstrate that CoLaR achieves 14.1% higher accuracy than latent-based baseline methods at comparable compression ratios, and reduces reasoning chain length by 53.3% with only 4.8% performance degradation compared to explicit CoT method. Moreover, when applied to more challenging mathematical reasoning tasks, our RL-enhanced CoLaR demonstrates performance gains of up to 5.4% while dramatically reducing latent reasoning chain length by 82.8%. The code and models will be released upon acceptance.

Paper Structure

This paper contains 22 sections, 5 equations, 8 figures, 3 tables.

Figures (8)

  • Figure 1: Our proposed Compressed Latent Reasoning Model (CoLaR) performs dynamic-speed reasoning by auto-regressively predicting latent variables, each compressing information from multiple word tokens. Simply prompting to reason faster enables CoLaR to predict more informative latents.
  • Figure 2: Our proposed method CoLaR consisting an LLM backbone and a Latent Head. During the SFT stage (left), for each training step, CoLaR first compresses embeddings $\mathbf{e}_r$ of the original reasoning chain into compressed embeddings $\mathbf{e}_c$ with a compression factor $c$ randomly selected from the range $[1, c_{max}]$. Then, CoLaR is trained to predict: i) the compressed reasoning embeddings via the Latent Head, and ii) the compressed reasoning tokens and answer tokens through the Language Head. During the RL stage (right), for every question input, CoLaR samples a group of $G$ outputs $o_{1:G}$ consisting of the latent reasoning chain and the predicted answer. We then calculate the relative rewards $a_{1:G}$ for each output, and the rewards are averaged on each token ($\bar{a}_i$), encouraging CoLaR to explore diverse latent reasoning pathways and exploit those more compact ones.
  • Figure 3: A case study on the GSM-8k validation set. We set the compression factor $c$ to 2 and 5, which produce two latent reasoning chains in length 7 and 3, respectively. We then retrieve tokens with the predicted latents by embedding cosine similarity, and underscore those informative tokens.
  • Figure 4: Accuracy and reasoning chain length (# L) of CoLaR on GSM8k dataset when trained with random $c\in[1, 5]$ (the lines) or trained solely on specific $c$ (the bars).
  • Figure 5: Accuracy and reasoning chain length (# L) of CoLaR on GSM8k dataset when trained with $c\in\{1,3,5,7\}$ and tested with extra $c\in\{2,4,6\}$ (under gray bars).
  • ...and 3 more figures