Source-linked AI summary
Soft Tokens, Hard Truths
Natasha Butt, Ariel Kwiatkowski, Ismail Labiad, Julia Kempe, Yann Ollivier
TL;DR
Continuous CoT reasoning is promising because mixtures of token representations may express multiple reasoning paths, but existing training approaches are inference-only, computationally costly, or dependent on discrete-CoT distillation. This work introduces scalable RL training with noisy soft tokens, matching discrete CoTs on pass@1 and improving pass@32 while preserving more of the base model’s out-of-domain behavior. Its best-performing setup trains continuously but uses discrete tokens for inference.
Problem
Existing continuous-CoT methods are limited by inference-only use, costly full-step training, or reliance on ground-truth discrete CoT traces.
Method
The paper introduces reinforcement-learning post-training for continuous CoTs using controlled noise on probability-weighted token mixtures, without ground-truth CoT annotations.
Results
Across mathematical reasoning benchmarks, continuous CoT training matches discrete-token training on pass@1 and improves pass@32, while the strongest configuration uses discrete CoTs at inference.
Takeaways & Limitations
Continuous CoT training provides a practical alternative for fine-tuning large models and better preserves out-of-distribution behavior than hard reasoning training.
Takeaways & Limitations
The hard-token model is deterministic and lacks the noise or random choices needed for direct Reinforce-style RL, while direct backpropagation through all CoT steps creates technical and memory challenges.
Abstract
from arXiv · showhide
The use of continuous instead of discrete tokens during the Chain-of-Thought (CoT) phase of reasoning LLMs has garnered attention recently, based on the intuition that a continuous mixture of discrete tokens could simulate a superposition of several reasoning paths simultaneously. Theoretical results have formally proven that continuous tokens have much greater expressivity and can solve specific problems more efficiently. However, practical use of continuous tokens has been limited by strong training difficulties: previous works either just use continuous tokens at inference time on a pre-trained discrete-token model, or must distill the continuous CoT from ground-truth discrete CoTs and face computational costs that limit the CoT to very few tokens. This is the first work introducing a scalable method to learn continuous CoTs via reinforcement learning (RL), without distilling from reference discrete CoTs. We use "soft" tokens: mixtures of tokens together with noise on the input embedding to provide RL exploration. Computational overhead is minimal, enabling us to learn continuous CoTs with hundreds of tokens. On math reasoning benchmarks with Llama and Qwen models up to 8B, training with continuous CoTs match discrete-token CoTs for pass@1 and surpass them for pass@32, showing greater CoT diversity. In systematic comparisons, the best-performing scenario is to train with continuous CoT tokens then use discrete tokens for inference, meaning the "soft" models can be deployed in a standard way. Finally, we show continuous CoT RL training better preserves the predictions of the base model on out-of-domain tasks, thus providing a softer touch to the base model.
1 Introduction
Continuous CoTs replace rigid intermediate language tokens with mixtures that can represent multiple reasoning paths, but practical training remains difficult. This work introduces RL training with controlled noise and reports competitive pass@1, stronger pass@32 diversity, and best performance when discrete CoTs are used at inference.
- Continuous reasoning represents a mixture of token embeddings and can implicitly follow multiple reasoning paths simultaneously, whereas standard CoT samples discrete tokens sequentially.
- Prior continuous-reasoning methods were limited by inference-only use, costly backpropagation through all reasoning steps, or distillation from ground-truth discrete traces.One cited approach was limited to six continuous CoT tokens because of computational constraints.
- The method trains continuous CoTs with controlled noise using reinforcement learning, without ground-truth discrete CoT annotations and with negligible overhead relative to discrete CoTs.The approach is evaluated with Llama-3.x and Qwen-2.5 models across GSM8K, MATH, DeepScaleR, and out-of-domain benchmarks.
- Continuous CoT training matches discrete-token CoTs for pass@1 and outperforms them on average for pass@32, indicating greater reasoning diversity.
- The strongest configuration trains with continuous CoTs but uses discrete CoTs during inference, allowing standard deployment of the resulting models.
2 Related Work
Related work spans inference-time soft tokens, post-training distillation, pretraining changes, and internal continuous computation. The literature motivates continuous CoTs through greater expressivity while also documenting adaptation and training constraints.
- A recurring challenge is that hard-token-trained language models may struggle to accept continuous representations outside their original training distribution.
- Inference: Inference-time methods replace hard tokens with probability-weighted embedding mixtures, but their reported gains depend on additional interventions and are not consistently confirmed.
- Post-training: Post-training methods include distillation from standard discrete CoTs, while Coconut requires ground-truth traces and computationally limits continuous reasoning to six tokens.
- Pretraining: Pretraining approaches introduce filler, latent, interspersed, or recurrent continuous tokens to support reasoning through internal activations.
- Theoretical arguments: Theoretical work proves that continuous CoTs can use superpositions to explore several reasoning paths simultaneously and solve graph reachability in O(n) rather than O(n^2).
3 Method
The method replaces discrete CoT sampling with probability-weighted token embeddings and adds input-embedding noise, making continuous reasoning trainable with reinforcement learning. It preserves standard hard-token generation for final answers while keeping computational overhead minimal.
- Soft thinking: During soft thinking, the model retains the full next-token probability distribution and feeds its probability-weighted embedding mixture into the transformer.This differs from hard-token generation, which samples a one-hot token and applies the process inductively.
- Transformer notation: The model obtains next-token probabilities by applying a temperature-controlled softmax to transformer outputs, then maps token distributions through the embedding matrix.The probability matrix is computed independently at each timestep before defining the next input embedding.
- Soft thinking: After the continuous CoT phase, the model samples normal hard tokens for the final answer.This retains a conventional discrete-token output stage despite continuous intermediate reasoning.
- Noisy soft thinking: Input-embedding noise makes continuous CoT amenable to reinforcement-learning algorithms by providing exploration absent from deterministic soft thinking.The method adds Gaussian noise with standard deviation σ > 0 to the soft-token computation; this enables Reinforce-like training.
- Reinforcement learning on soft tokens: Reinforce optimizes the expected answer reward by treating the noisy soft CoT and final answer as sampled model outputs, with the soft-CoT probability decomposed across timesteps.The approach supports standard Reinforce-like extensions, including RLOO, GRPO, and PPO, while requiring only probability-vector storage and first-layer noise injection.
4 Experiments
The experiments compare hard, soft, and fuzzy CoT training across Llama and Qwen models, datasets, inference settings, and out-of-domain benchmarks. Continuous-token training broadly matches hard training at pass@1, improves pass@32, and better preserves out-of-domain model behavior.
- Results: Soft and fuzzy training achieve broadly comparable pass@1 performance while providing a clear overall pass@32 advantage over hard training.The advantage is clearest for Llama models and is interpreted as preserved reasoning diversity.
- Inference settings: Hard inference generally performs best for both pass@1 and pass@32 across hard, soft, and fuzzy training methods.Thus, continuous-token training can be paired with standard discrete-token inference.
- Out-of-domain generalization: For Llama-8B trained on GSM8K, soft and fuzzy training recover MATH-500 performance to 44.6–44.7% greedy and 83.1–83.9% pass@32, versus 20.2% and 45.4% after hard training.The soft and fuzzy models maintain in-distribution GSM8K performance while improving this out-of-distribution result.
- Out-of-domain robustness: Soft and fuzzy training preserve the base model’s out-of-domain NLL better than hard training, especially on ARC and for Qwen on MMLU.Success rates remain comparable across the three training methods, while hard training degrades the base model NLL.
- Entropy analysis: Soft and fuzzy training retain roughly the base model’s entropy profile, whereas hard training lowers entropy toward greedy-sampling behavior.The authors connect this lower entropy with reduced pass@32, worse out-of-domain NLL, and occasional performance collapse.
5 Conclusion
The paper introduces reinforcement learning for continuous Chain-of-Thought training without ground-truth discrete CoT annotations or substantial computational overhead. Across mathematical reasoning benchmarks, continuous training matches discrete training for pass@1, improves pass@32, and better preserves out-of-distribution behavior.
- 5 Conclusion: Continuous CoT reinforcement learning requires neither ground-truth discrete CoTs nor substantial computational overhead.The framework is presented as the first reinforcement-learning approach for training continuous CoTs in LLMs under these conditions.
- 5 Conclusion: Continuous CoT training better preserves the base model’s out-of-distribution behavior than hard reasoning training.The conclusion describes this as a softer fine-tuning effect and evidence of behavioral differences between soft and hard reasoning processes.
- 5 Conclusion: Training uses Reinforce with a per-prompt leave-one-out group baseline over sampled CoT-and-answer sequences.For each prompt, the baseline averages the rewards of the other sampled sequences, while advantages are treated as constants during optimization.
- 5 Conclusion: The training prompt requires step-by-step reasoning before producing a boxed final answer, with hard and soft/fuzzy decoding using distinct stopping rules.Soft and fuzzy generation stop when a greedy shadow sequence reaches the final-answer phrase, while hard generation monitors generated text directly.
D Hyperparameter Search
Hyperparameters were selected using greedy validation performance, with shared learning-rate searches for hard and fuzzy training and an additional scale-factor search for fuzzy training.
- D Hyperparameter Search: All models used AdamW with a cosine learning-rate schedule and 20 warm-up steps.
- D Hyperparameter Search: Hyperparameters were tuned using greedy validation performance, and scale-factor values were searched only for fuzzy training.
- D Hyperparameter Search: Learning rates were swept from 1e−5 to 1e−6, with optimal values of 6e−6 for Llama 3B, 3e−6 for Llama 8B, and 8e−6 for Qwen 3B.The same optimal rates were found for hard and fuzzy training.
- D Hyperparameter Search: Fuzzy training additionally swept scale factors from 0.1 to 10 and found 0.33 best, although most values below 1 performed well.An ablation supports robustness to scale values below 1.
E Results on Soft and Fuzzy Inference
Across GSM8K, MATH-500, and OlympiadBench, hard inference performs best for all trained models, with no observed benefit from soft inference even after soft training.
- E Results on Soft and Fuzzy Inference: Hard inference achieves the best performance across base, hard-trained, fuzzy-trained, and soft-trained models.The comparison covers hard greedy, hard sampling, fuzzy greedy, fuzzy sampling, soft greedy, and soft sampling settings.
- E Results on Soft and Fuzzy Inference: Soft inference does not improve results on GSM8K, MATH-500, or OlympiadBench, contrary to previously reported benefits on hard-trained models.The reported evaluation includes six inference settings across the three test benchmarks.
- E Results on Soft and Fuzzy Inference: Table 3 compares GSM8K test performance by marking the best pass@1 and pass@32 results for each base-model and training-set pair.
- E Results on Soft and Fuzzy Inference: Table 4 applies the same best-pass@1 and best-pass@32 comparison to MATH-500 for each base-model and training-set pair.
F Ablations
All ablation experiments were conducted with two independent random seeds per experiment.
- F Ablations: Each ablation used two independent random seeds.
F.1 Noise Scale
Noise-scale training is robust across moderate settings, but excessive noise can collapse learning; placement and temperature also affect which variants learn effectively.
- Noise Scale: γ ≤ 1.0 preserves fuzzy-training performance, whereas γ = 3 causes a learning collapse from excessive noise.The scale factor γ is applied to the root mean square embedding norm to compute the noise scale σ.
- Noise Placement: Noise placement is most promising for top-k=5 logits, while final-hidden-layer and full-logit noise do not show comparable learning.The top-k=5 variant is the only placement showing learning similar to embedding noise, although test gains are inconsistent across metrics.
- Noise Placement: Top-k=50 noise placement can yield higher average test performance than the base model despite collapsed validation performance.The selected top-k=50 model was evaluated after only 50–100 training steps, so this result is narrower than the broader placement finding.
- Temperature: Fuzzy training remains robust across temperatures from 0.1 to 0.0001.Test performance is reported as comparable across these temperature values.
G.1 Training and Validation Performances
The experiments evaluate training and validation performance, pass@k, and CoT entropy across models, datasets, training methods, and inference modes. The supplied passages emphasize entropy preservation under soft or fuzzy training and practical training conditions.
- G.1 Training and Validation Performances: Training and validation success rates are reported for all RL-trained models across the Llama 3B, Llama 8B, and Qwen 3B configurations.The figures include models trained on MATH, DeepScaleR, and GSM8K datasets.
- G.1 Training and Validation Performances: Pass@k is evaluated across none, hard, fuzzy, and soft training methods and hard, fuzzy, and soft evaluation metrics.The comparisons cover GSM8K and MATH-500 settings across Llama and Qwen model–dataset combinations shown in Figures 11–15.
- G.3 Entropy Analysis: Soft or fuzzy training roughly preserves the base model’s CoT entropy profile under both greedy and sampled inference.Hard training produces a large change in the hard-sample entropy profile.
- G.3 Entropy Analysis: CoT entropy is measured at each token position as mean token-distribution entropy over non-pad test-generation tokens.Plots compare training and evaluation methods, with opacity indicating different seeds for the same method.
- G.1 Training and Validation Performances: Each RLOO run used eight H100 or H200 GPUs and required 48–96 hours end to end.Generation used a custom PyTorch-based function; paged attention is suggested for more optimized long-context or large-batch setups.