Why Limit the Residual Stream to Layers and Not Tokens? Persistent Memory for Continuous Latent Reasoning
Abstract
Large language models (LLMs) have demon- strated remarkable reasoning abilities on math- ematical and multi-hop planning tasks. The CoCoNuT (Chain of Continuous Thought) paradigm (Hao et al., 2024) extends this by en- abling models to reason in latent space, exploring multiple reasoning paths simultaneously rather than committing to a single chain early on. How- ever, we identify a limitation we term the con- cept bottleneck. At each reasoning pass, inter- mediate hidden states are overwritten, causing the model to lose critical facts computed in ear- lier steps as reasoning depth increases. We ob- serve this empirically. On HotpotQA, vanilla Co- CoNuT (10.4% EM) fails to improve over the CoT baseline (11.0% EM), and performance de- grades with curriculum depth on GSM8K. To ad- dress this, we propose AGCLR (Adaptive Gated Continuous Latent Reasoning), which augments CoCoNuT with a Gated Concept Stream. A per- sistent residual memory maintained across all rea- soning passes, controlled by three learned gates: a write gate that commits intermediate facts to memory, a read gate that retrieves relevant prior states, and a forget gate that prunes irrelevant context. Evaluated on GSM8K, HotpotQA, and ProsQA using GPT-2 as our base model, AG- CLR achieves consistent improvements across all types of datasets. With the performance gap compounding as curriculum depth increases, di- rectly resolving the concept bottleneck. Code available at https://anonymous.4open. science/r/JJJJ/README.md