Vishal Verma Labs
← All articles
AI Architecture & Mathematical Optimization

Sinkhorn-Aligned Latent Tokens Add Nothing Over Fine-Tuning: A Controlled Negative Result on DeepSeek-R1-Distill-Qwen-1.5B

/Vishal Verma20 min read

Abstract

Our earlier paper, Sinkhorn Alignments [1], proposed replacing a reasoning model’s long chain of thought with a few continuous “pause” tokens. Those tokens are refined by a looped computation (LOTUS) and trained with Sinkhorn optimal transport to match the model’s own reasoning. This paper tests that proposal.

On DeepSeek-R1-Distill-Qwen-1.5B [2] we train K=4K=4 latent tokens with R=4R=4 loop iterations. We score them on all 1,319 GSM8K test problems [3] and compare them with a control trained identically with no latent tokens (K=0K=0).

The result is negative:

GSM8K test, 1,319 problemsAccuracy95% CI
Untrained model, answering directly11.22%9.55–12.96
Fine-tuned control, no latents (K=0K=0)19.79%17.66–21.99
K=4K=4 latent tokens18.42%16.38–20.55
K=4K=4 latent tokens, collapse partly fixed19.56%17.44–21.68
Full reasoning (first 448 problems)86.16%82.81–89.29

The latent tokens add nothing measurable over the control. The paired difference is −1.36-1.36 points (95% CI −3.03-3.03 to +0.23+0.23, exact McNemar p=0.12p = 0.12).

A diagnosis explains why:

  • The four latents collapse into one. Their cosine similarity within a problem is 0.999.
  • They mostly re-encode the question. They identify their own problem 98.8% of the time, but the question’s final hidden state does that 96.1% of the time without them.

The alignment loss is part of the cause. Its targets are dominated by a component that every state shares, and by a few very large channels. We standardised the targets and sharpened the transport plan. That partly separated the latents (cosine 0.84), but the model then scores 19.56%, level with the control (p=0.86p = 0.86).

In both runs the latents never reached even the loss that four identical latents would achieve. So the bottleneck is what the latent pathway can produce from the question alone, not only the loss. Every number here is computed from saved run outputs.


1. The question and the control

An R1-style reasoning model spends most of its inference compute writing out its reasoning. On the GSM8K problems where we measured it, DeepSeek-R1-Distill-Qwen-1.5B generated a mean of 2,174 tokens per problem and scored 86%. Forced to answer immediately, it generated a median of 5 tokens and scored 11%.

The proposal in [1] is to keep most of that 75-point gap while paying for only a handful of extra positions:

  1. Insert KK learnable latent tokens where the reasoning would go.
  2. Refine them by looping the network over them.
  3. Train them so that they carry what the written reasoning carried.

This builds on pause tokens [4] and continuous-thought training [5].

Training the latent model necessarily fine-tunes the backbone on GSM8K answers, and that fine-tuning helps on its own. A comparison against the untrained model would credit that gain to the latents. Every latent result here is therefore paired with a K=0K=0 control: the same LoRA adapters, data, schedule and number of answer-training steps, with no latent positions. The question this paper answers is the paired difference K=4K=4 minus K=0K=0.

2. Method as implemented

The code is the open-source latent-harness library, release v0.1.0. This section describes exactly what it does.

2.1 Sequence layout

The latent model uses the direct-answer prompt with the latent tokens inside an empty reasoning block:

<question> <|Assistant|> <think>\n\n [z_1 … z_K] </think>\n\nThe final answer is $\boxed{ … }$.

With K=0K=0 this is exactly the direct-answer condition, so the K=0K=0 control and the untrained floor use the same prompt and the same 48-token answer budget.

2.2 The LOTUS loop

The KK latent positions are refined for RR iterations, all attending to the question’s cached keys and values Cpre\mathcal{C}_{\text{pre}}:

