A*-Thought-V2: Efficient Latent Reasoning via Geometric Dynamics of LLM

AI summary

Overview

Research area: Efficient reasoning for large language models, specifically continuous or "latent" reasoning as an alternative to long chain-of-thought (CoT) text generation.

Technical level: Advanced. The paper assumes familiarity with transformer hidden states, PCA, cross-entropy training objectives, and KV caching during autoregressive inference.

Scope: This paper proposes A*-Thought-V2, a framework that decides which reasoning steps stay as explicit text and which are compressed into continuous latent tokens, based on a geometric analysis of the model's hidden-state trajectory.

What This Paper Is About

Chain-of-thought prompting makes LLMs better at reasoning but produces very long outputs that are expensive in compute and context length. Prior compression methods either delete whole reasoning steps outright (hard pruning, which loses information) or compress into continuous representations without a principled rule for what to keep. This paper's goal is a principled, information-preserving compression scheme: measure how much each reasoning step points toward the final answer, keep the aligned steps as text, and encode the rest as dense latent vectors.

Key Contributions

1. A geometric dynamics analysis of CoT trajectories. The authors project question, intermediate step, and solution hidden states into a 3D PCA space, define a global question-to-solution direction, and measure the angle between each local step transition and that direction. They identify six 30° directional-angle intervals with distinct semantic tendencies and three qualitative reasoning stages (exploration, convergence, refinement).

2. An explicit–implicit interleaved latent architecture. Rather than A*-Thought's hard pruning, redundant steps are replaced by continuous latent tokens interleaved with retained text steps, producing an information-dense sequence that keeps the reasoning progression.

