Write Back the $Δ$: Revisiting the Same Tokens with Fresh Representations

Wencheng Ye; Anning Hu; Xiangdong Zhang; Tianyi Wang; Yikang Li; Hengyu Jin; Bing Li; Junchi Yan

cs.LGSep 26, 2026

Write Back the $Δ$ : Revisiting the Same Tokens with Fresh Representations

Authors: Wencheng Ye, Anning Hu, Xiangdong Zhang, Tianyi Wang, Yikang Li, Hengyu Jin, Bing Li, Junchi Yan

Organizations: Tongji University · Shanghai Jiao Tong University · Shanghai Innovation Institute

Abstract

Transformers process information strictly forward through depth, preventing deeper computation from revisiting and refining earlier representations. To augment the standard forward pass, existing approaches either re-execute depth, incurring additional computation, or modify the residual stream using predefined directions, limiting their instance-level adaptation. Recently, inference-time feedback offers a direct mechanism for recycling endogenously produced computation by writing deeper residual states back to earlier layers, yet what should be fed back remains unclear. We argue that the depth increment Delta, capturing newly accumulated computation between two layers, provides a more effective, composable, and scalable feedback signal than the full state. Building on this observation, we introduce ReFlux, a learnable feedback graph that dynamically selects and composes increment-carrying routes. ReFlux supports synchronous feedback to the same token and streaming feedback to subsequent tokens. Extensive experiments across various models, corpora, and benchmarks show that synchronous ReFlux consistently reduces perplexity across ten language-modeling corpora, and improves accuracy by 2.1-2.3 points, with gains reaching 4.7 points on multi-hop reasoning. Streaming ReFlux further retains most of these gains while preserving the base model's 1x theoretical backbone FLOPs. These results establish ReFlux as an efficient paradigm for unlocking the latent computational potential of LLMs, allowing them to revisit the same tokens with fresh representations. Code implementation can be found at https://github.com/gooogleshanghai/reflux.

Figures & tables

Figure 1: Revisiting earlier representations with new computation. Looped Transformers re-execute depth, while Recirculation returns full states to earlier layers. ReFlux instead isolates the newly accumulated computation and selectively routes depth increments $\Delta$ through a learned feedback graph, enabling streaming feedback without increasing the backbone’s theoretical FLOPs.

Figure 2: Cross-depth geometry of Gemma3-1B residual states. (a) Pairwise cosine similarity between $z_{d}$ and $z_{s}$ ; stars mark top-three positive source-target pairs. (b) Cosine similarity and norm-relative displacement versus layer separation $k$ . Adjacent states are similar, while successive layerwise rewrites accumulate into a larger change over longer depth intervals.

	Gemma3-1B		Gemma3-4B		Gemma3-12B		Qwen3-4B
Payload	PPL $\downarrow$	$\Delta\log$ PPL ( $\times 10^{-2}$ ) $\uparrow$	PPL $\downarrow$	$\Delta\log$ PPL ( $\times 10^{-2}$ ) $\uparrow$	PPL $\downarrow$	$\Delta\log$ PPL ( $\times 10^{-2}$ ) $\uparrow$	PPL $\downarrow$	$\Delta\log$ PPL ( $\times 10^{-2}$ ) $\uparrow$
None	30.772 ↓0.00	0.00	23.693 ↓0.00	0.00	42.973 ↓0.00	0.00	25.562 ↓0.00	0.00
Full $z_{s}$	29.699 ↓3.49	+3.55	22.790 ↓3.81	+3.89	41.462 ↓3.52	+3.58	24.626 ↓3.66	+3.73
Base $z_{d}$	31.110 ↑1.10	-1.09	23.899 ↑0.87	-0.87	43.451 ↑1.11	-1.11	25.599 ↑0.14	-0.14
$\Delta_{d\to s}$	29.210 ${}_{\color[rgb]{0,0.5117,0.4688}\downarrow\textbf{5.08}}$	+5.21	21.856 ${}_{\color[rgb]{0,0.5117,0.4688}\downarrow\textbf{7.75}}$	+8.07	40.231 ${}_{\color[rgb]{0,0.5117,0.4688}\downarrow\textbf{6.38}}$	+6.59	24.030 ${}_{\color[rgb]{0,0.5117,0.4688}\downarrow\textbf{5.99}}$	+6.18

