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

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

Source-linked AI summary

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

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

arXiv:2505.16552v6May 22, 2025cs.CL

TL;DR

Lengthy token-level Chain-of-Thought reasoning is computationally expensive, motivating more efficient reasoning. CoLaR compresses reasoning into latent space with variable compression and reinforcement learning, achieving higher accuracy with shorter reasoning chains across mathematical tasks.

Problem
Lengthy token-level Chain-of-Thought reasoning is computationally expensive and inefficient, creating a bottleneck for scalable LLM reasoning.
Method
CoLaR trains an auxiliary compressed-embedding objective and a probabilistic latent head, then uses reinforcement learning to explore diverse paths and exploit compact ones.
Results
Across four mathematical reasoning datasets, CoLaR improves accuracy by 14.1% over latent-based baselines at comparable compression ratios, while reinforcement learning yields up to 5.36% higher accuracy and 82.8% shorter reasoning on MATH.
Takeaways & Limitations
CoLaR enables dense latent reasoning with dynamically adjustable reasoning speed and shorter reasoning chains while preserving exploration and exploitation capabilities.
Takeaways & Limitations
CoLaR struggles to generalize to non-integer compression factors and values greater than the maximum training compression factor rmax.

Abstract

from arXiv · show

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%.

1 Introduction

CoLaR addresses the computational cost of lengthy token-level reasoning by dynamically compressing reasoning chains into latent space. Its probabilistic latent reasoning and reinforcement learning improve accuracy while reducing reasoning length.

CoLaR targets the computational burden and server load created by lengthy reasoning chains, particularly under high-concurrency conditions.
CoLaR dynamically compresses multiple reasoning tokens into latent variables, enabling reasoning at a denser representation.The framework uses an auxiliary next compressed embedding prediction task and variable compression factors during training.
A probabilistic Latent Head supports diverse reasoning pathways, while reinforcement learning encourages exploration and exploitation of shorter chains.
53.3% reduction in reasoning chain length comes with only 4.8% performance degradation relative to explicit CoT.
14.1% accuracy improvement is reported over existing latent-based methods at comparable compression ratios.

2 Related Work

Related work improves reasoning efficiency through token selection, architectural computation changes, and latent-space reasoning. CoLaR extends autoregressive latent reasoning with compressed embeddings and dynamic chain lengths.

Explicit reasoning methods shorten chains by skipping redundant tokens or encouraging compact patterns, but remain limited by sequential token prediction.
Latent reasoning methods operate in a denser continuous space and include knowledge internalization, architectural modifications, and autoregressive latent reasoning.
Coconut and CODI perform autoregressive latent reasoning, but their implicit latent variables constrain them to fixed-length reasoning chains.
CoLaR captures multiple word-token semantics in one latent variable through a next compressed embedding objective, enabling dynamic chain lengths.
Reinforcement learning with a probabilistic latent prediction head improves performance while substantially reducing latent reasoning length.

3 Method

CoLaR compresses reasoning-token embeddings into latent representations and trains an LLM with language and latent heads, then applies reinforcement learning to explore diverse paths and favor compact reasoning. Its pipeline combines random-factor compression, probabilistic next-embedding prediction, and group-relative rewards averaged across latent and token outputs.

3.1 Reasoning token compression and understanding: CoLaR represents mathematical problems with question, reasoning-chain, and answer tokens, whose embeddings are processed by an LLM backbone and Language Head.The backbone produces final-layer hidden states and next-token logits for these sequences.
3.1 Reasoning token compression and understanding: During SFT, CoLaR predicts compressed reasoning embeddings with the Latent Head and compressed reasoning plus answer tokens with the Language Head.The auxiliary token objective samples one token from each compressed group to approximate a multimodal distribution over the group’s potential tokens.
3.1 Reasoning token compression and understanding: Randomly sampled compression factors group consecutive reasoning embeddings into compressed embeddings, reducing the reasoning representation to length Lc = ⌈Lr/c⌉.The Embedding Compress module scales the sum of grouped embeddings by 1/√c to reduce distributional distortion.
3.2 Next compressed embedding prediction: The two-headed Latent Head predicts a distribution for the next compressed embedding, enabling stochastic latent generation rather than deterministic reasoning paths.Inference samples the next embedding with a re-parameterization trick; the probabilistic formulation supports uncertainty and diverse pathways.
3.2 Next compressed embedding prediction: A soft-MSE alternative adds entropy regularization because NLL empirically under-fits simpler mathematical reasoning datasets.The entropy term encourages larger predicted standard deviations and more diverse latents; the compressed-embedding and latent losses are summed during SFT.
3.3 Exploration with reinforcement learning: During RL, GRPO samples groups of latent reasoning chains and answers, uses group-relative rewards, and averages output-level rewards across latent or token positions.Correct answers receive reward 1 and incorrect answers 0; position-wise averaging encourages both exploration and exploitation of compact solutions.

4 Experiments

Experiments evaluate CoLaR against latent and token-level baselines across mathematical reasoning datasets, then analyze ablations, reinforcement learning, scaling, and dynamic compression. CoLaR improves accuracy while shortening reasoning chains, with performance depending on stochastic exploration, dense supervision, latent alignment, reward averaging, and base-model quality.

