P2N: Reusing Deep Representations for Greater Effective Depth

Yiqin Hao; Boyi Zeng; Zitong Wang; He Li; Shixiang Song; Feichen Song; Yifan Liu; Ran Tao

cs.CLOct 18, 2026

P2N: Reusing Deep Representations for Greater Effective Depth

Authors: Yiqin Hao, Boyi Zeng, Zitong Wang, He Li, Shixiang Song, Feichen Song, Yifan Liu, Ran Tao,

Organizations: LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai AI Laboratory

Abstract

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.

Zero-shot
Scale	Method	LAMBADA		PIQA	Wino Grande	SciQ	Hella Swag	ARC		RACE	Avg.
Scale	Method	OpenAI	Standard					Easy	Challenge
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\times$
4	16.983	16.991	235.5	235.4	$1.000\times$
16	17.079	17.143	936.8	933.3	$1.004\times$
32	22.298	22.352	1435.1	1431.7	$1.002\times$

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,

\Delta_{k}=\lVert H^{(k)}-H^{(k-1)}\rVert_{\mathrm{F}}/\sqrt{Td}

. (c) RMSE between Jacobi and sequential core states,

d_{k}=\lVert H^{(k)}-H_{\mathrm{seq}}\rVert_{\mathrm{F}}/\sqrt{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 $\downarrow$	Loss gap	Method	Val. loss $\downarrow$	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 $\downarrow$
Vanilla	0	$1.00\times$	2.0859
Middle $1/3$ ( P2N )	7	$1.83\times$	2.0511
Middle ${\approx}1/2$	11	$2.31\times$	2.0453
Full	21	$3.50\times$	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,\ldots,5

evaluation-time Jacobi iterations. (b) RMSE between consecutive core states,

\Delta_{k}

, over

k=1,\ldots,32

iterations. Random

K\in\{2,3\}

yields stable loss and state differences approaching zero, whereas fixed-

K

models deteriorate beyond their training iteration count.

Core placement	Val. loss $\downarrow$	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	$d_{\mathrm{model}}$	$d_{\mathrm{ff}}$	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	$d_{\mathrm{model}}$	$d_{\mathrm{ff}}^{\mathrm{expert}}$	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.
Scale	Method	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
Scale	Method	EM (%) $\uparrow$	Gain (pp.)	EM (%) $\uparrow$	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 $\times$
4	70.537	70.464	56.7	56.8	0.999 $\times$
16	71.038	71.013	225.2	225.3	1.000 $\times$
32	70.905	70.923	451.3	451.2	1.000 $\times$

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,\ldots,5

, where

K

counts updates after the initial warm pass. After the initial reduction, evaluation loss changes only slightly as further iterations are applied.

Efficient Language Model Inference

Efficient language model inference reduces the compute, memory, latency, and energy needed to generate useful outputs.

See all 321 papers

RecurTrace: Adaptive Latent Reasoning with Loop-Time Memory
Sep 3, 2026Yuxiang Wang, Kunyu Feng, Yingda Shen +3Memory-Augmented Language Models Recurrent Transformers
The Chinese University of Hong Kong, Shenzhen · The Chinese University of Hong Kong · Tianjin University +2
Depth-Attention: Cross-Layer Value Mixing for Language Models
Jun 3, 2026Boyi Zeng, Yiqin Hao, Zitong Wang +7Efficient Transformer Inference Cross-Layer Attention
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute +2
Trading Depth for Time in Recurrent Transformers
Sep 18, 2026Zeyi Huang, Xuehai He, Yelong Shen +1Transformer Recurrent Transformers
Microsoft · University of Wisconsin–Madison

Sep 3, 2026J/K move · Enter open · S save

RecurTrace: Adaptive Latent Reasoning with Loop-Time Memory

Yuxiang Wang, Kunyu Feng, Yingda Shen, Haoning Xu+2

The Chinese University of Hong Kong, Shenzhen · The Chinese University of Hong Kong · Tianjin University+2

Memory-Augmented Language Models Recurrent Transformers LLM Inference Acceleration Adaptive Computation

Open paper AskPDF arXiv

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.

Abstract

Efficient Language Model Inference

Explore similar work

RecurTrace: Adaptive Latent Reasoning with Loop-Time Memory

Depth-Attention: Cross-Layer Value Mixing for Language Models

Trading Depth for Time in Recurrent Transformers