Table 1: Matched-payload decomposition under architecture-specific stable feedback protocols on C4, with 64 windows of 512 tokens for each model. We use the best-performing source–target configuration. For Gemma3-1B, we use

s{=}11\to d{=}4

\alpha{=}0.25

\beta{=}0.9

; for Gemma3-4B,

s{=}18\to d{=}9

\alpha{=}0.15

\beta{=}0.9

; for Gemma3-12B,

s{=}35\to d{=}16

\alpha{=}0.15

\beta{=}1.0

; for Qwen3-4B,

s{=}26\to d{=}13

\alpha{=}0.2

\beta{=}1.0

Figure 3: Scaling behavior of feedback across three axes on Gemma3-4B.

Figure 4: Overview of ReFlux Architecture and Execution Schedules.

Corpus	Gemma3-1B					Gemma3-4B					Gemma3-12B					Qwen3-4B					Qwen3-8B
	Base	Orig.	TF-Loop	Stream	Sync.	Base	Orig.	TF-Loop	Stream	Sync.	Base	Orig.	TF-Loop	Stream	Sync.	Base	Orig.	TF-Loop	Stream	Sync.	Base	Orig.	TF-Loop	Stream	Sync.
arXiv	26.7	23.7	26.5	22.9	22.5	22.3	19.2	22.0	18.4	18.0	36.9	26.3	35.9	24.7	23.7	18.3	16.5	18.2	15.7	15.4	14.5	12.8	14.3	12.3	12.1
BigPatent	11.8	10.7	12.0	10.6	10.4	9.7	9.0	9.6	8.7	8.5	20.3	15.4	20.0	14.5	14.0	7.8	7.3	7.9	7.1	7.0	6.6	6.2	6.6	6.0	5.9
BillSum	5.2	5.1	5.1	4.9	5.0	4.0	3.8	4.0	3.7	3.6	4.8	4.3	4.9	4.1	4.0	4.7	4.5	4.6	4.6	4.3	3.9	3.7	4.0	3.6	3.5
BookSum/Book	30.6	25.6	30.8	25.0	24.4	29.7	25.0	29.1	23.9	23.3	78.2	53.8	75.3	50.0	47.2	17.9	16.3	17.7	15.8	15.5	12.4	11.4	12.2	10.9	10.8
C4/WebTextLike	25.4	24.1	25.2	23.5	23.3	23.8	22.0	23.5	21.2	20.9	40.6	37.8	40.1	35.4	34.8	21.6	19.9	21.8	19.0	18.7	16.3	15.6	16.2	14.7	14.5

Table 2: Perplexity on the ten language-modeling corpora (1024-token chunks). Lower is better. We highlight the best and second-best results.

Benchmark

Gemma3-1B

Gemma3-4B

Qwen3-8B

Base

Orig.

TF-Loop

Stream

Synchronous

Base

Orig.

TF-Loop

Stream

Synchronous

Base

Orig.

TF-Loop

Stream

Synchronous

ARC-Easy

73.2

73.4 ↑0.2

74.5 ↑1.3

74.8 ↑1.6

74.7 ↑1.5

82.9

83.0 ↑0.1

83.7 ↑0.8

84.6 ↑1.7

84.9 ↑2.0

80.9

81.2 ↑0.3

81.6 ↑0.7

82.1 ↑1.2

82.5 ↑1.6

ARC-Challenge

35.7

35.8 ↑0.1

36.5 ↑0.8

37.4 ↑1.7

37.8 ↑2.1

56.2

56.4 ↑0.2

57.5 ↑1.3

58.0 ↑1.8

58.4 ↑2.2

56.7

56.7 ↑0.0

57.4 ↑0.7

58.0 ↑1.3

58.5 ↑1.8

MMLU

26.2

26.3 ↑0.1

27.8 ↑1.6

27.5 ↑1.3

27.4 ↑1.2

59.5

59.9 ↑0.4

60.7 ↑1.2

60.6 ↑1.1

60.8 ↑1.3

72.9

73.1 ↑0.2