h(0)=0,h(t)=fθ ⁣(E+g(h(t−1)) | Cpre),Z=h(R)h^{(0)} = 0, \qquad h^{(t)} = f_\theta\!\left(E + g(h^{(t-1)}) \,\middle|\, \mathcal{C}_{\text{pre}}\right), \qquad Z = h^{(R)}

  • fθf_\theta is the whole backbone (frozen, plus LoRA rank 16 on every attention and MLP projection [6]). It returns final post-norm hidden states.
  • E∈RK×dE \in \mathbb{R}^{K \times d} are learnable pause embeddings, initialised with the statistics of real token embeddings.
  • g(h)=γ semb W (h/rms(h))g(h) = \gamma\, s_{\text{emb}}\, W\, (h / \mathrm{rms}(h)) maps hidden states back to embedding scale. WW is a d×dd \times d matrix under spectral normalisation [7] (∥W∥2≤1\lVert W \rVert_2 \le 1) and γ=0.9\gamma = 0.9.

Spectral normalisation is applied only to the feedback map WW, not to the transformer’s weights. It does not by itself make the loop contractive. We did not measure the loop’s Jacobian spectral radius; Section 5 reports how the loop’s residuals behave in practice.

The answer is decoded with the final iteration’s latent keys and values in context. With d=1,536d = 1{,}536, K=4K=4 and R=4R=4 the model has 20.8M trainable parameters: 18.5M LoRA, 2.4M feedback map and 6,144 pause-embedding values.

2.3 Alignment targets

The targets are the model’s own reasoning:

  1. We sampled reasoning traces on the first 2,500 GSM8K training problems (temperature 0.6, top-p 0.95, at most 4,096 tokens). The model solved 85.52% of them.
  2. We kept one correct, finished trace per problem: 2,136 traces. 362 were incorrect or unfinished, and 2 had no boxed answer.
  3. The kept reasoning sections average 1,236 tokens (median 1,034, 90th percentile 2,207, max 3,865).

For each trace, the frozen backbone’s final hidden states over the reasoning tokens are mean-pooled into at most 128 contiguous windows. That gives 273,361 target states of dimension 1,536.

2.4 Losses and training

Stage 1 (alignment). The loss is the entropic optimal-transport cost [8] between the KK latents ZZ and the trace’s pooled states HH:

  • The cost is squared Euclidean, after dividing both by the batch’s mean target norm.
  • Both sides have uniform marginals.
  • The plan comes from 50 log-domain Sinkhorn iterations at ε=0.1\varepsilon = 0.1.
  • The plan is held fixed when differentiating.

Stage 2 (answer). The loss is the answer’s cross-entropy, plus the stage-1 transport cost (weight 1.0), plus a frozen-head keyword loss (weight 0.1). The keyword loss decodes each latent through the frozen LM head and scores it against up to 32 token ids of the trace’s 16 most salient keywords.

SettingValue
Batch4 problems × 4 accumulation steps = 16 per optimizer step
OptimizerAdamW, lr 10−410^{-4}, 30 warm-up steps, cosine to 10−510^{-5}, gradient clip 1.0
Stage 1 / Stage 2 steps100 / 400 (stage 2 is 3.05 epochs of the 2,097 training problems)
Validation39 held-out training problems (2%)
K=0K=0 control400 stage-2 steps, no transport or keyword loss (no latents to apply them to)
HardwareApple M4, 16 GB unified memory, PyTorch MPS, bfloat16, seed 0
Time8.0 s per stage-1 step, 9.5 s per stage-2 step (5.8 s for K=0K=0); 1.3–1.4 h per latent run, 0.7 h for K=0K=0

2.5 Evaluation

Protocol:

  • All 1,319 GSM8K test problems.
  • One sample per problem at temperature 0.6 and top-p 0.95, the model’s recommended settings.
  • The answer is the last \boxed{} expression, graded against the reference.

