Hierarchical vs. Flat Iteration in Shared-Weight Transformers

Sang-Il Han

cs.CLApr 15, 2026

Hierarchical vs. Flat Iteration in Shared-Weight Transformers

Authors: Sang-Il Han

Organizations: Korea University of Technology and Education · Biomonster Co.

Abstract

We present an empirical study of whether hierarchically structured, shared-weight recurrence can match the representational quality of independent-layer stacking in a Transformer-based language model. HRM-LM replaces L independent Transformer layers with a two-speed recurrent pair: a Fast module operating at every step for local refinement, and a Slow module operating every T steps for global compression. This recurrent hierarchy is unrolled for M = N x T steps with shared parameters. The central and most robust finding, supported by a parameter-matched Universal Transformer ablation (UniTF, 1.2B) across five independent runs, is a sharp empirical gap between the two approaches.

Figures & tables

Property	Transformer	HRM-LM	Valid
Stored params	$O(Ld^{2})$	$O(d^{2})$	$\checkmark$
Inference KV cache	$O(Lnd)$	$O(Mnd)$	$\circ$ ( $L>M$ )
Parameter sets	$L$ (indep.)	$1$ (shared)	$\checkmark$
Step parallelism	Pipeline-opt.	Serial	—
Depth scaling	$+$ params	Free ( $N,T,S$ )	$\checkmark$
Gradient flow	$L$ layers	$K$ -step window	—

Table 1: Structural comparison: Transformer vs. HRM-LM.

\checkmark

: always valid.

\circ

: conditional on

L>M

Symbol	Meaning	Transformer	HRM-LM
$d$	model dim.	shared	shared
$L$	layers	$L$ param sets	—
$N$	Slow-module cycles	—	hyperpar.
$T$	L-steps/cycle	—	hyperpar.
$S$	supervision passes	—	hyperpar.
$K$	gradient window	—	hyperpar.

Table 2: Key hyperparameters, derived quantities, and parameter sets.

Update type	KV cache	Inter-token	Expressiveness
		interaction
Full causal attn (ours)	$O(Mnd)$	✓	High
RNN position-wise	$O(d)$	$\times$	Lower
Hybrid ( $k$ attn steps)	$O(knd)$ , $k{<}M$	Partial	Medium

Table 3: Design choices for the Fast-module update.

Dimension	Transformer	HRM-LM
Stored params	$O(Ld^{2})$	$O(d^{2})$
Train activations	$O(LA)$	$O(KA)$
Train FLOPs	$3LC$	$S(M\!+\!2K)C$
Inference KV cache	$O(Lnd)$	$O(Mnd)$

Table 4: Asymptotic complexity comparison.

	Exp	Fixed	Measured
Rep.	1 (equal params)	Stored params $\approx$ 1.23B	val CE
$\ddagger$	2 (equal FLOPs)	FLOPs/iter	val CE (T-L12 partial)
Rep.	3 (long train)	FLOPs/iter, 20k steps	val CE
$\ddagger$	A (LR search)	Best LR for T-L12	val CE (60k, converged)
Rep.	4 (inference)	Architecture	latency, peak memory
Rep.	5 (ablation)	$d$ , base config	val CE vs $K/S/N/T$

Table 5: Comparison axes. Bold : reported.

\ddagger

: reported, baseline incomplete.

Model	Params	val CE	$\Delta$ vs T-L4
HRM $N\!\times\!T\!=\!8$ (N=2,T=4)	1229M	4.239	$-$ 0.323
HRM $N\!\times\!T\!=\!12$ (N=4,T=3)	1229M	4.284	$-$ 0.278
HRM $N\!\times\!T\!=\!4$ (N=2,T=2)	1229M	4.450	$-$ 0.112
Transformer L=2	743M	4.509	$-$ 0.053
Transformer L=4	1280M	4.562	—

Table 6: EqualParam — equal stored parameters (

d=4096

, 10k steps).

\Delta

is relative to Transformer L=4. All models use scheduled LR (warmup + cosine decay) with scaled initialization and K/M-proportional gradient clipping. HRM NT=12 here uses the EqualParam LR schedule (lr

=4\!\times\!10^{-4}

, warmup

=500

); Ablation and UniTF use a lower LR (lr

=1.5\!\times\!10^{-4}

, warmup

=1000

), yielding 4.177—the 0.1-nat difference reflects per-configuration LR sensitivity.