74.3 ↑1.4

74.1 ↑1.2

73.9 ↑1.0

HellaSwag

62.1

62.0 ↓0.1

62.6 ↑0.5

62.5 ↑0.4

77.2

77.0 ↓0.2

77.6 ↑0.4

77.9 ↑0.7

74.9

74.8 ↓0.1

75.0 ↑0.1

75.3 ↑0.4

75.5 ↑0.6

Winogrande

58.3

58.2 ↓0.1

58.4 ↑0.1

58.6 ↑0.3

58.7 ↑0.4

64.4

64.4 ↑0.0

64.6 ↑0.2

64.9 ↑0.5

64.8 ↑0.4

67.7

67.8 ↑0.1

67.4 ↓0.3

68.0 ↑0.3

68.2 ↑0.5

Table 3: Accuracy on eight reasoning benchmarks. Each model group reports the base model, Original Recirculation, Training-Free Looped, ReFlux-streaming, and ReFlux-synchronous. Higher is better. We highlight the best and second-best results.

Figure 5: Deep computation as a predictive signal. In “The trophy does not fit in the suitcase because it is too big,” shallow layers may encode only coarse semantics when processing “it” , while previous deeper computation can establish the relation between the trophy and suitcase . Instead of revisiting the same token to apply this refinement, ReFlux-streaming carries the $\Delta$ forward and makes it available to subsequent tokens, allowing them to benefit from computation already performed.

	Gemma3-1B					Qwen3-4B
Method	Runtime( $\times$ )	FLOPs( $\times$ )	Memory(GB)	Latency(ms/tok)	PPL	Runtime( $\times$ )	FLOPs( $\times$ )	Memory(GB)	Latency(ms/tok)	PPL
Base	1.00	1.00	1.88	15.05	25.4	1.00	1.00	7.56	16.05	21.6
Orig.	1.88	2.00	1.90	28.26	24.1	1.94	2.00	7.58	31.14	19.9
TF-Loop	1.34	1.31	1.88	20.16	25.2	1.21	1.22	7.56	19.42	21.8
Streaming	1.02	1.00	1.88	15.34	23.5	1.08	1.00	7.57	17.35	19.0
Synchronous	1.91	2.00	1.89	28.71	23.3	1.97	2.00	7.64	31.59	18.7

Table 4: Decoding efficiency and C4 language-modeling quality. Efficiency metrics are color-coded.

Figure 6: Routing behavior summary. (a) Integer active-edge count for routing block $b{+}1$ against the base-model NLL of the preceding block $b$ ; the line connects binned means. (b) Source-target heatmap of aggregated activation frequencies, with the most frequently used pairs marked by stars. (c) Comparison of same-family and cross-family route overlap.

Figure 7: Case study (App. G ). (a) The question’s evidence chain, (b) how route support accumulates (we set the block size to 1 token for visualization), (c) the final active source-target edges, and (d) teacher-forced token probabilities for the gold answer path and the strongest wrong competitor.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Method	adds effective depth	no backbone retraining	any frozen checkpoint	single traversal	in-place write	endogenous payload	multi-edge feedback
UT / Looped LM / SMELT †	✓	✗	✗	✗	✓	✓	✗
MoR †	✓	✗	✗	$\sim$	✓	✓	✗
MoDA †	$\sim$	✗	✗	✓	✓	✓	$\sim$
T 2 MLR †	✓	✗	✗	✓	✓	✓	✗
TF-Loop	✓	✓	✓	✗	✓	✓	✗
CoLa	✓	✓	✓	✗	✓	✓	$\sim$

Appendix

Table 6: Positioning of ReFlux among depth-feedback and activation-editing methods. ✓ / ✗ : property holds / does not hold;

\sim

: partially.

\dagger

requires pretraining or continued training. The single-traversal column for ReFlux refers to its streaming schedule. Sources: UT ( Dehghani et al., 2018 ) ; Looped LM ( Yang et al., 2024a ) ; SMELT ( Wang et al., 2026a ) ; MoR ( Bae et al., 2026 ) ; MoDA ( Zhu et al., 2026 ) ; T 2 MLR ( Cai et al., 2026 ) ; TF-Loop ( Chen et al., 2026 ) ; CoLa ( Li et al., 2025 ) ; CAA ( Rimsky et al., 2024 ) ; ITI ( Li et al., 2023 ) ; Patchscopes ( Ghandeharioun et al., 2024 ) ; Recirculation ( Mozer et al., 2026 ) .