Statistics:

  • Each confidence interval is a bootstrap over problems (10,000 resamples).
  • Each comparison between two runs is paired on the same problems, with a bootstrap CI on the difference and an exact McNemar test on the problems where they disagree [9].

Reasoning baseline: the full-reasoning condition (16,384-token budget) has so far been run on the first 448 test problems. Section 3.2 compares every condition on exactly those 448.

3. Results

3.1 The latents add nothing

Comparison (1,319 problems)Difference95% CIOnly first rightOnly second rightpp
K=0K=0 control − untrained+8.57+6.44 to +10.6916148below 0.001
K=4K=4 − K=0K=0 control−1.36−3.03 to +0.2350680.117
K=4K=4 with fix − K=0K=0 control−0.23−1.90 to +1.4461640.858
K=4K=4 with fix − K=4K=4+1.14−0.30 to +2.6557420.159

Fine-tuning on 2,097 answers lifts the direct-answer model by 8.6 points. Adding latent tokens to that fine-tuning changes nothing we can detect.

The two latent models generate the same short answers as the control: a mean of 5.4 tokens, with every answer finished. The four latent positions and four loop iterations are the only extra compute.

3.2 Against full reasoning

On the 448 problems with all five conditions:

Figure 1. GSM8K accuracy on the first 448 test problems, where all five conditions have been run. Full reasoning generates a mean of 2,174 tokens; the other four generate about 5 answer tokens. The latent models sit with the fine-tuned control, not between it and full reasoning.
Condition (448 problems)AccuracyShare of the reasoning gap closed
Untrained, direct10.94%—
Fine-tuned control, K=0K=022.10%14.8%
K=4K=4 latents19.42%11.3%
K=4K=4 latents, collapse fix21.43%13.9%
Full reasoning86.16%100%

Full reasoning beats the best latent model by 64.7 points (95% CI 60.0–69.4). It solves 297 problems the latent model misses; the latent model solves 7 that it misses. Fine-tuning recovers about a seventh of what reasoning buys, and the latents recover none of it on top.

4. Why: a diagnosis of the latent states

To see what the latents learned, we ran each trained model on 256 training-set problems: the 39 held out from training plus 217 sampled at random. We compared each problem’s final latent states ZZ with its alignment targets. All numbers in this section come from that analysis or from CPU simulations of the loss on 64 random training traces.

4.1 The four latents collapse into one

In the K=4K=4 model the four latent vectors of a problem have mean pairwise cosine 0.999. On the batch’s 16 latent vectors (4 problems × 4 slots), effective rank [10] stays near 4 throughout training (3.56 at the start of stage 2, 3.85 at the end). So the model behaves like one latent token repeated four times.

The frozen-head keyword loss, which [1] proposed as the guard against this collapse, was on in every latent run. It fell from 7.90 to 5.00 but did not prevent the collapse. We did not train without it, so we cannot say whether it helped at all.

4.2 The latents carry the question, not the reasoning

The latents are problem-specific: their mean matches the right problem’s mean target state out of 256 candidates 98.8% of the time (chance 0.4%), and 100% on the 39 held-out problems.

A control shows this is mostly the question, though. The backbone’s final hidden state at the end of the question, computed with no latents at all, matches the right problem 96.1% of the time (97.4% held-out). After mean-centring, the latents’ cosine with that question-end state is 0.56.

What the latents add beyond the question is small, and nothing in the accuracy results suggests it is useful.

4.3 The loss pulls towards the mean

With the transport plan held fixed, the cost is minimised by moving each latent to the plan-weighted mean of the targets it is coupled to. So the plan’s sharpness decides whether the KK latents split the reasoning trace into parts or all sit at its mean. Two properties of the targets blur that plan:

  • A shared component dominates. On average, 89% of the energy of a problem’s mean target state is the direction every problem shares.
  • A few channels are huge. Five of 1,536 channels carry 46% of the target states’ energy, consistent with the massive activations reported in LLMs [11].

