Sparse Growing Transformer: Training-Time Sparse Depth Allocation via Progressive Attention Looping

Yao Chen; Yilong Chen; Yinqi Yang; Junyuan Shang; Zhenyu Zhang; Zefeng Zhang; Shuaiyi Nie; Shuohuan Wang

cs.CLMar 25, 2026

Sparse Growing Transformer: Training-Time Sparse Depth Allocation via Progressive Attention Looping

Authors: Yao Chen, Yilong Chen, Yinqi Yang, Junyuan Shang, Zhenyu Zhang, Zefeng Zhang, Shuaiyi Nie, Shuohuan Wang, +4 more

Organizations: Institute of Information Engineering, Chinese Academy of Sciences · School of Cyber Security, University of Chinese Academy of Sciences · Baidu Inc.

Abstract

Existing approaches to increasing the effective depth of Transformers predominantly rely on parameter reuse, extending computation through recursive execution. Under this paradigm, the network structure remains static along the training timeline, and additional computational depth is uniformly assigned to entire blocks at the parameter level. This rigidity across training time and parameter space leads to substantial computational redundancy during training. In contrast, we argue that depth allocation during training should not be a static preset, but rather a progressively growing structural process. Our systematic analysis reveals a deep-to-shallow maturation trajectory across layers, where high-entropy attention heads play a crucial role in semantic integration. Motivated by this observation, we introduce the Sparse Growing Transformer (SGT). SGT is a training-time sparse depth allocation framework that progressively extends recurrence from deeper to shallower layers via targeted attention looping on informative heads. This mechanism induces structural sparsity by selectively increasing depth only for a small subset of parameters as training evolves. Extensive experiments across multiple parameter scales demonstrate that SGT consistently outperforms training-time static block-level looping baselines under comparable settings, while reducing the additional training FLOPs overhead from approximately 16--20% to only 1--3% relative to a standard Transformer backbone.

Figures & tables

Figure 1: Visualization of attention heatmap in low- and high-entropy heads from Qwen3-0.6B. Along the horizontal axis, tokens highlighted in red denote the subset that receives the top 50% of the attention from the final query position (details in Appendix C.3.2 ).

Figure 2: Evolution of Intra-layer Variance in Head-wise Attention Entropy Across Training Steps for Models Pretrained from Scratch (we visualize four representative layers; see Figure 14(b) for detailed results)

Figure 3: The Sparse Growing Transformer (SGT) architecture. It implements Training-Phase Structural Sparsity via a deep-to-shallow progressive growth schedule (Left) and selective high-entropy head looping (Right).

Model	FLOPs	Reasoning & Knowledge ( $\uparrow$ )
Model	FLOPs	ARC-E	WG	SIQA	Hella.	OBQA	CSQA	BA	MMLU	Avg
Vanilla (275M)	42.83	44.56	49.57	41.30	30.89	27.40	29.24	22.70	24.80	33.81
Block Loop	51.50 ( +20.24% )	44.74	50.36	41.10	31.85	27.60	29.48	21.80	25.38	34.04
SGT(K=1)	43.46 ( 1.47% )	45.09	51.46	41.40	31.21	28.00	29.16	24.07	25.19	34.48
SGT(K=2)	43.88 ( +2.45% )	45.79	52.33	42.38	31.41	26.80	28.75	24.47	25.45	34.64
SGT(K=3)	44.21 ( +3.22% )	46.32	50.91	41.81	31.43	27.40	28.83	22.37	26.16	34.40

Table 1: Main experimental results. FLOPs denote the total training compute, measured in units of

10^{18}

. The best results at each model scale are shown in bold .

K

denotes

K_{\max}

, the upper bound on the loop depth for the selected loop layers.

Figure 4: Comparative convergence trajectories of training perplexity (PPL) versus cumulative training FLOPs for the 573M model scale.