Figure 8: Cross-depth residual-stream geometry across the eight profiled models. Each model panel shows pairwise cosine similarity, cosine similarity $C(k)$ , and norm-relative displacement $R(k)$ . Stars denote the best tested source-target edge(s) for each model.

Method	adds serial steps	extra training	extra token generation	plug-and-play	bottleneck addressed
CoT	✓	✗	5–50 $\times$	✓	computation
Filler tokens	✓	✓	$N\times$	✗	computation
Looped	✓	✓	0	✗	computation
ReFlux	✗	$\sim$	0	✓	access

Appendix

Table 7: Positioning of CoT, filler-token, looped, and increment-feedback methods. Filler tokens must be trained to be exploited ( Pfau et al., 2024 ) ; the router calibration of ReFlux is a single lightweight pass. The last column indicates which of the two bottlenecks each method primarily addresses. ✓ / ✗ : property holds / does not hold;

\sim

: partially. Sources: CoT ( Wei et al., 2022 ) ; Filler tokens ( Pfau et al., 2024 ) ; Looped ( Dehghani et al., 2018 ; Yang et al., 2024a ) .

	2WikiMultiHopQA		MMLU		GSM8K
Setting	Acc	Tokens	Acc	Tokens	Acc	Tokens
Base	44.5	97	70.2	64	88.6	189
+ CoT	44.7	326	71.4	508	91.8	1002
ReFlux (sync)	49.5	120	71.2	129	90.4	263
+ CoT	49.9	294	71.1	480	92.0	691

Appendix

Table 8: Four-setting comparison on Qwen3-8B. We report accuracy and mean generated tokens per example.

Figure 9: Pareto view of generation efficiency and quality for Gemma3-1B and Qwen3-4B. Streaming stays near the base runtime while retaining most of the PPL gain; Synchronous reaches the highest-quality frontier at roughly

2\times

cost; Original Recirculation is dominated by the delta-based variants; and TF-Loop remains lower in quality.

Schedule	Prefill	Decode	Prefill (ms)	TTFT (ms)	Decode (ms/tok)	Total (s)	Prefill $\times$	Decode $\times$	E2E $\times$
Base	Parallel, 1 pass	Standard	69.4	104.9	32.67	0.85	1.00	1.00	1.00
Strict Stream	Sequential	Chained	35378.8	35412.1	33.03	36.17	510.08	1.01	42.38
Parallel-Prefill Stream	Parallel, 1 pass	Chained	70.1	105.9	33.76	0.87	1.01	1.03	1.02
Synchronous	Parallel, 2 pass	2 pass	134.0	201.0	66.24	1.72	1.93	2.03	2.02
Orig. Recirculation	Parallel, 2 pass	2 pass	134.7	204.8	67.36	1.75	1.94	2.06	2.05
TF-Loop	Parallel, looped window	Looped	82.5	124.1	39.26	1.02	1.19	1.20	1.20

Appendix

Table 9: Prefill/decode decomposition on a supplementary multi-hop generation workload. This table uses Gemma3-4B on 128 HotpotQA and 2WikiMultiHopQA prompts (mean length 1107 tokens), batch size 1, and 24 teacher-forced decode steps. Ratios are relative to Base within this workload.

Panel of Figure 10 — Figure 10: ReFlux execution paths. Synchronous recurrence requires a second backbone traversal, whereas streaming carries the feedback increment across tokens. The two streaming variants differ only in prompt processing and share the same one-pass recurrent decoding mechanism.

Block $w$	C4 PPL $\downarrow$	ARC-C $\uparrow$	$\bar{\kappa}$	Runtime ( $\times$ )
32	21.1	58.1	2.73	1.08
64	21.0	58.2	2.66	1.06
128	21.2	58.0	2.58	1.04
256	21.4	57.8	2.31	1.04
512	21.8	57.4	2.44	1.03

Appendix

