NeurIPS 2026

Geometric Latent Reasoning Induces Shorter Generations in LLMs

Shashi Kumar1,2,* · Yacouba Kaloga1,* · Petr Motlicek1,3 · Ina Kodrasi1 · Andrea Cavallaro2
1Idiap Research Institute, Switzerland 2EPFL, Switzerland 3Brno University of Technology, Czech Republic * Equal contribution · shashi.kumar@epfl.ch, yacouba.kaloga@idiap.ch
Paper (PDF) arXiv 2606.02248 Code soon BibTeX

token embedding local neighborhood discrete CoT step learned latent update
TL;DR

We treat chain-of-thought as a path through the model's token-embedding space and train a tiny linear transition head to predict the next step's direction. At inference, the model takes K continuous latent steps before resuming normal decoding and needs ~80% fewer generation steps to reach a correct answer. We don't add a length penalty. The shorter generations emerge on their own.

~80%
fewer generation steps
Median reduction of 77–81% on GSM8K, MATH500 and SVAMP (Qwen3-1.7B, GLR-10), including latent steps.
5.3×
lower median latency
MATH500 on an H100, both methods in the same HuggingFace stack: 45.7 s → 8.6 s.
+66pp
accuracy under a tight budget
Qwen3-8B on GSM8K with a 512-step budget: 15.5% (CoT-SFT) → 81.6% (GLR).
0
extra models
No teacher, reward model or auxiliary network. GLR adds only a lightweight linear transition head.

Drop into an existing CoT-SFT pipeline

What it takes to train GLR on your own model.

What you need

  • Your existing CoT traces (10K in our experiments)
  • One linear layer (~4M params on Qwen3-1.7B)
  • Full fine-tuning on a single H100, up to 8B
  • K is a dial you set at inference, no retraining

What you don't need

  • No teacher model or distillation
  • No length penalty or RL
  • No multi-stage curriculum

The one extra cost: training runs two forward passes per example.

Abstract

Latent Reasoning in Model's embedding space

Large language models solve complex problems by generating lengthy chains of explicit reasoning tokens. While effective, this makes reasoning expensive, length-sensitive, and constrained to (discrete) natural language. While latent reasoning offers a continuous alternative, determining useful structures for intermediate latent states is an open challenge.

In this paper, we formulate latent reasoning as a geometric path-approximation problem within the model's pretrained token-embedding space. We introduce Geometric Latent Reasoning (GLR), which uses a lightweight transition head to predict iterative direction updates in embedding space. Using textual chain-of-thought traces as anchors, GLR learns to approximate discrete reasoning trajectories while permitting continuous deviations from exact token embeddings.

Evaluations on mathematical reasoning benchmarks using Qwen3 models reveal an emergent phenomenon: geometric latent reasoning induces substantially shorter generations without an explicit length objective. By replacing early explicit reasoning with continuous latent steps, models often reach correct answers using substantially fewer total generation steps. These findings suggest that continuous trajectories act as compact intermediate reasoning states, exposing a new tradeoff between latent computation budget, output length, and accuracy.

Method

Chain-of-thought is a discretized path. We learn to traverse it continuously.

Every token the model writes is fed back as its input embedding, so a textual reasoning trace traces out a trajectory in embedding space. GLR learns the local displacement along that trajectory. At inference it follows these learned directions for a few steps without committing to any vocabulary token.

1

CoT as an embedding-space trajectory

A reasoning trace <think> t1:m </think> a1:ℓ induces a sequence of input embeddings, and we treat this sequence as the path the model reasons along.

$$\mathbf{e}_i^t = E_{\text{in}}(t_i), \qquad \mathbf{e}_1^t,\ \mathbf{e}_2^t,\ \ldots,\ \mathbf{e}_m^t$$

The Transformer is continuous in its input, so small moves around a token embedding give nearby hidden states. Soft-token methods show that off-vocabulary inputs can still support coherent reasoning.

2

A lightweight transition head

A single linear layer $g_\phi:\mathbb{R}^d\!\to\!\mathbb{R}^d$ reads the last hidden state and predicts the displacement to the next reasoning embedding, not the embedding itself.