Distances between targets are therefore small relative to their norms, and at ε=0.1\varepsilon = 0.1 the plan is close to uniform. Our CPU simulation compares the loss with all four latents at the trace mean against the loss with them split at the loss’s own optimum. Splitting lowers the loss by only 12.5%, so there is little pull away from collapse.

Target space, ε\varepsilonLoss saved by splittingSlot cosine at the loss optimumTrace-position spread at optimum
Raw, 0.10 (the K=4K=4 run)12.5%0.940.078
Raw, 0.0328.6%0.870.142
Standardised, 0.1017.4%0.89 raw / 0.49 standardised0.108
Standardised, 0.03 (the fix)21.0%0.87 raw / 0.42 standardised0.12–0.13

Trace-position spread is the standard deviation, across the four slots, of the average position (0 = start, 1 = end) in the reasoning trace that each slot is coupled to.

The trained K=4K=4 model is more collapsed than even its loss’s optimum: cosine 0.999 against 0.94, and spread 0.034 against 0.078. So the loss alone does not explain the collapse.

4.4 A fix that partly works

The fix trains on per-channel standardised targets, using the mean and standard deviation over all 273,361 states, with ε=0.03\varepsilon = 0.03. Everything else is unchanged.

After trainingK=4K=4 (original)Fix, after stage 1Fix, finalLoss optimum (fix)
Slot cosine, raw0.9990.9640.9700.87
Slot cosine, standardised0.9930.7670.8420.42
Trace-position spread0.0340.0710.0830.12–0.13
Effective rank of 16 latents (stage 2)3.85—4.97—
Latents match own problem98.8%98.0%98.4%—
Question-end state matches (control)96.1%89.8%93.0%—
GSM8K test accuracy18.42%—19.56%—

The latents separate, and they separate in order: slot 1 is the most problem-specific and slot 4 is still almost generic. But they get only partway to the loss’s optimum, and stage 2 pulls them back together somewhat (standardised cosine 0.77 → 0.84). Accuracy moves from 1.4 points below the control to level with it. Neither change is statistically significant.

4.5 The latent pathway cannot produce the targets

The most telling number is the loss itself. In both runs, stage 1 ended above the loss of the collapsed solution, where all four latents sit exactly at the trace mean:

RunStage-1 loss, last 30 stepsLoss if all latents sat at the trace meanLoss at the split optimum
K=4K=4 (raw, ε=0.1\varepsilon = 0.1)0.3450.3080.270
Fix (standardised, ε=0.03\varepsilon = 0.03)0.8700.7840.619

The simulated values come from 64 random training traces, and the training values from training batches of the same distribution, so the comparison is close but not exact. Both runs’ losses stopped improving after about 70 of the 100 alignment steps.

So the model cannot even place its latents at the mean of each problem’s reasoning states, let alone spread them along the trace. A better loss cannot remove that limit.

We did not isolate the cause. The most direct explanation is about what the targets are. They are the model’s hidden states after it has read its own intermediate reasoning, while the latents must be computed from the question alone, in four positions and four passes. Producing those states requires doing the reasoning, which is exactly what the latents were meant to replace. Smaller contributing factors could be:

  • the capacity of rank-16 adapters;
  • the short 100-step alignment stage;
  • the loop settling all four slots towards similar fixed points.

5. Training dynamics

The loop converges in practice. We measured the relative change between successive iterations, ∥h(t)−h(t−1)∥/∥h(t−1)∥\lVert h^{(t)} - h^{(t-1)} \rVert / \lVert h^{(t-1)} \rVert:

WhenIteration 1→22→33→4
Start of training0.980.700.53
End of stage 1 (K=4K=4)0.600.340.038
End of stage 2 (K=4K=4)0.550.0720.008
End of stage 2 (fix)0.430.0330.007

By the fourth iteration the update is under 1% of the state, so more iterations would change little. We did not measure the Jacobian’s spectral radius or train without spectral normalisation. So we cannot say whether the normalisation is what makes the loop settle.