Loop Component	FLOPs ( $\times 10^{18}$ )	Perplexity $\downarrow$
Vanilla	87.41	23.973
Layer 2
Block Loop	92.20 (+5.48%)	23.560 (-0.413)
Attention Loop	89.01 (+1.83%)	23.503 (-0.470)
High-Ent. Loop	87.74 (+0.38%)	23.452 (-0.521)
Low-Ent. Loop	87.74 (+0.38%)	23.502 (-0.471)

Table 2: Ablation study of different loop components at varying layer depths. We report total training FLOPs (in units of

10^{18}

) and perplexity (PPL) on the validation set, where lower is better. High-Ent. Loop denotes looping the top-2 highest-entropy heads, while Low-Ent. Loop denotes looping the top-2 lowest-entropy heads.

Method	CSQA	SIQA	SCIQ	$\text{MMLU}_{\text{STEM}}$	Avg
Vanilla	29.24	41.86	71.00	26.28	42.10
S2D	30.39	41.40	71.80	27.40	42.75
D2S (Ours)	31.29	42.22	72.00	28.46	43.49

Table 3: Ablation study on growth direction. Experiments are conducted on the 573M model with SGT (

K=1

) configuration. D2S denotes deep-to-shallow growth, while S2D represents shallow-to-deep growth.

Figure 5: Attention entropy dynamics of two heads in Layer 6 during training, including the warm-up phase and layer selection phase, for models trained with D2S and S2D strategies (more visualizations in Figure 15 ).

Setting	Vanilla	Block Loop	High-Ent. Loop
1024 $\times$ 1	24.21	24.18 ( $-0.03$ )	24.16 ( $-0.05$ )
1024 $\times$ 2	57.38	58.80 ( $+1.42$ )	56.62 ( $-0.76$ )
1024 $\times$ 3	116.94	119.48 ( $+2.54$ )	111.66 ( $-5.28$ )
1024 $\times$ 4	180.70	175.05 ( $-5.65$ )	171.90 ( $-8.80$ )

Table 4: Long Context Extrapolation Results (PPL). Setting denotes the sequence length for extrapolation evaluation. The configurations for High-Ent. Loop and Block Loop follow the Layer 2 settings in Table 2 .

Figure 6: Attention contribution flow across loops in a high-entropy head (Layer 2, Head 15). Each cell shows the weighted attention contribution $C_{j}=A_{N,j}\cdot\|V_{j}\|_{2}$ from the final token to position $j$ , with color intensity indicating magnitude. Red and blue borders mark tokens with the largest contribution increase and decrease from Loop 0 to Loop 2, respectively. Arrows with percentage labels denote inter-loop changes. The visualization shows attention progressively shifting from syntactic elements toward task-critical and answer-relevant tokens.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Benchmarks	Avg. Input Length
Mathematical Reasoning
AIME Balunović et al. (2025)	160 tokens
MATH-500 Lightman et al. (2023)	89 tokens
GSM8K Cobbe et al. (2021)	78 tokens
General Knowledge QA
MMLU Hendrycks et al. (2020)	117 tokens

Appendix

Table 5: Overview of the 11 diverse reasoning benchmarks, categorized by task family, used for analyzing the model’s attention entropy. The tasks span various domains and sequence lengths.

Model	layers	heads	hidden size	intermediate size
Qwen3-0.6B-Base	28	16	1024	3072
LLaMA3.2-1B-Base	16	32	2048	8192
Qwen3-1.7B-Base	28	16	2048	6144

Appendix

Table 6: Architectural hyperparameters of three representative open-source models used in our observational study of Section 3 . heads denotes the number of query heads.

Figure 7: Layer-wise attention entropy across 11 diverse tasks for three open-source models (details in Appendix C.1 and C.2.1 ).

Figure 8: Layer–head attention entropy pattern of Qwen3-0.6B and Qwen3-1.7B. The x-axis denotes the head index, and the y-axis denotes the layer index (details in Appendix C.1 and C.2.2 ).

Figure 9: Layer–head attention entropy pattern of LLaMA3.2-1B-Base. The x-axis denotes the head index, and the y-axis denotes the layer index (details in Appendix C.1 and C.2.2 ).