$$\mathcal{L}_{\Delta} = \frac{1}{m}\sum_{i=1}^{m} \gamma^{\,i-1}\,\big\lVert g_\phi(\mathbf{h}_{i-1}^t) - (\mathbf{e}_i^t - \mathbf{e}_{i-1}^t)\big\rVert_2^2$$

The discount $\gamma \le 1$ (0.999) relaxes the penalty for later positions, so the latent path is allowed to drift further from the literal text as reasoning proceeds.

3

Two-pass training with latent replacements

Pass 1 runs the original sequence and collects the true and predicted displacements. Pass 2 replaces every thought embedding with its latent version and trains on

$$\hat{\mathbf{e}}_i = \mathbf{e}_{i-1} + g_\phi(\mathbf{h}_{i-1}^t), \qquad \mathcal{L} = \mathcal{L}_{\text{CE}}^{\text{answer}} + \lambda\,\mathcal{L}_{\Delta}$$

Cross-entropy is applied only to answer tokens, never at latent positions. $\mathcal{L}_\Delta$ anchors the path, and leaving out CE there lets latent states sit between vocabulary tokens. Input embeddings stay frozen, so the displacement targets do not move during training.

4

K latent steps, then ordinary decoding

After <think>, the model rolls out $K$ continuous updates. Each one is fed back as the next input, skipping the ~300M-weight vocabulary head. Normal token decoding then resumes.

$$\hat{\mathbf{e}}_i = \hat{\mathbf{e}}_{i-1} + g_\phi(\mathbf{h}_{i-1}), \quad i = 1,\ldots,K$$

$K$ can be changed at inference time without retraining, and it sets how much of the reasoning budget goes to latent computation versus explicit text.

GLR two-pass training pipeline: first forward pass collects exact and predicted embedding displacements; second forward pass replaces thought embeddings with latent states and trains with answer cross-entropy plus the transition objective.
Training pipeline. Left (1st forward): the model processes the discrete sequence and collects the exact displacements Δk and the predicted displacements Δ̂k between consecutive reasoning tokens. Right (2nd forward): thought embeddings are replaced by continuous latent states ê built from the first-pass predictions. The objective combines answer-token cross-entropy with a transition loss that anchors the second-pass predictions to the true discrete updates. The embedding layer is frozen.

Results

Shorter generations emerge, and accuracy holds up under tight budgets

We fine-tune Qwen3 models on 10K math traces from Open-R1 Mixture-of-Thoughts and compare against CoT-SFT trained on the same data. Generation length counts every step after the prompt, latent steps included, which is a conservative accounting for GLR.

Accuracy vs generation budget and generated-token distributions on GSM8K
Where CoT still wins. At the largest budgets (2048 steps on GSM8K, 4096 on MATH500), CoT-SFT catches up and often surpasses GLR in final accuracy. GLR's gains are concentrated in the constrained-budget regime, which suggests the two modes are complementary. Latent steps compress early reasoning, and explicit tokens are a more stable scratchpad when the budget is large.

Understanding latent dynamics

The learned directions carry the reasoning, not the extra positions

Is the gain just from inserting a few extra positions? Is GLR secretly emitting the answer early? Several controls separate these explanations.

1 · Turn latent rollout off (K = 0)

Take the same GLR-trained model, with the same training recipe and loss masking, and simply decode with exact discrete tokens from the first step. Generations jump back to ~1,000 tokens. A few learned latent steps (K = 5 or 10) bring the median under 200.

Box plot of generated tokens on GSM8K for GLR-0, GLR-5 and GLR-10: medians near 1000, 170 and 160.

Qwen3-1.7B on GSM8K, correct answers only. The reduction appears only when the learned continuous transitions are actually used.

2 · Replace the learned directions

All controls keep K = 10. Zero / repeated-<think> keep the positions but remove movement. A norm-matched random head and norm-preserving noise keep the movement but scramble its direction.

Accuracy vs. median total steps

upper-left is better · 100 fixed examples per benchmark

Removing movement keeps accuracy but raises median length from 183.5 to 948 steps. Random or noisy directions shorten outputs but lose 19.3 and 13.7 accuracy points. Only the learned directions give the favorable tradeoff.

3 · Where does the latent state land?

We teacher-force the CoT-SFT trace through the GLR model and align each latent state to the discrete trajectory in hidden-state space. Progress is how far along the CoT the final latent state lands (0 to 1).