All three models overfit after about two epochs. On the 39 validation problems, answer cross-entropy was lowest at 100–200 answer-training steps, then rose:

Answer-training step100200300400
K=0K=0 control0.6180.6180.7490.726
K=4K=40.6010.6080.6190.694
K=4K=4 with fix0.6140.6150.6160.688

We evaluated the final checkpoint of every run, as fixed in advance. An earlier checkpoint might score slightly higher for all three runs, but 39 validation problems cannot tell the runs apart, and nothing here suggests it would favour the latents.

6. What [1] predicted, and what we found

Prediction in [1]Finding
K=4K=4 latents trained with Sinkhorn alignment retain most of chain-of-thought accuracyNot supported. 19.6% at best, against 86% for full reasoning and 19.8% for fine-tuning without latents.
Multi-token transport preserves the reasoning’s structure, unlike mean poolingNot supported. At the default ε\varepsilon the loss’s own optimum is close to the mean, and the trained latents were even closer: effectively mean pooling.
The frozen-head constraint prevents dimensional collapseNot supported. It was on throughout, and the latents still collapsed (cosine 0.999, effective rank about 4 of 16).
Spectral normalisation makes the loop contract to a fixed pointUntested. The loop does settle (under 1% change by iteration 4), but we neither measured ρ(J)\rho(J) nor ablated the normalisation.
Latent reasoning cuts generated tokens by over 90%True but empty. The latent models generate 5.4 tokens against 2,174, but only because they answer without reasoning, at the accuracy of answering without reasoning.

7. Limitations

  • One model, one benchmark. DeepSeek-R1-Distill-Qwen-1.5B on GSM8K only. We did not run MATH-500, AIME or GPQA for the latent models.
  • One recipe at one scale. 2,136 training traces from the first 2,500 training problems, about 3 epochs, K=4K=4, R=4R=4, LoRA rank 16, one seed. A larger model, all 7,473 training problems, or other KK and RR might behave differently. Given a null result with a diagnosed cause, we chose not to spend compute on sweeps.
  • No curriculum. Continuous-thought methods that work [5] replace reasoning steps with latents gradually over several stages. We tested the one-shot alignment proposed in [1], not a curriculum.
  • One sample per problem. Each accuracy has a confidence interval of about ±2 points. The paired tests account for this.
  • Partial reasoning baseline. Full reasoning is measured on the first 448 of 1,319 test problems. The rest is still to run.
  • No ablations of the components. We did not train without spectral normalisation, without the frozen-head loss, or with plain mean pooling in place of transport. With no effect to explain, ablations would not have been informative.

8. Conclusion

On a 1.5B reasoning model, four Sinkhorn-aligned latent tokens refined by a LOTUS loop add nothing to GSM8K accuracy beyond what the same fine-tuning achieves without them. The latents collapse into one and carry mostly the question.

Fixing the loss’s blurring partly separates them but changes nothing measurable. In both runs the latent pathway could not reach even the collapsed solution of its own alignment loss.

The proposal’s central assumption is that a few latent positions computed from the question can stand in for states that the model reaches only by reasoning. That assumption fails here. The method as described in [1] does not work at this scale, and we do not recommend building on it without a training signal that gives the latents intermediate steps to learn, such as a curriculum.

9. Reproduction

The code is at github.com/vishal-dehurdle/latent-harness (release v0.1.0). Every number in this paper is written by one script, scripts/analysis/paper_numbers.py, from the saved run outputs. The same directory holds the diagnostics (collapse_check.py), the loss simulations (ot_gap.py, ot_attractor.py) and the paired comparison (compare_runs.py).

The commands, as recorded in each run’s manifest:

# Reasoning traces on 2,500 GSM8K training problems (MLX backend)
python scripts/eval_baseline.py --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
    --benchmarks gsm8k_train --conditions think --max-tokens 4096 --batch-size 16 --limit 2500 \
    --out-dir outputs/traces/DeepSeek-R1-Distill-Qwen-1.5B

