Latent Chain-of-Thought Improves Structured-Data Transformers

Technical Summary: Latent Chain-of-Thought Improves Structured-Data Transformers

Problem Statement
Recent advancements in large language models (LLMs) have demonstrated that scaling test-time compute, particularly through Chain-of-Thought (CoT) reasoning, can significantly enhance reasoning capabilities, sometimes outperforming simple parameter scaling. While mechanisms like "Coconut" (feeding hidden states back as embeddings) and recurrent depth (iteratively applying shared transformer blocks) have been explored in natural language, it remains unclear if these benefits transfer to structured data domains. State-of-the-art models for tabular prediction (e.g., TabPFN, TabICL) and time-series forecasting (e.g., Chronos, TimesFM) currently operate via a single forward pass over structured observations. This paper investigates whether introducing latent chain-of-thought—multiple rounds of latent computation before prediction—can improve performance in these structured-data settings.

Methodology
The authors propose a recurrent scheme for structured-data transformers that allows for multiple rounds of computation before generating a final prediction. The core mechanism operates as follows:

1. Initial Pass: A transformer $f_\theta$ processes the input (context and query tokens) to produce hidden states.
2. Compression: The hidden states at the query positions ($H_q$) are compressed via a small Multi-Layer Perceptron (MLP) $\phi_\theta$ into "feedback tokens" ($Z$).
3. Recurrence: These feedback tokens are appended to the original input sequence. The transformer is then run again on this augmented sequence.
4. Iteration: This process repeats for $R$ recurrences. The final prediction is decoded from the hidden states at the original query positions after the $R$-th pass.

The study evaluates this approach across two domains:

• Tabular Prediction: Utilizing the three-stage architecture of TabICL (column-wise, row-wise, and in-context learning transformers). Recurrence is applied specifically at the in-context learning stage, where feedback tokens are appended to the row-embedding sequence.
• Time-Series Forecasting: Using a patch-based forecasting transformer. Future-query hidden states are compressed and appended to the input sequence before re-running the transformer stack.

Experimental Design and Baselines
To isolate the specific contributions of latent CoT, the authors train models from scratch on individual datasets and compare against three distinct baselines:

1. Same-Depth Baseline: A standard transformer with no recurrence ($R=0$) and the same number of layers. This controls for parameter count.
2. Deeper Baseline: A transformer with double the layers ($2L$) trained without recurrence. This matches the effective forward-pass depth and FLOPs of a CoT model with $R=1$ but doubles the parameter count, isolating gains from depth/capacity versus recurrence.
3. Looped Baselines: Models that share weights across passes (recurrent depth) but do not append feedback tokens to the input. These include weight-tied single-block looping and Universal-Transformer-style looping over multi-block stacks. This isolates the contribution of the explicit CoT tokens versus mere weight-tied recurrence.

Key Results
The approach was evaluated on 36 datasets: 9 from the LTSF-9 time-series benchmark and 27 from an OpenML subset used in TabPFN evaluations.

• Performance Gains: Latent CoT improved over the same-depth baseline on 8/9 time-series datasets (average gain of +10.99% in quantile loss) and 22/27 tabular datasets (average gain of +5.31% in accuracy).
• Comparison to Deeper Baselines: CoT models outperformed the deeper, non-recurrent baselines by an average of 21.01% (time-series) and 4.75% (tabular). The deeper baselines often performed worse than the original same-depth baseline on time-series tasks, suggesting overfitting on smaller datasets.
• Comparison to Looped Baselines: CoT models outperformed weight-tied looped models (which lack feedback tokens) by 7.74% (time-series) and 3.82% (tabular).
• Depth Scaling:
  • Time-Series: Performance improved across all tested depths ($R \in \{1, 2, 4, 8\}$), with no clear monotonic trend, though gains plateaued. Notably, a model trained at $R=4$ generalized to $R=8$ at evaluation.
  • Tabular: Performance showed a non-monotonic trend; a single recurrence ($R=1$) slightly degraded performance, but additional recurrences ($R=2, 4$) recovered and surpassed the baseline.

Significance and Claims
The paper claims that chain-of-thought is a useful axis for scaling test-time compute for structured data, distinct from simply increasing model depth or parameter count. The results demonstrate that:

1. Mechanism of Improvement: The gains are not solely due to increased compute or weight tying. The specific mechanism of compressing query-position hidden states into feedback tokens and reprocessing them is critical. The underperformance of looped baselines suggests that the benefit comes from the ability to "write" intermediate computation into a sequence that subsequent passes can attend to, rather than just iterative refinement of a hidden state.
2. Transferability: The latent CoT paradigm, previously successful in language modeling, successfully transfers to tabular and time-series forecasting tasks.
3. Limitations: The authors note that their findings are based on models trained from scratch; the efficacy of latent CoT on large-scale pretrained foundation models remains an open empirical question. Additionally, the current implementation uses a fixed recurrence depth per model rather than an adaptive mechanism per input, and the interaction with distribution shift is unexplored.