Comparison to baseline methods on GSM datasets: CoLaR-5 improves average accuracy over Coconut by 14.1% while using fewer reasoning steps, 4.57 versus 6.00.This comparison covers four grade-school mathematical reasoning datasets.
Comparison to baseline methods on GSM datasets: CoLaR-2 achieves 48.8% accuracy, only 4.8% below explicit CoT, while reducing reasoning chain length by 53%.The same trained model can be evaluated with a different test-time compression factor.
Comparison to baseline methods on GSM datasets: CoLaR shows minimal performance degradation relative to CoT on MultiArith, whereas other latent-based methods suffer significant drops.This result indicates robust out-of-domain behavior on that dataset.
Ablation studies: A deterministic latent head fits simple datasets but lacks test-time exploration, while NLL training adds excessive randomness and worsens overall performance.The ablation supports balancing exploration and exploitation in latent reasoning.
Ablation studies: Removing compressed reasoning-chain tokens from cross-entropy labels degrades performance by 1.6% at c = 5 and 0.6% at c = 2.The comparison identifies dense reasoning supervision as important for latent-based reasoning.
Ablation studies: Mean pooling causes 3.4% and 2.2% performance degradation, attributed primarily to compression-induced distribution shifts.The result supports sharing a common latent space across compression factors.
Reinforcement learning results: On MATH, RL-enhanced CoLaR gains up to 5.36% accuracy while reducing reasoning length by 82.8%.The experiments compare deterministic and stochastic latent reasoning and highlight the value of exploration for difficult problems.
Reinforcement learning results: Without averaged rewards, Qwen-1.5B improves from 8.94% to 13.8% but reasoning length reaches its upper limit, while Llama-1B tends toward collapse.The averaged design encourages exploitation of more efficient reasoning pathways.

5 Limitations

CoLaR has demonstrated efficiency and effectiveness in latent reasoning, but its performance generally remains comparable to explicit CoT and its compression-factor generalization is limited.

CoLaR generally approximates explicit CoT performance without surpassing it on benchmarks other than GPQA.
CoLaR struggles with non-integer compression factors and factors exceeding the maximum training compression factor rmax.The authors attribute this limitation primarily to discrete tokenization constraints that restrict continuous compression-factor representations.
The authors recommend monitoring downstream applications because CoLaR may amplify reasoning biases or support convincing misinformation.

6 Conclusion

The paper presents CoLaR as a framework for dynamically compressing reasoning into latent space while retaining exploration and exploitation. It reports improved accuracy and substantially shorter reasoning chains, with further gains from reinforcement learning on MATH.

CoLaR compresses the semantics of multiple tokens into latent representations for more efficient reasoning.
Variable compression factors support flexible reasoning-chain lengths and fully parallelized processing during training and inference.
14.1% accuracy improvement over state-of-the-art latent-based reasoning methods accompanies a 53.3% shorter reasoning chain and 4.8% performance degradation relative to explicit CoT.
On MATH, reinforcement learning increases performance by 5.36% while reducing reasoning-chain length by 82.8%.

A More implementation details

The implementation uses LoRA-augmented language-model backbones, specified optimization settings, normalized latent-head targets, and several grade-school and advanced mathematics datasets.

Model hyperparameters: Experiments use frozen Llama-3.2-1B-Instruct or DeepSeek-distill-Qwen-1.5B backbones with tunable LoRA modules.The Latent Head is a three-layer MLP with hidden dimensions matching the backbone dimension d = 2048.
Training hyperparameters: SFT uses learning rate 1e-4 and RL uses 1e-6, with AdamW weight decay 1e-2 throughout.SFT uses distributed training across eight A100 GPUs with total batch size 256; RL uses a single A100 GPU.
Training hyperparameters: Latent-head targets are normalized by the embedding standard deviation during training and rescaled during inference.Because embeddings are approximately centered at zero, the procedure applies no mean shift.
Dataset information: Evaluation covers GSM8K-Aug, GSM8k-Hard, SVAMP, MultiArith, and the more challenging MATH dataset.Because MATH lacks an official validation set, 10% of its shuffled training samples are allocated for validation.

B Layer-wise analyses on compression factors

Layer-wise activation patterns vary with compression factor: smaller factors affect shallow layers more, whereas larger factors preserve stronger activity in deeper layers. The analysis links this pattern to the denser semantic content of highly compressed latents.

The analysis tests compression factors from 1 to 5 and measures relative activation-norm differences between consecutive LLM layers.The results are presented in Figure 6 using the same sample as Section 4.6.
Shallow layers show higher activation for smaller compression factors, intermediate layers behave similarly, and deeper layers retain stronger patterns for larger factors.
With c = 1, less informative token prediction requires minimal thinking, so shallow layers can determine the next token while deeper layers are underused.
Higher compression factors pack richer semantic content into each latent, requiring deeper layers to analyze condensed information and predict subsequent compressed latents.

C RL training curves

CoLaR’s RL training first expands latent reasoning to explore correct pathways, then shortens it while maintaining higher validation accuracy. Early stopping preserves the best checkpoint as overfitting begins.

Exploration: Accuracy rises from 9% to 14% while latent reasoning steps expand from 40 to 60 during initial exploration.GRPO encourages broader exploration to discover correct reasoning pathways.
Exploitation: Validation accuracy fluctuates between 14% and 16% while latent reasoning length decreases from 60 to 20 during exploitation.Per-token averaged reward/loss reinforces shorter effective reasoning pathways.
Early stopping: Early stopping preserves the best-performing checkpoint at approximately 4k steps when CoLaR begins to overfit.

D Scaling properties of CoLaR

CoLaR maintains predictable performance improvements as the underlying foundation model grows from 1B to 8B parameters, indicating scaling across tested model sizes.

Scaling properties: Performance improves predictably as CoLaR’s underlying foundation model increases from 1 billion to 8 billion parameters.The reported scaling behavior follows established neural scaling laws.