BenchmarkProgress: correct / incorrectDrift: correct / incorrectFutureHit@10
GSM8K0.905 / 0.8570.985 / 0.9840.20%
MATH5000.810 / 0.7090.989 / 0.9870.00%
SVAMP0.892 / 0.8520.987 / 0.9820.13%

Correct generations get further along the reasoning trajectory: +1 SD of progress corresponds to an accuracy odds ratio of 1.31 [1.21, 1.45], while drift is not predictive (OR 1.08 [0.99, 1.19]). FutureHit@10 is only 0–0.2%, so latent states are genuinely continuous rather than copies of upcoming token embeddings.

4 · Not an early-answer shortcut

Learned GLR closes the reasoning span immediately (</think> right after the latent phase) in only 8% of cases. In the other 92% it keeps reasoning explicitly, but from a more advanced state that needs far fewer discrete steps.

A clarification on the word “shortcut”: GLR does not jump to an endpoint. Each of the K latent states is generated autoregressively and stays in context for all later decoding. The shortcut is past redundant discrete transitions, while the computation remains sequential and contextual.

5 · How many latent steps?

The effect of $K$ is non-monotonic. Moderate values (5–20) give the best tradeoff. Very long uninterrupted rollouts (K = 80 or 100 at 1.7B) lower accuracy, likely because a locally trained transition model accumulates drift without discrete grounding. Choosing K adaptively per problem is a natural next step.

Qualitative example

From latent steps back to text, mid-solution

These are real Qwen3-1.7B generations on a GSM8K problem. The orange squares are continuous latent steps. They are fed back as inputs and never decoded into text. After 20 of them, the model's first explicit tokens are 54 - 20 = 34., the final operation of the solution, as if the earlier steps had already been worked out in latent space.

So the latent prefix does more than shorten the visible text. It moves the model into a partially advanced reasoning state before standard decoding resumes. With larger K the model closes </think> right away and writes a compact solution.

GSM8K · Qwen3-1.7B
user
The ratio of men to women on a bus is 5:9. If the total number of passengers on the bus is 84, and 20 women alight from the bus at the next stop, how many women have remained on the bus?
assistant
<think>
</think>

Notes

Details & limitations

How is efficiency measured, and is it fair to GLR's baseline?

The primary metric counts every post-prompt step: for GLR that is the K latent steps plus all text tokens. This is conservative for GLR, since a latent step still needs a full Transformer forward pass but replaces the ~300M-weight vocabulary projection with a ~4M-parameter transition head. Wall-clock measurements with both methods in the same HuggingFace stack confirm that the savings are real (4.6–5.3× lower median latency).

Why freeze the embedding layer?

The transition head learns displacements between input embeddings. If the embeddings were updated at the same time, the targets would be non-stationary and training would be unstable. CoT-SFT is frozen too for a matched comparison, and an unfrozen CoT-SFT baseline yields the same conclusions (see “vs. latent baselines”).

Does it work outside math, or with much larger models?

We show consistent gains from 0.6B to 8B across six math benchmarks, ranging from arithmetic word problems to olympiad-level problems. Science, code, theorem proving, planning and multi-hop reasoning are untested so far. Training so far uses at most 10K CoT traces, which may understate what GLR can do.

What are the main limitations?

K is fixed before the latent rollout instead of adapting per problem, and very long latent rollouts drift. At large generation budgets, CoT-SFT can reach higher final accuracy. Replacing explicit reasoning with continuous states also makes part of the computation less interpretable, so GLR-style inference should be paired with answer verification and drift monitoring.

What comes next?

We plan learned stopping criteria that choose K per problem, larger training corpora, diffusion- or flow-style latent trajectory models (where test-time scaling becomes sampling multiple continuous paths), and evaluation beyond mathematics.

Citation

BibTeX

The paper has been accepted at NeurIPS 2026. If you find this work useful, please cite:

@article{kumar2026geometric,
  title={Geometric Latent Reasoning Induces Shorter Generations in LLMs},
  author={Kumar, Shashi and Kaloga, Yacouba and Motlicek, Petr and Kodrasi, Ina and Cavallaro, Andrea},
  journal={arXiv preprint arXiv:2606.02248},
  year={2026}
}