Mathematical problem solving and autonomous agents are increasingly important applications of language models, demanding stronger reasoning capabilities. Computational depth is a key resource for such reasoning, but adding model layers increases both training and inference costs, while test-time scaling incurs substantial additional inference computation. We introduce P2N (Previous to Next), which reuses the previous token's deep representations, already available during generation, to increase effective computational depth with negligible decoding overhead. Specifically, P2N adds the previous token's deep hidden state to the current token's shallow representation, allowing subsequent layers to build on previously completed deep computation. This recurrence extends effective computational depth across tokens, while introducing only a single vector addition per decoding step. We train the recurrence using parallel Jacobi updates restricted to a middle-layer core. Experiments with Qwen-style models show that, at 1.5B parameters, P2N improves GSM8K accuracy by over 10 percentage points under both zero-shot and few-shot prompting, with nearly unchanged decoding latency. The benefits also extend to language modeling and average performance on general downstream benchmarks. on general downstream benchmarks.
Figures & tables
Figure 1: Mathematical reasoning at the 1B and 1.5B scales. At 1.5B, P2N improves GSM8K accuracy over Vanilla by 10.9 and 11.1 percentage points under zero-shot and eight-shot prompting, respectively, corresponding to relative gains of 52.2% and 65.8%.
Figure 2: Overview of P2N . The Transformer is divided into prefix, core, and suffix layers. Autoregressive decoding (right). At each token, the previous token’s core output is added to the current token’s prefix representation before entering the core layers. Since the previous token’s state is already available, this cross-token feedback is naturally compatible with autoregressive decoding; however, directly applying it during teacher-forced training would introduce a sequential dependency across token positions. Parallel training (left). The prefix processes the full sequence once, and an initial warm pass through the core produces the starting states. A small number of Jacobi updates then refines these states in parallel: in each iteration, the preceding iteration’s core outputs are shifted one position to the right and added to the fixed prefix representations to form the new inputs. Finally, the suffix layers are applied once to the final states to compute the language-modeling loss.
Figure 3: Scaling behavior. (a) Final Pile validation loss across parameter scales (lower is better). (b) Pile validation loss over training tokens at the 1.5B scale (lower is better). (c) Mean zero-shot accuracy across parameter scales (higher is better). Overall, P2N exhibits consistently favorable behavior across all three scaling views.
Figure 4: Language modeling at the 1B and 1.5B scales. Perplexity of Vanilla and P2N dense models on (a) Pile validation, (b) WikiText-2, (c) LAMBADA OpenAI, and (d) LAMBADA Standard; lower is better. WikiText-2 uses word-level perplexity, and LAMBADA uses final-word perplexity. P2N uses three Jacobi updates at evaluation time. Numbers above the bars denote perplexity values. P2N achieves lower perplexity at both scales on every benchmark, indicating the effectiveness of this cross-token feedback method.
Scale
Method
LAMBADA
PIQA
Wino Grande
SciQ
Hella Swag
ARC
RACE
Avg.
OpenAI
Standard
Easy
Challenge
Zero-shot
1B
Vanilla
61.73
49.91
71.06
59.91
88.70
41.62
60.02
27.39
33.30
54.85
1B
P2N
64.22
57.13
72.31
58.64
89.50
43.45
61.45
27.05
34.16
56.43
1.5B
Vanilla
66.23
57.36
74.16
59.91
88.60
45.05
64.14
27.99
35.50
57.66
1.5B
P2N
67.92
61.46
73.18
62.27
90.70
46.74
64.90
29.86
37.03
59.34
Table 1: Downstream performance at the 1B and 1.5B scales. Zero-shot and five-shot raw accuracy (%) of dense models. Avg. is the unweighted mean over the nine task configurations. Bold indicates the higher score within each matched scale and shot setting. P2N consistently improves average accuracy at both scales in both settings.
Figure 5: Instruction-following ability. Category-wise MT-Bench scores for 1B (left) and 1.5B (right) dense models after Alpaca fine-tuning, evaluated by GPT-5.5. Each category score averages both conversation turns; legends report the overall mean across all 80 questions and both turns. P2N improves the overall score at both scales, with gains across all eight categories at 1.5B. Both panels use the same radial scale.
Batch
TPOT (ms)
Throughput (tok/s)
Rel. latency
size
Vanilla
P2N
Vanilla
P2N
P2N / Vanilla
1
16.442
16.494
60.8
60.6
1.003×
4
16.983
16.991
235.5
235.4
1.000×
16
17.079
17.143
936.8
933.3
1.004×
32
22.298
22.352
1435.1
1431.7
1.002×
Table 2: Autoregressive decoding efficiency of the 1.5B dense model on a single GPU. TPOT denotes time per output token, and throughput counts generated tokens per second across the batch. Results are averaged over five paired runs generating 2,048 tokens per sequence. Rel. latency is the ratio of P2N TPOT to Vanilla TPOT. P2N maintains near-vanilla latency across batch sizes.
Figure 6: Convergence diagnostics of Jacobi iteration on the 1.5B P2N model. (a) Validation loss on 4096 eval samples under different numbers of Jacobi iterations. (b) RMSE between consecutive core states, Δk=∥H(k)−H(k−1)∥F/Td . (c) RMSE between Jacobi and sequential core states, dk=∥H(k)−Hseq∥F/Td . Validation loss quickly stabilizes, while both state errors continue decreasing toward zero as additional iterations are applied.
(a) Residual-scale treatment
(b) Learnable mixing
Method
Val. loss ↓
Loss gap
Method
Val. loss ↓
Loss gap
Direct sum ( P2N )
2.0511
0.0000
Direct sum ( P2N )
2.0511
0.0000
RMS matching
2.0543
+0.0032
Channel-wise gate
2.0491
−0.0020
Compute-only
2.0665
+0.0154
Scalar gate
2.0494
−0.0017
Table 3: Feedback formulation. Validation loss on 600M model trained for 120B tokens to ablate residual-scale treatments (a) and learnable mixing strategies (b). Loss gaps are measured relative to direct sum; negative values indicate improvement. Scale treatments increase loss, while learnable gates offer only small gains.
Method
Core layers
Relative compute
Val. loss ↓
Vanilla
0
1.00×
2.0859
Middle 1/3 ( P2N )
7
1.83×
2.0511
Middle ≈1/2
11
2.31×
2.0453
Full
21
3.50×
2.1046
Table 4: Recurrent core depth. Validation loss and estimated training compute relative to Vanilla on 600M model trained for 120B tokens. Compute ratios use the layer-based approximation 1+2.5C/21 , where C is the number of recurrent layers. The middle-third core provides most of the gain from the 11-layer core at a lower estimated compute cost.
Figure 7: Effect of sampling the number of Jacobi iterations. (a) Validation loss over k=0,…,5 evaluation-time Jacobi iterations. (b) RMSE between consecutive core states, Δk , over k=1,…,32 iterations. Random K∈{2,3} yields stable loss and state differences approaching zero, whereas fixed- K models deteriorate beyond their training iteration count.
Core placement
Val. loss ↓
Loss gap
Vanilla (no recurrence)
2.0859
0.0000
First 1/3
2.0792
−0.0067
Middle 1/3 ( P2N )
2.0511
−0.0348
Last 1/3
2.0969
+0.0110
Table 5: Recurrent core placement. Validation loss and loss gaps relative to Vanilla on 600M model trained for 120B tokens to ablate different seven-layer core placements. Among recurrent configurations with the same compute budget, the middle-third placement achieves the lowest loss.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Model
Parameters
Layers
dmodel
dff
Q heads
KV heads
P2N core
Prefix/Core/Suffix
150M
149,541,120
12
768
3,328
12
4
L4–L7
4/4/4
300M
303,203,200
15
1,024
4,608
16
4
L5–L9
5/5/5
600M
604,627,328
21
1,280
5,632
20
5
L7–L13
7/7/7
1B
1,011,626,496
24
1,536
7,168
24
6
L8–L15
8/8/8
1.5B
1,496,092,288
27
1,792
8,192
28
7
L9–L17
9/9/9
Appendix
Table 6: Dense model configurations. Vanilla and P2N use identical physical architectures and parameter counts. The P2N core is reported using zero-based layer indices.
Model
Active params.
Total params.
Layers
dmodel
dffexpert
Q/KV heads
P2N core
Prefix/Core/Suffix
A0.15B
150,130,944
794,218,752
12
768
416
12/4
L4–L7
4/4/4
A0.3B
304,186,240
1,790,542,720
15
1,024
576
16/4
L5–L9
5/5/5
A0.6B
606,347,648
3,785,499,008
21
1,280
704
20/5
L7–L13
7/7/7
A1B
1,013,985,792
6,563,049,984
24
1,536
896
24/6
L8–L15
8/8/8
Appendix
Table 7: MoE model configurations. “Active parameters” denote the parameters activated for each token rather than the total number of model parameters. All models use 64 routed experts with Top-8 routing.
Figure 8: Scaling behavior of MoE models. (a) Final Pile validation loss across active parameter scales, from A0.15B to A1B. (b) Training loss over training tokens for the matched A1B models, with each point averaged over 500 consecutive optimizer steps. (c) Mean raw accuracy over all nine downstream task configurations across active parameter scales. Lower is better in panels (a) and (b), and higher is better in panel (c). P2N achieves lower final validation loss, lower training loss at the reported token budgets, and higher average downstream accuracy.
Figure 9: Language modeling at the A0.6B and A1B scales. Perplexity of Vanilla and P2N MoE models on (a) Pile validation, (b) WikiText, (c) LAMBADA OpenAI, and (d) LAMBADA Standard; lower is better. Numbers above the bars denote perplexity values. P2N achieves lower perplexity at both scales on every benchmark, with the largest relative improvements on the LAMBADA variants.
Scale
Method
LAMBADA
PIQA
Wino Grande
SciQ
Hella Swag
ARC
RACE
Avg.
OpenAI
Standard
Easy
Challenge
A0.6B
Vanilla
64.43
52.61
72.85
59.04
88.60
44.12
63.05
26.79
34.83
56.26
A0.6B
P2N
66.68
57.11
73.18
61.25
90.60
45.08
63.51
29.69
34.93
58.00
A1B
Vanilla
67.77
58.74
75.03
62.59
90.20
48.44
68.06
32.94
35.50
59.92
A1B
P2N
69.84
63.11
75.79
65.35
92.30
49.62
69.70
34.98
35.50
61.80
Appendix
Table 8: Downstream performance at the A0.6B and A1B scales. Raw accuracy (%) of Vanilla and P2N MoE models. Avg. denotes the unweighted mean over the nine task configurations. Bold indicates the best score within each matched scale, including ties. P2N improves average accuracy at both scales.
Scale
Method
Zero-shot
Eight-shot
EM (%) ↑
Gain (pp.)
EM (%) ↑
Gain (pp.)
A1B
Vanilla
21.68
–
20.02
–
A1B
P2N
27.98
+6.29
26.38
+6.37
Appendix
Table 9: Mathematical reasoning at the A1B MoE scale. GSM8K test strict final-answer EM (%) under zero-shot and eight-shot prompting after supervised fine-tuning. Gains are measured in percentage points relative to Vanilla and calculated before rounding. P2N improves accuracy in both settings.
Batch
TPOT (ms)
Throughput (tok/s)
Rel. latency
size
Vanilla
P2N
Vanilla
P2N
P2N / Vanilla
1
70.454
70.470
14.2
14.2
1.000 ×
4
70.537
70.464
56.7
56.8
0.999 ×
16
71.038
71.013
225.2
225.3
1.000 ×
32
70.905
70.923
451.3
451.2
1.000 ×
Appendix
Table 10: Autoregressive decoding efficiency of the A1B MoE model on a single GPU. Each entry is the mean of four timed runs generating 2,048 tokens per sequence after one untimed warm-up run. Each run uses a 256-token prompt. Relative latency is normalized to Vanilla TPOT.
Figure 10: Evaluation loss under additional Jacobi iterations at the A1B MoE scale. Loss is evaluated on 4,096 examples for K=0,…,5 , where K counts updates after the initial warm pass. After the initial reduction, evaluation loss changes only slightly as further iterations are applied.
Repeating a small block of middle layers increases a language model's effective inference depth without adding parameters or generating extra tokens, and recent work shows that this latent recurrence improves reasoning. However, two design choices limit these gains. Each iteration sees only the previous output and cannot directly access earlier computations. Moreover, a fixed loop count wastes depth on easy inputs while leaving hard ones with too little computation. We introduce RecurTrace, which addresses both limitations using the loop's own trajectory. Specifically, Loop Memory Attention lets each looped layer attend to its own states from previous iterations along the loop-time axis, so the model can revisit earlier computations instead of relying on the latest state alone. A halting head then reads the loop state and predicts whether to continue, with supervision from an oracle that identifies when additional depth still reduces loss. In a controlled MathQA comparison on the same looped backbone, RecurTrace achieves 56.9% accuracy with an average of 2.0 loops, exceeding the best fixed loop depth by 2.2 points at matched compute. By comparison, ACT and PonderNet collapse to one loop, and CALM reaches only 54.1% with 5.6 loops, while the stronger LoopUS-Conf and TaH-Mismatch baselines reach 55.3% at 3.2 loops and 55.7% at 2.1 loops. Finally, RecurTrace improves generation accuracy over same-budget fine-tuned baselines at 0.6B, 1.7B, 4B, and 8B, with the gain growing with model size from 0.6 to 3.4 points.
Yuxiang Wang, Kunyu Feng, Yingda Shen +3
The Chinese University of Hong Kong, Shenzhen · The Chinese University of Hong Kong · Tianjin University +2
Self-attention selects information freely across the sequence, but across depth, Transformers merely add each layer's output to the residual stream, so later layers cannot selectively reuse earlier-layer representations. Recent cross-layer methods improve this flow but operate on hidden states outside attention, adding state beyond the key-value cache at inference--a cost that becomes increasingly salient as modern LLMs compress the cache with grouped-query and multi-head latent attention. We introduce Depth-Attention, which performs this selection inside the attention module itself: before a layer attends over the sequence, its query attends over the keys of earlier layers at the same token position and mixes their values into the value that self-attention then reads. Because Depth-Attention reuses the standard attention queries, keys, and value-cache slots, storing depth-mixed values in place of the original values, it adds no parameters and introduces no persistent inference state beyond the standard key-value cache--the same cache size as a vanilla decoder and less than hidden-state-based cross-layer methods. On Qwen3-style decoders at 1.5B and 3B parameters, Depth-Attention attains the lowest perplexity and the highest average downstream accuracy, improving over the vanilla Transformer by up to 2.3 accuracy points and surpassing strong cross-layer baselines in perplexity and average accuracy, while adding under 0.01% extra arithmetic FLOPs and no additional persistent inference state. The gains hold from 360M to 3B parameters and extend to looped Transformers.
Boyi Zeng, Yiqin Hao, Zitong Wang +7
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute +2
Recurrent Transformers increase computational depth through temporal recurrence, feeding each token's high-level hidden state into the computation of the next. This raises a natural question: is additional computation better spent on more temporal steps or greater physical depth? We investigate this question using Latent Recurrent Transformers (LRT), which retain one backbone forward pass per vocabulary token during decoding and provide a controlled setting for comparing these two ways of adding computation. We insert a latent thought token between consecutive vocabulary tokens, adding a hidden-state refinement step through the shared backbone before predicting the next token. On 16- and 20-layer NanoChat MoE backbones, we compare this model with double-depth vanilla Transformers and LRT, two-loop Transformers, and PonderLM-2 with one thought token. All configurations execute 2L blocks per vocabulary token during decoding. Under matched training FLOPs, LRT with one thought token outperforms PonderLM-2 and overtakes the two-loop Transformer as compute increases. Its gap to vanilla double-depth narrows from 0.014 to 0.001 bits per byte, despite approximately 46% fewer parameters and approximately half as many unique training tokens. At the training budget used for our main comparison, one thought token recovers an average of 89% of the gains from doubling LRT depth across the two backbone sizes, with approximately 48% fewer parameters. These results highlight temporal thinking as a parameter-efficient way to extend computation in recurrent Transformers.