# Alignment targets: pooled final hidden states of each kept trace
python scripts/build_latent_data.py --traces outputs/traces/DeepSeek-R1-Distill-Qwen-1.5B/gsm8k_train__think.jsonl \
    --out-dir outputs/latent_data/gsm8k_train

# K=4, the K=0 control, and the collapse fix
python scripts/train_latent.py --data-dir outputs/latent_data/gsm8k_train --out-dir outputs/latent/k4 \
    --num-latents 4 --num-iterations 4 --stage1-steps 100 --stage2-steps 400
python scripts/train_latent.py --data-dir outputs/latent_data/gsm8k_train --out-dir outputs/latent/k0 \
    --num-latents 0 --stage1-steps 0 --ot-weight 0 --frozen-head-weight 0 --stage2-steps 400
python scripts/train_latent.py --data-dir outputs/latent_data/gsm8k_train --out-dir outputs/latent/k4_std \
    --num-latents 4 --num-iterations 4 --stage1-steps 100 --stage2-steps 400 --ot-standardize --ot-epsilon 0.03

# GSM8K test: each latent model, the untrained direct floor, and full reasoning
python scripts/eval_baseline.py --backend latent --model outputs/latent/k4/final --benchmarks gsm8k --batch-size 16
python scripts/eval_baseline.py --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --benchmarks gsm8k \
    --conditions direct --max-tokens 16384 --batch-size 16
python scripts/eval_baseline.py --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --benchmarks gsm8k \
    --conditions think --max-tokens 16384 --batch-size 8

# Diagnostics (Section 4), loss simulations, and every number in this paper
python scripts/analysis/collapse_check.py outputs/latent/k4/final outputs/latent_data/gsm8k_train
python scripts/analysis/ot_gap.py outputs/latent_data/gsm8k_train
python scripts/analysis/ot_attractor.py outputs/latent_data/gsm8k_train
python scripts/analysis/paper_numbers.py

The model revision is ad9f0ae0864d7fbcd1cd905e3c6c5b069cc8b562. Each run directory holds its manifest: the exact command, git state, package versions and host. Interrupted training resumes from its last checkpoint and reproduces the same losses.


References

[1]  Verma, V. (2026). "Sinkhorn Alignments: Information Bottleneck Bounds on Span-Level Latent Reasoning." Research Pieces.
[2]  DeepSeek-AI. (2025). "DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning." arXiv:2501.12948.
[3]  Cobbe, K., et al. (2021). "Training Verifiers to Solve Math Word Problems." arXiv:2110.14168.
[4]  Goyal, S., et al. (2024). "Think before you speak: Training Language Models With Pause Tokens." ICLR 2024.
[5]  Hao, S., et al. (2024). "Training Large Language Models to Reason in a Continuous Latent Space." arXiv:2412.06769.
[6]  Hu, E. J., et al. (2022). "LoRA: Low-Rank Adaptation of Large Language Models." ICLR 2022.
[7]  Miyato, T., et al. (2018). "Spectral Normalization for Generative Adversarial Networks." ICLR 2018.
[8]  Cuturi, M. (2013). "Sinkhorn Distances: Lightspeed Computation of Optimal Transport." NeurIPS 2013.
[9]  McNemar, Q. (1947). "Note on the sampling error of the difference between correlated proportions or percentages." Psychometrika 12(2).
[10] Roy, O., & Vetterli, M. (2007). "The Effective Rank: A Measure of Effective Dimensionality." EUSIPCO 2007.
[11] Sun, M., Chen, X., Kolter, J. Z., & Liu, Z. (2024). "Massive Activations in Large Language Models." arXiv:2402.17762.
latent-reasoningsinkhorn-alignmentoptimal-transportdeepseek-r1lotuspause-tokensnegative-resultgsm8krepresentation-collapsetest-time-compute
In this article