Mask Settings	Original	High-Entropy	Low-Entropy	Random Heads (Various Seeds)
Mask Settings	Model	Heads	Heads	Set 1	Set 2	Set 3	Set 4
Masked Heads	-	L1-H1, L1-H7 L1-H12, L2-H6	L4-H9, L25-H2 L28-H4	L4-H1, L4-H5 L4-H12, L21-H8	L2-H5, L2-H6 L2-H8, L10-H4	L1-H9, L4-H8 L21-H8	L2-H5, L10-H15 L23-H12
MATH-500	27.0	5.0 $\downarrow 81.5\%$	25.0 $\downarrow 7.4\%$	20.0 $\downarrow 25.9\%$	24.0 $\downarrow 11.1\%$	25.0 $\downarrow 7.4\%$	23.0 $\downarrow 14.8\%$
GSM8K	32.0	7.0 $\downarrow 78.1\%$	29.0 $\downarrow 9.4\%$	37.0 $\uparrow 15.6\%$	32.0 $\pm 0.0\%$	32.0 $\pm 0.0\%$	27.0 $\downarrow 15.6\%$
MMLU	50.5	18.0 $\downarrow 64.4\%$	46.0 $\downarrow 8.9\%$	45.5 $\downarrow 9.9\%$	47.0 $\downarrow 6.9\%$	46.0 $\downarrow 8.9\%$	50.5 $\pm 0.0\%$
NarrativeQA	13.7	11.0 $\downarrow 19.7\%$	14.2 $\uparrow 3.7\%$	13.6 $\downarrow 0.7\%$	12.2 $\downarrow 11.0\%$	12.6 $\downarrow 8.0\%$	14.5 $\uparrow 5.8\%$

Appendix

Table 7: Impact of head masking on Qwen3-0.6B-Base performance across reasoning and long-context tasks. The Original Model column represents the unmasked model baseline. Mask Settings compare the removal of High-Entropy heads against Low-Entropy and four groups of Random head baselines. The specific indices of removed heads are detailed in the Masked Heads row, denoted as L x -H y (Layer x , Head y , 1-indexed). Cells report absolute performance scores and the relative percentage change compared to the Original Model (details in Appendix C.3.1 and C.3.3 ).

Figure 10: Layer-wise Attention Entropy at Different Training Steps. The horizontal axis represents the model’s layer index. Each line shows the entropy distribution recorded at 250-step intervals during the early stages of training (details in Appendix C.4 ).

Model Setting	OLMo-275M	OLMo-573M	OLMo-1.2B
hidden size	768	1024	1536
intermediate size	6144	8192	12288
attention heads	12	16	16
layers	12	16	16

Appendix

Table 8: Detailed architecture configurations for the three model scales (275M, 573M, and 1.2B) pre-trained from scratch in our main experiments. We pre-train models from scratch at three scales using the OLMo framework. All models adopt the LLaMA architecture with Multi-Head Attention (MHA).

Model	ARC-E	WG	SIQA
Vanilla	44.56	49.57	41.30
Block Loop	44.74	50.36	41.10
SGT ( $t_{\text{start}}=\Delta t=500$ )	45.26	50.91	41.56
SGT ( $t_{\text{start}}=\Delta t=250$ )	45.09	51.46	41.40
SGT ( $t_{\text{start}}=\Delta t=100$ )	45.44	51.62	41.97

Appendix

Table 9: Analysis of growth schedule parameters (

t_{\text{start}}

and

\Delta t

) for the 275M model.

Batch	Model	Prefill TPS	Decode TPS
1	Vanilla (275M)	14381.0	127.9
	SGT (275M)	13805.7	122.7
	Block Loop (275M)	11504.8	102.3
	Vanilla (573M)	10969.2	100.5
	SGT (573M)	10718.0	98.2
	Block Loop (573M)	9237.2	84.6

Appendix

Table 10: Inference throughput comparison (tokens/s) under different batch sizes. Prefill throughput measures tokens processed during context encoding (KV cache construction), while decode throughput measures autoregressive token generation speed.

