Improving Latent Reasoning in LLMs via Soft Concept Mixing

TL;DR

This paper introduces Soft Concept Mixing (SCM), a training scheme that enhances large language models' latent reasoning by integrating soft concept representations, leading to improved reasoning performance while maintaining training stability.

Contribution

SCM is a novel training method that incorporates soft concept vectors into LLMs' hidden states, bridging the gap between soft reasoning and discrete token training.

Findings

SCM improves reasoning accuracy on five benchmarks.

SCM maintains stable training dynamics.

Soft concept integration enhances latent reasoning in LLMs.

Abstract

Unlike human reasoning in abstract conceptual spaces, large language models (LLMs) typically reason by generating discrete tokens, which potentially limit their expressive power. The recent work Soft Thinking has shown that LLMs' latent reasoning via soft concepts is a promising direction, but LLMs are trained on discrete tokens. To reduce this gap between the soft concepts in reasoning and the discrete tokens in training, we propose Soft Concept Mixing (SCM), a soft concept aware training scheme that directly exposes the model to soft representations during training. Specifically, SCM constructs a soft concept vector by forming a probability-weighted average of embeddings. Then, this vector is mixed into the model's hidden states, which embody rich contextual information. Finally, the entire latent reasoning process is optimized with Reinforcement Learning (RL). Experiments on five reasoning benchmarks demonstrate that SCM improves the reasoning performance of LLMs, and simultaneously maintains a stable training dynamic.