Technical Summary: AGCLR – Adaptive Gated Continuous Latent Reasoning

Problem Identification: The Concept Bottleneck

Large Language Models (LLMs) have demonstrated strong reasoning capabilities, particularly when guided by Chain-of-Thought (CoT) prompting. However, CoT is constrained to a single forward pass where each generated token becomes the input for the next, forcing early commitment to a reasoning path. To address this, the CoCoNuT (Chain of Continuous Thought) paradigm was introduced, enabling models to reason in a continuous latent space by feeding the last hidden state back as the next input embedding. This allows for implicit breadth-first search over reasoning paths without discrete tokens.

Despite its promise, the authors identify a critical limitation in vanilla CoCoNuT termed the concept bottleneck. In multi-pass inference, intermediate hidden states are progressively overwritten as new latent tokens are generated. Without a mechanism to preserve information, critical facts computed in earlier reasoning steps are lost as the reasoning depth increases. Empirical evidence shows that while vanilla CoCoNuT performs comparably to CoT on shallow tasks, its performance degrades on deeper reasoning chains (e.g., on GSM8K and ProsQA), failing to improve over standard CoT baselines on multi-hop QA tasks like HotpotQA.

Methodology: AGCLR

To resolve the concept bottleneck, the authors propose AGCLR (Adaptive Gated Continuous Latent Reasoning). This method augments CoCoNuT with a Gated Concept Stream, a persistent residual memory vector ($c_t$) maintained across all reasoning passes.

Architecture

At each latent token position, the model computes a hidden state ($h_t$) which is processed through three learned sigmoid gates to control information flow between the current state and the persistent memory:

1. Write Gate ($w_t$): Determines how much of the current gated hidden state ($h'_t$) is committed to the concept stream. This prevents low-confidence or irrelevant states from polluting the memory.
2. Read Gate ($r_t$): Controls how much of the prior concept stream ($c_{t-1}$) is retrieved and integrated into the current hidden state, allowing the model to access facts from earlier passes.
3. Forget Gate ($f_t$): Determines how much of the current hidden state is preserved versus replaced by the retrieved memory, enabling the selective pruning of irrelevant context.

The update equations are defined as:
$$\hat{h}_t = \text{LayerNorm}(h_t)$$
$$r_t, f_t, w_t = \sigma(W \hat{h}_t)$$
$$h'_t = (1 - f_t) \odot h_t + r_t \odot c_{t-1}$$
$$c_t = \text{LayerNorm}(c_{t-1} + w_t \odot h'_t)$$

The gated hidden state $h'_t$ replaces $h_t$ as the input for the next latent token position, while $c_t$ serves as the persistent memory for subsequent passes.

Training Protocol

AGCLR utilizes a multi-stage curriculum learning approach inspired by CoCoNuT:

• Stage 0: The model is trained on standard CoT data with explicit reasoning steps.
• Subsequent Stages: Explicit reasoning steps are progressively replaced with continuous latent tokens. At stage $k$, the first $k$ steps are replaced by $k \times c$ latent tokens.
• Final Stage: All reasoning is conducted in the latent space.
  The model is trained on GPT-2 (117M parameters) using this curriculum across three datasets: GSM8K (arithmetic), HotpotQA (multi-hop QA), and ProsQA (planning).

Key Contributions

1. Identification of the Concept Bottleneck: The authors empirically demonstrate that vanilla CoCoNuT suffers from information loss in intermediate states as reasoning depth increases, leading to performance degradation on complex tasks.
2. AGCLR Architecture: The proposal of a gated residual memory mechanism that resolves this bottleneck with only a 1.41% increase in parameters over the base GPT-2 model.
3. Performance Gains: AGCLR consistently outperforms vanilla CoCoNuT across all tested datasets, with the performance advantage compounding as the curriculum depth (reasoning complexity) increases.

Results

Evaluation on three benchmarks reveals the following:

• GSM8K (Math): AGCLR achieves 34.0% accuracy, outperforming vanilla CoCoNuT (31.4%) and the CoT baseline (40.6% is the CoT baseline, but CoCoNuT variants are generally lower on single-step math; AGCLR improves significantly over the latent-only baseline).
• HotpotQA (Multi-hop QA): AGCLR achieves 14.0% Exact Match (EM), a +3.6% improvement over vanilla CoCoNuT (10.4%) and a significant gain over the CoT baseline (11.0%).
• ProsQA (Planning): AGCLR reaches 96.0% accuracy, compared to 92.0% for vanilla CoCoNuT. Notably, vanilla CoCoNuT's performance degrades at the deepest curriculum stage (Stage 6), whereas AGCLR sustains improvement.

Analysis of Memory Retention

Cosine similarity analysis between the initial hidden state and subsequent passes shows that vanilla CoCoNuT suffers from monotonic memory decay (similarity drops from 1.0 to 0.126 by pass 6, representing an 87% information loss). In contrast, AGCLR stabilizes after pass 3, retaining a similarity of ~0.22, which corresponds to 71% more information retention than the baseline.

Ablation studies confirm that the gains stem from the persistent memory mechanism rather than just additional parameters:

• Removing the write gate drops performance to 8.8% EM (below vanilla CoCoNuT).
• Freezing the write gate after pass 2 results in only a 0.8% drop (13.2% vs 14.0% EM), suggesting that early information capture is sufficient, and the read/forget gates maintain the quality of this stored information.
• Removing the forget gate causes a significant drop (8.4% EM), highlighting the necessity of pruning irrelevant context to prevent memory pollution.

Significance and Claims

The paper claims that AGCLR directly addresses the concept bottleneck by introducing a persistent, gated residual memory that allows models to maintain critical facts across iterative reasoning cycles. Unlike previous approaches that operate in language space (e.g., iCoT, pause tokens) or lack persistent state, AGCLR enables reasoning in an unconstrained latent space while preserving the computational states necessary for multi-hop tasks.

The authors emphasize that the improvement is driven by the architecture's ability to store, maintain, and retrieve information across passes, rather than simply increasing model capacity. While the results are consistent across diverse reasoning tasks (math, QA, planning), the paper maintains modesty regarding its scope, noting that experiments were conducted on a single seed with a GPT-2 124M base model. Scalability to larger models and broader task distributions remains a subject for future validation.