Figure 11: Attention contribution flow across loops in a high-entropy head (Layer 2, Head 15). Each cell shows the weighted attention contribution $C_{j}=A_{N,j}\cdot\|V_{j}\|_{2}$ from the final token to position $j$ , with color intensity indicating magnitude. Red and blue borders mark tokens with the largest contribution increase and decrease from Loop 0 to Loop 2, respectively. Arrows with percentage labels denote inter-loop changes. The visualization shows attention progressively shifting from syntactic elements toward task-critical and answer-relevant tokens.

Figure 12: Visualization of attention heatmap in low- and high-entropy heads from Qwen3-0.6B, with a sample from MATH-500. Along the horizontal axis, tokens highlighted in red denote the subset that receives the top 50% of the attention from the final query position (details in Appendix C.3.1 and C.3.2 ).

Figure 13: Visualization of attention heatmap in low- and high-entropy heads from Qwen3-0.6B, with a sample from MMLU. Along the horizontal axis, tokens highlighted in red denote the subset that receives the top 50% of the attention from the final query position (details in Appendix C.3.1 and C.3.2 ).

Panel of Figure 14 — Figure 14: (a) Head-wise attention entropy trajectories across all layers during training . Within each layer-specific subplot, individual curves correspond to distinct attention heads. (b) Intra-layer variance of attention entropy across heads over the course of training (details in Appendix C.4 ) .

Figure 15: Attention entropy dynamics of Layer 6 heads during training, including the warm-up phase (steps 0–250) and layer selection phase (steps 250–1000), for models trained with D2S and S2D strategies in Table 3 .

Explore similar work

DeepLoop: Depth Scaling for Looped Transformers
Jul 15, 2026Shuzhen Li, Yifan Zhang, Jiacheng Guo +2Residual Learning Recurrent Transformers
Princeton University · University of California, Los Angeles
Prism Transformer: Progressive Head Schedules for Hierarchical Attention Processing
Jun 25, 2026Shubham AggarwalHierarchical Representation Learning Transformer Attention
CascadeFormer: Depth-Tapered Transformers Motivated by Gradient Fan-in Asymmetry
Jun 25, 2026Huzama Ahmad, Cao Viet Hai Nam, Se-Young YunEfficient Transformer Inference Transformer
KAIST

Jul 15, 2026J/K move · Enter open · S save

DeepLoop: Depth Scaling for Looped Transformers

Shuzhen Li, Yifan Zhang, Jiacheng Guo, Quanquan Gu+1

Princeton University · University of California, Los Angeles

Residual Learning Recurrent Transformers

Open paper AskPDF arXiv

Looped Transformers scale sequential computation by applying a compact stack of physical blocks for multiple rounds, increasing unrolled depth without increasing stored parameters. This reuse changes the residual-scaling problem: in an untied Transformer, each residual branch receives and applies its own parameter update, whereas in a looped Transformer one shared update aggregates gradients from repeated visits and is read back by those same visits in the next linearized forward pass. We formalize this tied-depth effect through a first-order perturbation bound controlled by a visit-alignment coefficient

κ_R

. The bound recovers the DeepNorm exponent when visits decorrelate, but in the conservative aligned regime it requires the exponent to increase from

1/4

1/2

as loop count grows at fixed physical depth. The resulting method, DeepLoop, keeps the Post-LN DeepNorm architecture and sets

α=(2N)^{1/2}

and

β=(8N)^{-1/2}

for unrolled depth

N

. On GPT-style looped language models at GPT-2 small and GPT-2 medium scale, DeepLoop is neutral when no physical block is revisited and improves validation loss and downstream accuracy once recurrent depth is activated. These results show that stable recurrent depth requires residual scaling rules that account for parameter visits, not only nominal layer count.

Abstract

Explore similar work

DeepLoop: Depth Scaling for Looped Transformers

Prism Transformer: Progressive Head Schedules for Hierarchical Attention Processing

CascadeFormer: Depth-Tapered Transformers Motivated by Gradient Fan-in Asymmetry