Panel of Figure 1 — Figure 1: EqualParam learning curves (val CE, 10k steps). All three HRM variants outperform both Transformer baselines. HRM $N\!\times\!T\!=\!4$ ( $T\!=\!2$ ), which previously exploded under a higher LR, now converges stably to 4.450 after applying scaled initialization, extended warmup, and K/M-proportional gradient clipping (max_norm $=0.5$ ). The three HRM variants outperform Transformer L=4 by $0.11$ – $0.32$ nats. Initial values (iter 0, $\approx$ 10.9–11.7) are off-chart.

Model	Params	val CE	$\Delta$ vs T-L12
HRM $d\!=\!5120$ , $N\!\times\!T\!=\!9$	1857M	4.165	$-$ 1.231
HRM $d\!=\!4096$ , $N\!\times\!T\!=\!12$	1229M	4.233	$-$ 1.163
HRM $d\!=\!2048$ , $N\!\times\!T\!=\!36$	359M	4.364	$-$ 1.032
Transformer L=12	3427M	5.396 ‡	—

Table 7: EqualFLOPs — equal FLOPs (

\approx

Transformer L=12), 10k steps.

\ddagger

: partial convergence (grad spikes at iter 500/1500, plateau after iter 8500). HRM advantage is real but T-L12 is not fully converged at 10k; LongTrain extends to 20k.

∗ EqualFLOPs result; LongTrain run uses different config (warmup=2000).
Model	Params	val CE (10k)	val CE (20k)	val CE (60k)	$\Delta$ (10k→20k)
HRM $d\!=\!4096$ , $N\!\times\!T\!=\!12$	1229M	4.233	3.901	—	$-$ 0.332
Transformer L=12, $d\!=\!4096$	3427M	5.396 ‡	4.929 ‡	4.251 ‡	$-$ 0.467
HRM $d\!=\!2048$ , $N\!\times\!T\!=\!36$	359M	(4.364 ∗ )	6.113 †	—	$+$ 1.749
Universal TF $M\!=\!12$ , $d\!=\!4096$	474M	6.673	6.614 ⋆	—	$-$ 0.059
T-L12 at 60k (4.250) trails HRM NT=12 at 10k (4.177) by 0.073 nats; confirmed converged ( $\Delta$ CE $<$ 0.001 nats over final 3,500 steps).

Table 8: LongTrain — extended training (20k steps; T-L12 extended to 60k, confirmed converged).

\ddagger

: partial convergence, still slowly improving.

\dagger

: training failure (warmup

\times

2 + K/M

=5.6\%

incompatible).

\star

: spike-triggered descent from long plateau. HRM NT=12 here uses the EqualFLOPs/3 LR schedule (lr

=5\!\times\!10^{-5}

, warmup

=1000

, different from EqualParam and Ablation/B); the 10k value (4.233) lies between EqualParam (4.284, higher LR) and Ablation/B (4.177, lower LR), all consistent with LR sensitivity.

Config	Steps	val CE	Status
T-L12, lr $=5\!\times\!10^{-5}$ (EqualFLOPs)	10k	5.396	[email protected]
T-L12, lr $=5\!\times\!10^{-5}$ (LongTrain)	20k	4.929	slow descent
T-L12, lr $=3\!\times\!10^{-5}$ (TF-Baseline)	60k	4.250	converged
HRM NT=12 (1229M, EqualFLOPs LR)	10k	4.233	converged
HRM NT=12 (1229M, LongTrain)	20k	3.901	converging