3. Two training mechanisms: Embedding Forcing and Label Forcing. Embedding Forcing mean-pools all token embeddings of a redundant step into a single latent embedding; Label Forcing supervises that latent position with a soft, multi-modal vocabulary distribution (the average of the one-hot labels of the compressed step's tokens) instead of a hard one-hot target.

4. Extensive empirical validation across two model scales and six benchmarks, with reported gains in average accuracy, halved response length, higher Accuracy per Computation Unit (ACU), and large reductions in preprocessing and training time relative to A*-Thought and standard supervised fine-tuning.

Main Findings

• Up to 2.6-point accuracy gain and 2.29× ACU improvement. A*-Thought-V2 improves average accuracy by up to 2.6 percentage points over same-data SFT across both model scales. For Qwen3.6-27B with a 90° threshold, ACU rises from 0.35 to 0.80 over the original backbone, a 2.29× improvement.

• Shorter responses than SFT with higher accuracy. On Qwen3.6-27B, the 90° variant improves average accuracy by 1.0 point (92.9 to 93.9) while reducing response length by 16.0% (14023.88 to 11785.45 average tokens).

• Best overall average results in the main table. Qwen3.5-9B with a 90° threshold reaches 91.5 average accuracy and 0.68 ACU; Qwen3.6-27B with 90° reaches 93.9 average accuracy and 0.80 ACU. A*-Thought-V2 surpasses SwiReasoning, CopT, and A*-Thought in average accuracy and ACU.

• Semantically grounded thresholds generalize. The 90° threshold preserves forward-aligned steps while compressing checking, correction, and exploration into latent tokens; 60° additionally compresses mixed execution and transition. Both improve average accuracy and shorten responses over same-data SFT, with 90° giving the best trade-off.

• Angle intervals carry distinct semantics. Measured shares of steps: [0°, 30°) 7.57% (direct reduction/conclusion), [30°, 60°) 18.46% (routine derivation/checking), [60°, 90°) 24.72% (mixed execution/transition), [90°, 120°) 24.50% (condition checking/revision), [120°, 150°) 17.80% (reinterpretation/correction), [150°, 180°] 6.96% (reconsideration/branching).

• Geometry-guided selection beats alternatives. Against random-angle and reversed-selection variants (compression rates 48.16% and 47.15% vs 48.09% for the geometry-guided variant), the proposed selection achieves higher average accuracy and ACU while reducing average generation length by 7.54% and 12.91% respectively.

• Both forcing mechanisms matter. Removing Embedding Forcing, Label Forcing, or both reduces average accuracy (on AIME 2024/2025/2026) to 92.8%, 72.5%, and 61.7%, respectively, versus 94.5% for the full method.

• Large efficiency gains. Compression time drops from 5:16:22 to 0:16:57, a 94.6% reduction over A*-Thought. The 60° variant achieves the strongest compression rate (31.67%) and cuts training time by 80.3% (Qwen3.5-9B) and 68.7% (Qwen3.6-27B); the 90° variant cuts it by 69.3% and 51.8%.

• Latent states occupy a distinct region. PCA visualization on an AIME2024 case (case 79) shows latent tokens forming a compact cluster separate from the broader text-token distribution, occupying a coherent intermediate region along the reasoning trajectory.

• Latent positions show higher predictive entropy. Predictive entropy is higher at latent-token positions than at explicit-text positions for both model scales, which the authors attribute to fitting broader soft vocabulary targets rather than to uncertainty about a single next token.

• Training converges lower. A*-Thought-V2 starts with higher loss because it must learn latent reasoning segments, but converges rapidly and reaches the lowest final loss at both model scales.

• Larger latent length caps help. Increasing the maximum latent length generally yields higher accuracy with shorter responses, and the decoded latent-length distribution closely matches the training distribution, indicating the model learned the latent reasoning format.

• Extractor scale is not critical. Comparing Qwen3.5 extractors at 0.8B, 2B, 4B, and 9B parameters shows slightly different angle distributions but no clear downstream performance gap; Qwen3.5-0.8B is used as the default extractor.

Methodology in Plain English

The authors start from the intuition that a reasoning step either moves toward the answer or does something else (checking, backtracking, reinterpreting). To measure this, they take the hidden states the model produces for the question, each reasoning step, and the solution, and project them down to three dimensions with PCA. They define one global direction from the question to the solution, and for each step define a local direction from the previous step. The angle between the local and global directions acts as a "does this step move toward the answer" score.

That angle drives the compression decision. Steps whose angle is below a threshold stay as ordinary text. Steps above the threshold get compressed: all token embeddings of the step are averaged into one vector (Embedding Forcing), which sits in the sequence in place of the original tokens, wrapped by begin-of-latent and end-of-latent tags. Because one latent vector must stand in for many tokens, it cannot be trained against a single correct next token; instead the targets of all tokens in that step are averaged into one soft probability distribution (Label Forcing), and the latent position is trained against that soft distribution with a scaled cross-entropy term added to the normal text loss.

At inference, the latent positions are generated differently from text positions: instead of looking up an embedding for a token, the model feeds its own previous last-layer hidden state back in as the input vector, while the KV cache still attends over the entire preceding context. Once the latent span ends, normal text generation resumes.

Training used the OpenR1-Math-3k reasoning dataset, two backbones (Qwen3.5-9B and Qwen3.6-27B), 3 epochs on 8 NVIDIA A100 80GB GPUs, AdamW with a learning rate of 1.0×10⁻⁵, warmup ratio 0.1, total batch size 64, and cutoff length 20480. Evaluation used temperature 1.0, top-p 0.95, and 4 repeats per benchmark.

Why This Matters

Impact on research. The paper reframes CoT compression as a trajectory-geometry problem rather than an importance-scoring or search problem, and it offers a middle path between hard pruning and full latent reasoning. It also supplies a concrete training recipe (Embedding Forcing plus Label Forcing) for teaching a model to emit latent tokens that carry step-level semantics, plus representation-level evidence (PCA clusters, entropy profiles, training curves) that the compressed states are structured rather than arbitrary.

Real-world applications:

• Reducing inference cost for math and science tutoring systems that rely on long chains of reasoning.
• Fitting longer, more complex queries into fixed context windows in document analysis or multi-step problem solving.
• Lowering serving cost for reasoning-heavy assistants where output tokens are the dominant expense.
• Speeding up the data-preparation pipeline for teams that build their own reasoning training sets, since preprocessing time falls by 94.6% relative to A*-Thought.

Industry relevance. The reported ACU improvements (0.35 to 0.80 on Qwen3.6-27B) translate directly into serving economics: fewer generated tokens per correct answer means less GPU time per query. The 94.6% reduction in preprocessing time and up to 80.3% reduction in training time matter for teams that repeatedly re-curate and retrain reasoning models. The released artifacts (GitHub repository, Hugging Face model AStar-Thought-V2-OpenR1-Math-3k) lower the barrier to reproducing and adapting the approach.

Future Directions

• Reinforcement learning. The conclusion states explicitly that future work will explore reasoning under reinforcement learning, extending beyond the supervised fine-tuning paradigm used here.
• Threshold selection. The 60° and 90° thresholds are compared empirically; whether the optimal threshold can be chosen automatically, or per-problem rather than per-dataset, is left open.
• Interpretability of latent states. PCA shows latent tokens cluster separately, but what individual latent tokens encode and whether they are readable or editable is not established.
• Generalization beyond the tested setting. Evaluation covers four math benchmarks (Math500, AIME 2024, AIME 2025, AIME 2026) and two out-of-domain benchmarks (ARC-Challenge, GPQA-Diamond) on two Qwen backbones; behavior on other model families, modalities, and long-horizon agentic tasks is not reported.

Target Audience

This paper is most useful to machine learning researchers and engineers working on inference efficiency, chain-of-thought compression, and continuous or latent reasoning in LLMs, particularly those with the infrastructure to fine-tune 9B to 27B parameter models. It also suits practitioners who need to cut per-query token costs in production reasoning systems, and researchers interested in using hidden-state geometry as a diagnostic tool for how models reason. Readers without background in transformer internals, PCA, and supervised fine-tuning objectives will find the methodology section difficult, though the paper's central intuition about aligned versus deviating reasoning steps is accessible.

Authors’ abstract

Chain-of-Thought (CoT) improves the reasoning ability of Large Language Models (LLMs) but incurs substantial computation and context costs. Existing methods either lose intermediate information through hard pruning or lack a principled criterion for continuous compression. We present A*-Thought-V2, a geometric dynamics of LLM guided framework that models CoT as a hidden-state trajectory and replaces hard deletion with an explicit-implicit interleaved latent architecture. After projecting question, step, and solution representations into a 3D PCA space, it measures alignment between each local transition and global question-to-solution direction. Aligned steps remain explicit text, whereas deviating steps are compressed into continuous latent tokens. Directional angles capture both local semantics and reasoning dynamics: small angles indicate direct execution and answer formation, while large angles more frequently involve checking, correction, and branch exploration; their temporal variation reveals exploration, convergence, and refinement stages. To train this architecture, we introduce stepwise embedding forcing, which pools each redundant step into a single latent embedding, and label forcing, which supervises that latent token with a soft multi-modal vocabulary distribution instead of a hard one-hot label. Experiments on Qwen3.5-9B and Qwen3.6-27B across six in-domain and out-of-domain benchmarks show that A*-Thought-V2 improves average accuracy by up to 2.6% while reducing response length by up to half, increasing Accuracy per Computation Unit by 2.29$\times$, and reducing preprocessing and training time by 94.6% and up to 80.3%, respectively. Representation analyses suggest that latent states form a compact region distinct from textual states, while higher entropy at latent-token positions reflects broader soft targets that encourage richer step-level feature learning.