Table 10: Sensitivity to routing-block length on Gemma3-4B using ReFlux-streaming. Runtime is normalized to the base model; lower C4 PPL and higher ARC-C accuracy are better.

Threshold $\xi$	C4 PPL $\downarrow$	ARC-C $\uparrow$	$\bar{\kappa}$	Runtime ( $\times$ )
0.4	21.5	57.6	4.61	1.06
0.5	21.1	58.1	3.37	1.06
0.6	21.2	58.0	2.58	1.04
0.7	21.5	57.7	1.74	1.03
0.8	22.0	57.1	0.91	1.03

Appendix

Table 11: Sensitivity to the inference-time routing threshold on Gemma3-4B using ReFlux-streaming. The router is not retrained between settings. Runtime is normalized to the base model.

Task family	$\kappa$	$H$	Same ovlp.	Cross ovlp.	Depth span
C4 / WebTextLike	2.18	1.09	0.64	0.19	0.27
BookSum / PG19	2.87	1.23	0.58	0.33	0.36
MMLU / ARC	2.33	1.14	0.67	0.30	0.24
HellaSwag / Winogrande	1.96	1.06	0.69	0.38	0.22
GSM8K	2.54	1.19	0.60	0.34	0.31
HotpotQA / 2WikiMultiHopQA	3.12	1.34	0.54	0.22	0.42

Appendix

Table 12: Routing summary by task family on Gemma3-1B.

\kappa

is the active-edge count,

H

is route entropy, overlap compares top-

k

soft-route sets, and depth span is the source-minus-target layer distance normalized by model depth.

	Params	Train time	C4 PPL $\downarrow$	ARC-C $\uparrow$
Base	0	–	21.6	53.8
Learned constant	1,260	1.3 h	20.9	54.2
Conditional-vector MLP	5.2M	1.8 h	19.9	54.9
LoRA-r8 ( $q,v$ )	2.9M	2.2 h	19.6	55.1
ReFlux router	2.2M	1.5 h	18.7	55.8

Appendix

Table 13: Budget-matched adaptation baselines on Qwen3-4B. All methods use the same calibration data, optimizer, and update count; the backbone is frozen throughout. ReFlux is evaluated with synchronous execution.

Efficient Transformer Inference

Efficient Transformer Inference studies methods and systems that reduce the latency, memory, and energy needed to run Transformer models.

See all 122 papers

LoopFold: Parallelizing Recurrent Depth in Looped Transformers
Oct 26, 2026Feichen Song, Boyi Zeng, Yiqin Hao +7Recurrent Transformers Efficient Transformer Inference
Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute
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

Oct 26, 2026J/K move · Enter open · S save

LoopFold: Parallelizing Recurrent Depth in Looped Transformers

Feichen Song, Boyi Zeng, Yiqin Hao, Shixiang Song+6

Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute

Recurrent Transformers Efficient Transformer Inference LLM Inference Acceleration

Open paper AskPDF arXiv

Looped Transformers increase computation per token by repeatedly applying a shared backbone, offering a way to improve language modeling and downstream task performance at a fixed backbone parameter count. However, each loop depends on the preceding loop, so increasing recurrent depth also increases decoding latency. Existing approaches include Parallel Loop Transformers, which reduce latency by overlapping loop computation across token positions, but yield rapidly diminishing marginal accuracy gains at larger depths. To address this challenge, we introduce LoopFold, which folds computation at multiple recurrent depths into parallel batches for the same token prefix. A lightweight predictor estimates inputs to later loops so they can start before earlier loops finish. A corrector then refines the outputs of these later loops using newly available results from earlier loops. All components are trained jointly using only the final next-token prediction loss. In experiments with 1.4B models, deeper LoopFold configurations achieve progressively higher average zero-shot accuracy across eight tasks, rising from 53.41% to 56.38%. Vanilla looped models and PLT show limited gains over the same depth range. Compared with a nine-loop vanilla model, LoopFold uses five sequential stages and additional parallel computation to reduce decoding latency by 46%.

Abstract

Efficient Transformer Inference

Explore similar work

LoopFold: Parallelizing Recurrent Depth in Looped Transformers

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

Trading Depth for Time in Recurrent Transformers