Table 9: TF-Baseline — T-L12 extended training with best-found LR (lr

=3\!\times\!10^{-5}

, warmup

=1000

\ddagger

: still slowly improving.

Model	Params	Lat.	Peak	$\times$ T-L4
T L=2	743M	6.6 ms	1670 MB	$0.63\times$
T L=4	1280M	10.5 ms	2698 MB	$1\times$
T L=12	3427M	25.4 ms	6810 MB	$2.42\times$
HRM $M\!=\!4$	1229M	20.5 ms	2656 MB	$1.96\times$
HRM $M\!=\!8$	1229M	31.1 ms	2656 MB	$2.97\times$
HRM $M\!=\!12$	1229M	46.8 ms	2656 MB	$4.47\times$

Table 10: Latency — inference metrics (seq

\!=\!1024

, bfloat16, RTX 6000, no KV cache).

\dagger

d\!=\!2048

, others

d\!=\!4096

. Ratio vs. T-L4.

K ablation (N=4, T=3, S=1)
Config	val CE	$\Delta$	Notes
$K=1$ (K/M $=8\%$ )	4.238	$+$ 0.012
$K=2$ (base, 17%)	4.226	—
$K=4$ (33%)	4.098	$-$ 0.128
$K=8$ (67%)	4.052	$-$ 0.174	best overall
S ablation (N=4, T=3, K=2)

Table 11: Ablation results (base:

N=4

T=3

S=1

K=2

d=4096

, lr

=1.5\!\times\!10^{-4}

, warmup

=1000

, 10k steps).

\Delta

relative to respective base. The base HRM NT=12 val CE (4.177) differs from EqualParam (4.284) due to the lower LR used here; both are the same architecture and data.

Model	Params	val CE @ 10k	Note
HRM $N\!\times\!T\!=\!12$ , $d\!=\!4096$	1229M	4.177	UniTF (10k complete)
T-L4, $d\!=\!4096$	1280M	4.562	EqualParam
UniTF $M\!=\!12$ , $d\!=\!4096$	474M	6.403
UniTF $M\!=\!12$ , $d\!=\!5760$	820M	7.624	3 runs; see App. B
UniTF $M\!=\!12$ , $d\!=\!7296$	1218M	7.578	10k complete; plateau throughout

Table 12: UniTF — parameter-matched UniTF vs HRM at 10k steps. HRM from EqualParam; UniTF

d\!=\!5760

complete (3 runs);

d\!=\!7296

complete.

Panel of Figure 2 — Figure 2: UniTF val CE curves. HRM NT=12 (1229M, green) converges steadily to 4.177 at 10k steps. All UniTF variants plateau at $\approx$ 7.6 regardless of width (820M or 1218M), learning rate, warmup, or initialization. v2b and v2c started from a bad random seed (iter-0 CE $>$ 90); curves shown from iter 500 after recovery. The $\approx$ 3.4-nat gap at matched parameter count (1218M vs 1229M) is consistent with a structural rather than parameter-count advantage for HRM.

HRM NT=12 (1229M)			T-L4 (1280M)
lr	warmup	val CE	lr	warmup	val CE
$1.5\!\times\!10^{-4}$	1000	4.308	$2\!\times\!10^{-4}$	1000	4.543
$1.5\!\times\!10^{-4}$	500	4.485	$2\!\times\!10^{-4}$	500	4.545
$8\!\times\!10^{-5}$	500	4.367 †	$10^{-4}$	1000	4.687
$8\!\times\!10^{-5}$	1000	4.381	$10^{-4}$	500	4.706
$4\!\times\!10^{-5}$	1000	4.920	$3\!\times\!10^{-4}$	1000	4.997

Table 13: HPSearch: budget-matched LR/warmup grid search, 10k steps, seed=42. Best per model in bold . HRM NT=12 best (4.308) outperforms T-L4 best (4.543) by 0.235 nats under matched hyperparameter search.

Panel of Figure 3 — Figure 3: HPSearch val CE curves (10k steps, seed=42). Left : HRM NT=12 best config (lr $=1.5\!\times\!10^{-4}$ , w $=1000$ ) achieves 4.308 (thick teal). Right : T-L4 best config (lr $=2\!\times\!10^{-4}$ , w $=1000$ ) achieves 4.543 (thick blue). At their respective best hyperparameters, HRM leads by 0.235 nats.

Aggregation	HRM NT=12	T-L4	HRM advantage
Best (top-1)	4.308	4.543	$+0.235$
Top-2 mean	4.338	4.544	$+0.206$
Median (of 8)	4.703	4.852	$+0.149$
Worst (of 8)	5.727	5.799	$+0.072$

Table 14: HPSearch budget summary (8 configs each, 10k steps, seed=42). HRM NT=12 outperforms T-L4 at every aggregation level.

Model	seed=42	seed=1337	seed=2024	mean	$\pm$ std
T-L4 (1280M)	4.702	4.698	4.692	4.697	$\pm$ 0.004
HRM NT=12 (1229M)	4.909	4.892	4.889	4.896	$\pm$ 0.009

Table 15: MultiSeed — multi-seed reproducibility (10k steps). Both models show low seed variance; results are reliable across seeds.

Model	tok/sec	ms/tok	Peak GPU	vs T-L4
T-L2 (743M)	542.7	1.84 ms	1483 MB	$1.70\times$
T-L4 (1280M)	318.6	3.14 ms	2511 MB	$1\times$
T-L12 (3427M)	120.2	8.32 ms	6623 MB	$0.38\times$
HRM $N\!\times\!T\!=\!4$ (1229M)	144.3	6.93 ms	2428 MB	$0.45\times$
HRM $N\!\times\!T\!=\!12$ (1229M)	61.3	16.33 ms	2429 MB	$0.19\times$
HRM $N\!\times\!T\!=\!36^{\dagger}$ ( $d\!=\!2048$ , 359M)	55.7	17.96 ms	744 MB	$0.17\times$

Table 16: GenSpeed — autoregressive generation benchmark (prompt

=64

, gen

=128

tokens, full-recompute per token, bfloat16, RTX 6000, 50 reps).

\dagger

d\!=\!2048

, all others

d\!=\!4096

. Ratio relative to T-L4.

Model	CE (normal)	CE (H frozen)	$\Delta$ CE
HRM NT=4 (N=2,T=2)	4.375	5.281	+0.906
HRM NT=8 (N=2,T=4)	4.375	5.938	+1.563
HRM NT=12 (N=4,T=3)	4.313	8.313	+4.000

Table 17: HLRole(A): Slow-module ablation. CE (H frozen) measures quality with Slow-module disabled at inference time;

\Delta

CE is the performance penalty from losing the Slow-module.

Step	H?	gate	$\\|z_{L}\\|$	$\\|z_{H}\\|_{\text{bef}}$	$\\|z_{H}\\|_{\text{aft}}$	$\cos(z_{H},z_{L})$
0		0.281	16.36	56.57	56.57	0.048
1		0.268	10.32	56.57	56.57	0.040
2	$\checkmark$	0.281	8.73	56.57	10.21	0.084
3		0.281	7.94	10.21	10.21	0.085
4		0.288	7.76	10.21	10.21	0.080
5	$\checkmark$	0.292	7.68	10.21	6.85	0.134

Table 18: HLRole(B–C) step statistics for HRM NT=12 (N=4, T=3) on real val text.

\checkmark

: H fires. Gate values 0.28–0.33 (vs. 0.50 on zero tokens) confirm content-sensitive gating.

Panel of Figure 4 — Figure 4: EqualFLOPs learning curves (10k steps, corrected LR). Transformer L=12 (3427M, dashed red) now converges with the corrected LR ( $5\!\times\!10^{-5}$ , max_norm $=0.3$ ), reaching val CE 5.396 at 10k steps, but plateaus after iter 8500 (best 5.378). Initial grad spike at iter 500 (norm $=109$ ) delays convergence, leaving a gap of 1.03–1.23 nats versus HRM variants at 10k. HRM variants converge steadily to 4.16–4.36 (4–10 $\times$ fewer parameters).

Model	Blocks	Params	Mem	Steps †	Saving
TF L=4	4 indep.	1.07B	2.1 GB	4	—
TF L=12	12 indep.	3.22B	6.4 GB	12	—
HRM	3 shared	0.81B	1.6 GB	$N\!\times\!T$	$4\times$

Table 19: Attn block weight memory (

d\!=\!4096

, bf16). † Times

\mathcal{W}_{L}

executes per forward pass.

Metric	TF L=4	HRM $N\!\times\!T\!=\!8$	HRM $N\!\times\!T\!=\!12$
Stored params	1280M	1229M	1229M
KV cache (theory)	$O(4nd)$	$O(8nd)$	$O(12nd)$
Latency (seq=1024)	10.5 ms	31.1 ms	46.8 ms
Latency ratio	$1\times$	$2.97\times$	$4.47\times$
Peak GPU (seq=1024)	2698 MB	2656 MB	2656 MB
val CE (10k)	4.562	4.239	4.233

Table 20: Measured trade-offs (

d=4096

, seq=1024, bfloat16). Latency: no-cache forward pass, single RTX 6000.

Property	HRM-LM	TurboQuant
Target	Stored model weights	KV cache activations
Stored param. mem.	$O(d^{2})$ , $L$ -fold $\checkmark$	Unchanged
KV cache memory	$M/L$ reduction ( $L{>}M$ only)	${\geq}4.5\times$ $\checkmark$
Quality cost	None (CE improved)	Quality-neutral at 3.5 bits
Requires retraining	Yes	No
Latency overhead	$2\text{--}5\times$ slower	Negligible

Table 21: HRM-LM vs. TurboQuant [ 18 ] : orthogonal targets.

Model	M	K/M	3k	5k	7k	10k	grad 10k
HRM $N\!=\!8,T\!=\!6$	48	33%	5.054	4.443	4.255	4.172	1.281
HRM $N\!=\!12,T\!=\!4$	48	25%	4.819	4.741	4.509	4.519	1.570
UniTF flat †	48	33%	5.735	6.461	6.823	6.101	0.812
HRM $N\!=\!4,T\!=\!3$ (M=12)	12	67%	—	—	—	4.177	—
UniTF (M=12)	12	17%	—	—	—	7.578	—

Table 22: M=48 results (10k steps,

d=4096

\approx

1.23B params). M=12 reference values (10k, converged) shown for comparison.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Run	lr / warmup	iter 0 CE	val CE plateau
v2	$5\!\times\!10^{-5}$ / 2000	11.98 (normal)	7.624 @ 10k
v2b	$10^{-4}$ / 500	91.7 (bad seed)	7.613 @ 2k
v2c	$10^{-4}$ / 500, v3-init	109.9 (bad seed)	7.646 @ 4.5k

Appendix

Table 23: UniTF — all UniTF

d\!=\!5760

runs (820M,

M\!=\!12

). All three converge to the same

\approx

7.6 plateau.

iter	val CE
0	139.04
500	7.909
1000	7.651
1500	7.627
2000	7.612
2500	7.553

Appendix

Table 24: UniTF

d\!=\!7296

val CE trajectory (1218M, 10k complete).

Transformer

Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.

See all 528 papers

The Surprising Effectiveness of Shared Memory in Looped Transformers
Oct 1, 2026Giovanni Monea, Keshav Ramji, Yousef El-Kurdi +4Memory-Efficient Inference Recurrent Transformers
Cornell University · IBM Research
The Recurrent Transformer: Greater Effective Depth and Efficient Decoding
Apr 23, 2026Costin-Andrei Oncescu, Depen Morwani, Samy Jelassi +3Autoregressive Language Modeling Self-Attention
Harvard University
Anatomical Heterogeneity in Transformer Language Models
Mar 19, 2026Tomasz WietrzykowskiTransformer Efficient Language Model Training
Independent Researcher, Wroclaw, Poland

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

The Surprising Effectiveness of Shared Memory in Looped Transformers

Giovanni Monea, Keshav Ramji, Yousef El-Kurdi, Luis A. Lastras+3

Cornell University · IBM Research

Memory-Efficient Inference Recurrent Transformers

Open paper AskPDF arXiv

Looped Transformers apply the same layers several times per token, adding compute to improve quality without more parameters. Each recursion, however, writes its own key-value cache, so memory still grows with compute. Inference-time techniques can shrink this cache at a cost in quality. We pretrain looped language models to share memory: only the first recursion writes a cache, and later recursions read it while keeping a short window of their own. Surprisingly, we find that sharing memory does not cost quality and instead improves it. At 150M-1B parameters, our Looped Prediction Transformer (LPT) and its hybrid variant set a new quality-memory frontier for looped models: with five recursions, the hybrid lowers validation perplexity on FineWeb-Edu by 1.12-1.82 relative to a same-size standard Transformer while using 76-79% less context memory. Through an extensive analysis, we investigate why memory sharing helps. Shared and local memory develop different representations, and later recursions attend mostly to the shared memory, which also acts as a gradient highway to the first recursion.

Abstract

Transformer

Explore similar work

The Surprising Effectiveness of Shared Memory in Looped Transformers

The Recurrent Transformer: Greater Effective Depth and Efficient Decoding

Anatomical Heterogeneity in Transformer Language Models