cs.LGOct 29, 2026
SaveHow to Loop Your Transformer: Loop Scaling, Specialization, and Parallel Inference
Organizations: FAIR at Meta · MIT
Abstract
Looping in transformer models reuses weights by repeatedly applying the same transformer, or parts of it, making the model deeper without increasing its parameter count. Whether looping is useful depends on which constraint dominates: model size, training compute, or autoregressive inference efficiency. We study these trade-offs jointly in dense decoder-only transformers trained from scratch, varying which parts of the model are looped, how many loops are applied, and how information is passed between loops. Our experiments span models with up to 1B parameters, trained on up to 155B tokens with up to FLOPs and 32 loops. We find that looping the full transformer generally gives the strongest quality at a fixed parameter budget, whereas looping a middle block of layers gives a better return when training FLOPs are constrained. Adding loops is not uniformly beneficial: gains eventually saturate, reverse, or become unstable. This motivates asking whether forcing every loop to use the same weights limits the benefit of deeper looping. We address this by adding small loop-specific low-rank adapters that improve quality, with larger gains as more loops are added. Looped transformers also suffer a decoding slowdown proportional to the number of loops. Parallel Loop Transformers address this by passing hidden states between loops across adjacent token positions, allowing loop stages to run concurrently. However, in our experiments the original formulation does not consistently retain the quality gains of ordinary looping. We find that preserving the current-token embedding and separately normalizing it and the previous-token hidden state before combining them recovers strong performance across model scales and training budgets. Together, these results show that effective looping requires matching the design to the resource constraint and inference setting.
Figures & tables
Figure 1: Looping can improve data and compute efficiency (in over-trained regimes) while retaining downstream quality. We compare Vanilla (no looping), Standard Looping (repeatedly applying the same Transformer sequentially), and our modified versions Parallel Loop Transformer (PLT) ( Wu et al., 2025 ) and Looped Transformers. Left: for 1B models trained across TPP 10–160, all looping methods achieve a given DCLM ( Li et al., 2024 ) bits-per-byte (BPB) value with fewer training tokens than the fitted Vanilla baseline, with the modified PLT variants providing the largest relative data-efficiency gains. Center: across model sizes from 100M to 1B, we compare Vanilla models trained at TPP with two-loop models trained at TPP . All looping methods achieve a given DCLM BPB with fewer training FLOPs than the fitted Vanilla baseline, and the modified PLT and Looped variants are more efficient than Standard Looping. Right: for 1B models across TPP 10–160, looping retains or improves mean NLL-ranked accuracy across eight reasoning and question-answering tasks; the modified PLT variants match or exceed Standard Looping. See App. D for the efficiency calculation and task-level results.
Figure 2: Looping designs: what is repeated, where, and in what order. Top: the looped component can be the complete transformer block, attention alone, the MLP alone, or attention and the MLP separately. Bottom: full-stack looping traverses the complete transformer for loops, whereas per-layer looping completes loops of each block before advancing to the next. Middle-stack and middle-layer looping restrict the corresponding schedule to the central blocks. For a model containing blocks, the first and last blocks are run once.
Figure 3: Full-Stack Looping is favored under fixed parameter-count and training-data budgets, whereas Middle-Stack Looping is favored under FLOP constraints. DCLM validation BPB versus the number of loops or total training FLOPs for 100M and 500M models. Each panel holds the model size and training-data budget constant, so the differences reflect how each strategy uses the additional computation introduced by looping. The advantage of full-stack looping becomes clearer at larger model scales and training budgets. Middle-Stack Looping with fewer loops achieves lower BPB when compute is limited, whereas Full-Stack Looping can attain lower BPB when more compute is available. Increasing the number of loops does not improve quality indefinitely: the 100M results exhibit saturation, regression, and instability when more loops are used. Lower BPB is better; red triangles indicate unstable or clipped runs.
Figure 4: Three ways to vary weights across loops. All designs retain the same base weights. Loop LoRA gives each loop a separate low-rank adapter. Loop Matrix lets each loop learn a combination of these adapters, while Loop Softmax restricts this combination to be nonnegative and sum to one.
Figure 5: Loop-specific corrections become more useful with deeper looping while adding few parameters. DCLM BPB for 300M models with Full-Stack Looping at TPP 100. The left, middle, and right panels plot BPB against loop count, non-embedding parameter count, and total training FLOPs, respectively. The break in the logarithmic parameter axis shows both parameter ranges. FLOPs include the corrections and the training tokens allocated at TPP 100. Lower BPB is better.
Figure 6: PLT shifts the connection between loops across tokens. Loop at token receives the current-token embedding and the state from loop at token . KV sharing and G-SWA are omitted for clarity.
Figure 7: Separate normalization recovers much of the quality lost by the shifted transition. DCLM BPB for 200M models at TPP 40 and 80. Each raw panel is followed by its normalized counterpart. Except for Normalized Addition, normalized variants use separate RMS normalization and fixed 50/50 mixing. Red stars mark unstable runs clipped for visualization. Lower BPB is better.
Figure 8: Normalized shifted PLT retains the benefit of looping at 1B. DCLM BPB at TPP 10, 20, 40, and 80. Shifted and unshifted PLT use separate RMS normalization and a rank- linear gate. At , shifted PLT closely tracks the unshifted control and improves upon Full-Stack Looping ; at , a clearer gap to the unshifted control remains. Lower BPB is better.
Appendix figures & tables128 assets
Supplementary material from the paper’s appendix.
Appendix
| Architecture | Parameter count (M) | |||||
| Model label | Transformer layers | Model width | MLP width | Attention heads (Q/KV) | Excluding input embedding | Including input embedding |
| 70M | 4 | 512 | 1,536 | 4/1 | 77.726 | 143.397 |
| 100M | 6 | 768 | 2,048 | 6/1 | 135.070 | 233.580 |
| 200M | 8 | 1,024 | 2,816 | 8/1 | 219.415 | 350.766 |
| 300M | 10 | 1,280 | 3,584 | 10/1 | 337.838 | 502.033 |
| 500M | 12 | 1,536 | 4,096 | 12/2 | 489.554 | 686.594 |
Table 1: Model configurations used in our experiments. Transformer layers give the number of Transformer blocks. Model width is the dimension of the representation passed between layers, and MLP width is the intermediate dimension of the SwiGLU MLP. Q/KV gives the numbers of query heads and key–value heads, respectively. The first parameter column excludes the input embedding and is used to assign the model label and set the training-token budget. The final column includes the input embedding and gives the total number of trainable parameters. Parameter counts are reported in millions. All models use a vocabulary of 128,256 tokens, a context length of 8,192, and an attention-head dimension of 128.
Figure 9: Raw scaling curves underlying the relative-efficiency analysis. We plot observed DCLM validation BPB together with unconstrained Vanilla scaling-law fits. Left: BPB as a function of actual training tokens for fixed 1B models across training-token budgets. Middle: BPB as a function of total training FLOPs across model families from 100M to 1B. Vanilla models use TPP , whereas looping models use two loops and TPP . Points show observed runs, connecting lines show empirical trajectories, and dashed black curves show fitted Vanilla scaling laws. Lower BPB is better.
Figure 10: Task-level NLL-ranked accuracy for the 1B models. Results are shown for ARC Easy, ARC Challenge, HellaSwag, PIQA, WinoGrande, OpenBookQA, CommonsenseQA, and Social IQA. For each question, the model predicts the candidate with the lowest token-normalized NLL. Dashed horizontal lines denote random-choice accuracy. Lines connect measured checkpoints and are not fitted scaling laws. Higher is better.
Figure 11: Candidate-normalized correct-answer NLL for the same eight tasks. Candidate scores are computed from token-normalized NLL and normalized across the available answers. This continuous metric measures how strongly the model prefers the correct answer relative to the alternatives, including changes that do not alter the top-ranked answer. Dashed lines show the random-choice value for a task with candidates. Because choice counts and prompt formats differ, values should be compared within, rather than across, datasets. Lower is better.
Figure 12: MMLU and MMLU-Pro results for the 1B models. Top: macro-averaged NLL-ranked accuracy, where the predicted answer is the candidate with the lowest token-normalized NLL. Bottom: candidate-normalized correct-answer NLL. Accuracy remains near the respective random-choice baselines, while candidate-normalized correct-answer NLL remains at or above the random-choice reference, particularly on MMLU-Pro. Dashed lines denote random-choice performance. Higher accuracy and lower NLL are better.
Figure 13: DCLM BPB versus loop factor : 70M and 100M models. Results span TPP 10 to 160 and loop factors . Each panel fixes model scale and TPP, allowing the effect of increasing the loop factor to be compared across looping designs. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 14: DCLM BPB versus loop factor : 200M, 300M, and 500M models. Results span TPP 10 to 100 and loop factors . Each panel fixes model scale and TPP, allowing the effect of increasing the loop factor to be compared across looping designs. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 15: DCLM BPB versus total training FLOPs: 70M and 100M models. Results span TPP 10 to 160 and loop factors .
Figure 16: DCLM BPB versus total training FLOPs: 200M, 300M, and 500M models. Results span TPP 10 to 100 and loop factors . Middle-Stack Looping often achieves strong gains for a smaller increase in compute, while Full-Stack Looping can attain lower BPB when more training compute is available.
Figure 17: C4 English BPB versus loop factor : 70M and 100M models. Results span TPP 10 to 160 and loop factors . At these smaller scales, several looping designs remain competitive, while the improvement from increasing generally diminishes at larger loop factors. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 18: C4 English BPB versus loop factor : 200M, 300M, and 500M models. Results span TPP 10 to 100 and loop factors . As model scale and TPP increase, Full-Stack Looping more consistently attains lower BPB, while the marginal improvement from increasing becomes smaller. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 19: C4 English BPB versus total training FLOPs: 70M and 100M models. Results span TPP 10 to 160 and loop factors . Middle-Stack Looping often captures much of the quality improvement with a smaller increase in training compute, while spending additional FLOPs through larger loop factors eventually shows diminishing returns. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 20: C4 English BPB versus total training FLOPs: 200M, 300M, and 500M models. Results span TPP 10 to 100 and loop factors . Middle-Stack Looping often provides strong gains for a smaller increase in training compute, while Full-Stack Looping can attain lower BPB when more compute is available, particularly at the larger model scales. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 21: WikiText-103 BPB versus loop factor : 70M and 100M models. Results span TPP 10 to 160 and loop factors . At these smaller scales, several looping designs remain competitive, while the improvement from increasing generally diminishes at larger loop factors. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 22: WikiText-103 BPB versus loop factor : 200M, 300M, and 500M models. Results span TPP 10 to 100 and loop factors . As model scale and TPP increase, Full-Stack Looping more consistently attains lower BPB, while the marginal improvement from increasing becomes smaller. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 23: WikiText-103 BPB versus total training FLOPs: 70M and 100M models. Results span TPP 10 to 160 and loop factors . Middle-Stack Looping often captures much of the quality improvement with a smaller increase in training compute, while increasing the loop factor further eventually gives diminishing returns. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 24: WikiText-103 BPB versus total training FLOPs: 200M, 300M, and 500M models. Results span TPP 10 to 100 and loop factors . Middle-Stack Looping often provides strong gains for a smaller increase in training compute, while Full-Stack Looping can attain lower BPB when more compute is available, particularly at larger model scales. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 25: Dolma v1.5 BPB versus loop factor : 70M and 100M models. Results span TPP 10 to 160 and loop factors . At these smaller scales, several looping designs remain competitive, while the improvement from increasing generally diminishes at larger loop factors. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 26: Dolma v1.5 BPB versus loop factor : 200M, 300M, and 500M models. Results span TPP 10 to 100 and loop factors . As model scale and TPP increase, Full-Stack Looping more consistently attains lower BPB, while the marginal improvement from increasing becomes smaller. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 27: Dolma v1.5 BPB versus total training FLOPs: 70M and 100M models. Results span TPP 10 to 160 and loop factors . Middle-Stack Looping often captures much of the quality improvement with a smaller increase in training compute, while increasing the loop factor further eventually gives diminishing returns. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 28: Dolma v1.5 BPB versus total training FLOPs: 200M, 300M, and 500M models. Results span TPP 10 to 100 and loop factors . Middle-Stack Looping often provides strong gains for a smaller increase in training compute, while Full-Stack Looping can attain lower BPB when more compute is available, particularly at larger model scales. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 29: RedPajama BPB versus loop factor : 70M and 100M models. Results span TPP 10–160 and loop factors . At these smaller scales, several looping designs remain competitive, while the improvement from increasing generally diminishes at larger loop factors. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 30: RedPajama BPB versus loop factor : 200M, 300M, and 500M models. Results span TPP 10–100 and loop factors . As model scale and TPP increase, Full-Stack Looping more consistently attains lower BPB, while the marginal improvement from increasing becomes smaller. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 31: RedPajama BPB versus total training FLOPs: 70M and 100M models. Results span TPP 10–160 and loop factors . Middle-Stack Looping often captures much of the quality improvement with a smaller increase in training compute, while increasing the loop factor further eventually gives diminishing returns. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 32: RedPajama BPB versus total training FLOPs: 200M, 300M, and 500M models. Results span TPP 10–100 and loop factors . Middle-Stack Looping often provides strong gains for a smaller increase in training compute, while Full-Stack Looping can attain lower BPB when more compute is available, particularly at larger model scales. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 33: Dolma programming-language BPB versus loop factor : 70M and 100M models. Results span TPP 10–160 and loop factors . At these smaller scales, Middle-Stack Looping is often competitive with Full-Stack Looping, while the improvement from increasing generally diminishes at larger loop factors. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 34: Dolma programming-language BPB versus loop factor : 200M, 300M, and 500M models. Results span TPP 10–100 and loop factors . As model scale and TPP increase, Full-Stack Looping becomes more consistently favorable, while the marginal improvement from increasing becomes smaller. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 35: Dolma programming-language BPB versus total training FLOPs: 70M and 100M models. Results span TPP 10–160 and loop factors . Middle-Stack Looping often provides strong quality improvements for a smaller increase in training compute, while spending additional FLOPs through larger loop factors eventually shows diminishing returns. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 36: Dolma programming-language BPB versus total training FLOPs: 200M, 300M, and 500M models. Results span TPP 10–100 and loop factors . Middle-Stack Looping remains competitive for smaller increases in training compute, while Full-Stack Looping can attain lower BPB when more compute is available, particularly at the larger model scales. Each panel fixes model scale and TPP. Lower BPB is better; red stars denote unstable runs as described in Subsec. E.1 .
Figure 37: GSM8K BPB versus loop factor : 70M and 100M models. Results span TPP 10–160 and loop factors . Several looping designs remain competitive, while the improvement from increasing generally diminishes at larger loop factors. Each panel fixes model scale and TPP. BPB measures teacher-forced likelihood rather than answer accuracy; lower BPB is better. Red stars denote unstable runs as described in Subsec. E.1 .
Figure 38: GSM8K BPB versus loop factor : 200M, 300M, and 500M models. Results span TPP 10–100 and loop factors . Full-Stack Looping becomes more favorable as model scale and TPP increase, although the differences across looping designs remain less uniform than on the general pretraining corpora. Each panel fixes model scale and TPP. BPB measures teacher-forced likelihood rather than answer accuracy; lower BPB is better.
Figure 39: MATH BPB versus loop factor : 70M and 100M models. Results span TPP 10–160 and loop factors . Several looping designs remain competitive at smaller training budgets, while the improvement from increasing generally diminishes at larger loop factors. Each panel fixes model scale and TPP. BPB measures teacher-forced likelihood rather than answer accuracy; lower BPB is better. Red stars denote unstable runs as described in Subsec. E.1 .
Figure 40: MATH BPB versus loop factor : 200M, 300M, and 500M models. Results span TPP 10–100 and loop factors . Full-Stack Looping becomes more consistently favorable as model scale and TPP increase. Each panel fixes model scale and TPP. BPB measures teacher-forced likelihood rather than answer accuracy; lower BPB is better.
Figure 41: HumanEval BPB versus loop factor : 70M and 100M models. Results span TPP 10–160 and loop factors . The preferred design varies across training budgets, and increasing does not always improve BPB. Each panel fixes model scale and TPP. BPB measures teacher-forced likelihood rather than execution accuracy; lower BPB is better. Red stars denote unstable runs as described in Subsec. E.1 .
Figure 42: HumanEval BPB versus loop factor : 200M, 300M, and 500M models. Results span TPP 10–100 and loop factors . Full-Stack Looping becomes more favorable at larger model scales and higher TPP, although it does not dominate every setting. Each panel fixes model scale and TPP. Lower BPB is better.
Figure 43: MBPP BPB versus loop factor : 70M and 100M models. Results span TPP 10–160 and loop factors . Several looping designs remain competitive, while the improvement from increasing generally diminishes at larger loop factors. Each panel fixes model scale and TPP. BPB measures teacher-forced likelihood rather than execution accuracy; lower BPB is better. Red stars denote unstable runs as described in Subsec. E.1 .
Figure 44: MBPP BPB versus loop factor : 200M, 300M, and 500M models. Results span TPP 10–100 and loop factors . Full-Stack Looping becomes more favorable as model scale and TPP increase, although the preferred design remains more task-dependent than on the general pretraining corpora. Each panel fixes model scale and TPP. Lower BPB is better.
Figure 45: DCLM BPB versus loop factor at 1B. Full-stack, middle-stack, and Per-Layer Looping are evaluated at TPP 10 and 40 over . All three improve over the non-looped model, with Full-Stack Looping attaining the lowest BPB at matched loop counts. Lower is better.
Figure 46: DCLM BPB versus training FLOPs at 1B. The same models are compared at TPP 10 and 40. Middle-Stack Looping improves BPB with a smaller increase in FLOPs, while full-stack looping reaches lower BPB with more compute. The horizontal axis is logarithmic; marker shape denotes . Lower is better.
Figure 47: C4 BPB versus loop factor at 1B. Non-looped, full-stack, middle-stack, and per-layer models are compared at TPP 10 and 40 over . Each panel fixes TPP; lower BPB is better.
Figure 48: C4 BPB versus training FLOPs at 1B. The same checkpoints are plotted against total training FLOPs at TPP 10 and 40. The horizontal axis is logarithmic; marker shape denotes . Lower BPB is better.
Figure 49: Dolma v1.5 BPB versus loop factor at 1B. Non-looped, full-stack, middle-stack, and per-layer models are compared at TPP 10 and 40 over . Each panel fixes TPP; lower BPB is better.
Figure 50: Dolma v1.5 BPB versus training FLOPs at 1B. The same checkpoints are plotted against total training FLOPs at TPP 10 and 40. The horizontal axis is logarithmic; marker shape denotes . Lower BPB is better.
Figure 51: Dolma programming-languages BPB versus loop factor at 1B. Non-looped, full-stack, middle-stack, and per-layer models are compared at TPP 10 and 40 over . Each panel fixes TPP; lower BPB is better.
Figure 52: Dolma programming-languages BPB versus training FLOPs at 1B. The same checkpoints are plotted against total training FLOPs at TPP 10 and 40. The horizontal axis is logarithmic; marker shape denotes . Lower BPB is better.
Figure 53: RedPajama BPB versus loop factor at 1B. Non-looped, full-stack, middle-stack, and per-layer models are compared at TPP 10 and 40 over . Each panel fixes TPP; lower BPB is better.
Figure 54: RedPajama BPB versus training FLOPs at 1B. The same checkpoints are plotted against total training FLOPs at TPP 10 and 40. The horizontal axis is logarithmic; marker shape denotes . Lower BPB is better.
Figure 55: WikiText-103 BPB versus loop factor at 1B. Non-looped, full-stack, middle-stack, and per-layer models are compared at TPP 10 and 40 over . Each panel fixes TPP; lower BPB is better.
Figure 56: WikiText-103 BPB versus training FLOPs at 1B. The same checkpoints are plotted against total training FLOPs at TPP 10 and 40. The horizontal axis is logarithmic; marker shape denotes . Lower BPB is better.
Figure 57: Specialization helps more as the number of loops increases, consistently across TPP. DCLM BPB versus loop count for 70M–300M full-stack models at TPP 10, 40, and 100. Specialization provides a small improvement at and clearer gains at and . Increasing TPP also generally lowers BPB. Lower is better.
Figure 58: Specialization improves DCLM BPB with relatively few additional parameters across TPP. Results cover 70M–300M full-stack models at TPP 10, 40, and 100. The horizontal axis reports additional non-embedding parameters relative to the identical-weight model. Higher TPP generally lowers BPB. Lower is better.
Figure 59: Adding independently parameterized layers can achieve lower BPB but requires substantially more storage. Rows show 100M, 200M, 300M, and 500M models; columns show TPP 10, 40, and 100. The purple diamond is the matched doubled-depth model. The final tick after the broken axis gives its true increase in non-embedding parameters. Lower BPB is better.
Figure 60: When parameters can be stored, additional physical depth can be a better use of training FLOPs. DCLM BPB versus total training FLOPs for 70M–300M models at TPP 10, 40, and 100. Training FLOPs include repeated loop execution, loop-specific corrections, attention-score computation, and the realized token budget. Lower BPB is better.
Figure 61: Specialization generally helps more as the number of loops increases. C4 English BPB versus loop count for 70M–300M full-stack models at TPP 10, 40, and 100. Specialization is generally more useful at and than at . Increasing TPP also generally lowers BPB. Lower is better.
Figure 62: Specialization improves C4 English BPB with relatively few additional parameters. Results cover 70M–300M full-stack models at TPP 10, 40, and 100. The horizontal axis reports additional non-embedding parameters relative to the identical-weight model. Higher TPP generally lowers BPB. Lower is better.
Figure 63: Adding independently parameterized layers can achieve lower C4 English BPB but requires substantially more storage. Rows show 100M, 200M, 300M, and 500M models; columns show TPP 10, 40, and 100. The purple diamond is the matched doubled-depth Transformer. The final tick after the broken axis gives its true increase in non-embedding parameters. Lower BPB is better.
Figure 64: Low parameter overhead does not necessarily mean better training-compute efficiency. C4 English BPB versus total training FLOPs for 70M–300M models at TPP 10, 40, and 100. Training FLOPs include repeated loop execution, loop-specific corrections, attention-score computation, and the realized token budget. Lower is better.
Figure 65: Specialization generally helps more as the number of loops increases. WikiText-103 BPB versus loop count for 70M–300M full-stack models at TPP 10, 40, and 100. Specialization is generally more useful at and than at . Increasing TPP also generally lowers BPB. Lower is better.
Figure 66: Specialization improves WikiText-103 BPB with relatively few additional parameters. Results cover 70M–300M full-stack models at TPP 10, 40, and 100. The horizontal axis reports additional non-embedding parameters relative to the identical-weight model. Higher TPP generally lowers BPB. Lower is better.
Figure 67: Adding independently parameterized layers can achieve lower WikiText-103 BPB but requires substantially more storage. Rows show 100M, 200M, 300M, and 500M models; columns show TPP 10, 40, and 100. The purple diamond is the matched doubled-depth Transformer. The final tick after the broken axis gives its true increase in non-embedding parameters. Lower BPB is better.
Figure 68: Low parameter overhead does not necessarily mean better training-compute efficiency. WikiText-103 BPB versus total training FLOPs for 70M–300M models at TPP 10, 40, and 100. Training FLOPs include repeated loop execution, loop-specific corrections, attention-score computation, and the realized token budget. Lower is better.
Figure 69: Specialization generally helps more as the number of loops increases. Dolma v1.5 BPB versus loop count for 70M–300M full-stack models at TPP 10, 40, and 100. Specialization is generally more useful at and than at . Increasing TPP also generally lowers BPB. Lower is better.
Figure 70: Specialization improves Dolma v1.5 BPB with relatively few additional parameters. Results cover 70M–300M full-stack models at TPP 10, 40, and 100. The horizontal axis reports additional non-embedding parameters relative to the identical-weight model. Higher TPP generally lowers BPB. Lower is better.
Figure 71: Adding independently parameterized layers can achieve lower Dolma v1.5 BPB but requires substantially more storage. Rows show 100M, 200M, 300M, and 500M models; columns show TPP 10, 40, and 100. The purple diamond is the matched doubled-depth Transformer. The final tick after the broken axis gives its true increase in non-embedding parameters. Lower BPB is better.
Figure 72: Low parameter overhead does not necessarily mean better training-compute efficiency. Dolma v1.5 BPB versus total training FLOPs for 70M–300M models at TPP 10, 40, and 100. Training FLOPs include repeated loop execution, loop-specific corrections, attention-score computation, and the realized token budget. Lower is better.
Figure 73: Specialization generally helps more as the number of loops increases. RedPajama BPB versus loop count for 70M–300M full-stack models at TPP 10, 40, and 100. Specialization is generally more useful at and than at . Increasing TPP also generally lowers BPB. Lower is better.
Figure 74: Specialization improves RedPajama BPB with relatively few additional parameters. Results cover 70M–300M full-stack models at TPP 10, 40, and 100. The horizontal axis reports additional non-embedding parameters relative to the identical-weight model. Higher TPP generally lowers BPB. Lower is better.
Figure 75: Adding independently parameterized layers can achieve lower RedPajama BPB but requires substantially more storage. Rows show 100M, 200M, 300M, and 500M models; columns show TPP 10, 40, and 100. The purple diamond is the matched doubled-depth Transformer. The final tick after the broken axis gives its true increase in non-embedding parameters. Lower BPB is better.
Figure 76: Low parameter overhead does not necessarily mean better training-compute efficiency. RedPajama BPB versus total training FLOPs for 70M–300M models at TPP 10, 40, and 100. Training FLOPs include repeated loop execution, loop-specific corrections, attention-score computation, and the realized token budget. Lower is better.
Figure 77: Specialization generally helps more as the number of loops increases. Dolma Programming Languages BPB versus loop count for 70M–300M full-stack models at TPP 10, 40, and 100. Specialization is generally more useful at and than at . Increasing TPP also generally lowers BPB. Lower is better.
Figure 78: Specialization can improve Dolma Programming Languages BPB with relatively few additional parameters. Results cover 70M–300M full-stack models at TPP 10, 40, and 100. The horizontal axis reports additional non-embedding parameters relative to the identical-weight model. Higher TPP generally lowers BPB. Lower is better.
Figure 79: Adding independently parameterized layers can achieve lower Dolma Programming Languages BPB but requires substantially more storage. Rows show 100M, 200M, 300M, and 500M models; columns show TPP 10, 40, and 100. The purple diamond is the matched doubled-depth Transformer. The final tick after the broken axis gives its true increase in non-embedding parameters. Lower BPB is better.
Figure 80: Low parameter overhead does not necessarily mean better training-compute efficiency. Dolma Programming Languages BPB versus total training FLOPs for 70M–300M models at TPP 10, 40, and 100. Training FLOPs include repeated loop execution, loop-specific corrections, attention-score computation, and the realized token budget. Lower is better.
Figure 81: Model scale and TPP reliably improve GSM8K likelihood, whereas the benefit of specialization depends on the configuration. Teacher-forced GSM8K BPB versus loop count for 70M–300M full-stack models. Rows vary model size and columns show TPP 10, 40, and 100. Specialization often helps at or , but the gain is not uniform. Lower BPB is better; these values do not measure final-answer accuracy.
Figure 82: The doubled-depth Transformer gives the lowest GSM8K BPB, while specialization uses far fewer additional parameters. Rows show 70M, 100M, 200M, and 300M models; columns show TPP 10, 40, and 100. The horizontal axis reports additional non-embedding parameters relative to the identical-weight model. Insets enlarge the specialization range, and the separated reference point denotes the doubled-depth Transformer. Lower BPB is better.
Figure 83: At comparable training FLOPs, doubling depth is stronger than adding and specializing more loops on GSM8K. Teacher-forced BPB versus total training FLOPs for 70M–300M models at TPP 10, 40, and 100. Training FLOPs include repeated loop execution, loop-specific corrections, attention-score computation, and the realized token budget. The doubled-depth Transformer is included as a reference. Lower BPB is better.
Figure 84: MATH likelihood improves with model scale, TPP, and, more modestly, loop-specific specialization. Teacher-forced MATH BPB versus loop count for 70M–300M full-stack models. Rows vary model size and columns show TPP 10, 40, and 100. Specialization is most consistently useful at and beyond the smallest scale. Lower BPB is better; these values do not measure final-answer accuracy.
Figure 85: Specialization improves MATH BPB with little added storage, while the doubled-depth Transformer remains strongest. Rows show 70M, 100M, 200M, and 300M models; columns show TPP 10, 40, and 100. The horizontal axis reports additional non-embedding parameters relative to the identical-weight model. Insets enlarge the specialization range, and the separated reference point denotes the doubled-depth Transformer. Lower BPB is better.
Figure 86: Doubling depth provides lower MATH BPB than adding and specializing loops at comparable training FLOPs. Teacher-forced BPB versus total training FLOPs for 70M–300M models at TPP 10, 40, and 100. Training FLOPs include repeated loop execution, loop-specific corrections, attention-score computation, and the realized token budget. The doubled-depth Transformer is included as a reference. Lower BPB is better.
Figure 87: Specialization usually improves HumanEval likelihood, but the best loop count and method depend on model size and TPP. Absolute teacher-forced BPB versus loop count for 70M–300M full-stack models. Rows vary model size and columns show TPP 10, 40, and 100. Lower BPB is better; these values do not measure execution accuracy or pass@ .
Figure 88: The doubled-depth Transformer achieves lower HumanEval BPB, while specialization improves identical weights with far less storage. Results cover 70M–300M full-stack models at TPP 10, 40, and 100. The horizontal axis reports additional non-embedding parameters relative to the identical-weight model. Insets enlarge the specialization range, and the separated reference point denotes the doubled-depth Transformer. Lower BPB is better.
Figure 89: At comparable training FLOPs, doubling depth gives lower HumanEval BPB than adding and specializing more loops. HumanEval BPB versus total training FLOPs for 70M–300M full-stack models at TPP 10, 40, and 100. Training FLOPs include repeated loop execution, loop-specific corrections, attention-score computation, and the realized token budget. The doubled-depth Transformer is included as a reference. Lower BPB is better.
Figure 90: Specialization improves MBPP likelihood in most settings, while the best loop count and method remain configuration-dependent. Absolute teacher-forced BPB versus loop count for 70M–300M full-stack models. Rows vary model size and columns show TPP 10, 40, and 100. Lower BPB is better; these values do not measure execution accuracy or pass@ .
Figure 91: The doubled-depth Transformer usually reaches lower MBPP BPB, while specialization is competitive at 70M and requires far less storage. Results cover 70M–300M full-stack models at TPP 10, 40, and 100. The horizontal axis reports additional non-embedding parameters relative to the identical-weight model. Insets enlarge the specialization range, and the separated reference point denotes the doubled-depth Transformer. Lower BPB is better.
Figure 92: Doubling depth gives lower MBPP BPB than adding and specializing loops at comparable training FLOPs. MBPP BPB versus total training FLOPs for 70M–300M full-stack models at TPP 10, 40, and 100. Training FLOPs include repeated loop execution, loop-specific corrections, attention-score computation, and the realized token budget. The doubled-depth Transformer is included as a reference. Lower BPB is better.
Figure 93: At , Full-Stack Looping is strongest and receives the clearest benefit from specialization. Results use 300M models at TPP 40. Each panel shows one evaluation corpus. Within each looping schedule, bars compare Loop Matrix, Loop LoRA, and Loop Softmax; stars mark the lowest BPB. Per-layer specialization is less consistently beneficial. Lower BPB is better.
Figure 94: At , specialization further strengthens full-stack looping but remains less useful for Per-Layer Looping. Results use 300M models at TPP 40. Each panel shows one evaluation corpus. Within each looping schedule, bars compare Loop Matrix, Loop LoRA, and Loop Softmax; stars mark the lowest BPB. Lower BPB is better.
Figure 95: At , Full-Stack Looping obtains the largest gains from specialization, while per-layer gains remain inconsistent. Results use 300M models at TPP 40. Each panel shows one evaluation corpus. Within each looping schedule, bars compare Loop Matrix, Loop LoRA, and Loop Softmax; stars mark the lowest BPB. Lower BPB is better.
Figure 96: Loop-specific specialization remains useful at 500M and 1B on DCLM. Full-stack results are shown against loop count for the completed TPP settings. The improvement is generally clearest at the larger loop counts, although the absolute differences are smaller than in the 70M–300M sweep. Blank or disconnected cells are unavailable and are not imputed. Lower BPB is better.
Figure 97: The larger-scale specialization gains require relatively few additional stored parameters on DCLM. BPB is plotted against additional non-embedding parameters relative to the identical-weight full-stack model. Lower BPB is better.
Figure 98: Specialization continues to improve Full-Stack Looping after accounting for its additional training cost on DCLM. Training FLOPs include repeated execution, loop-specific corrections, attention-score computation, and realized training tokens. Lower BPB is better.
Figure 99: Loop-specific specialization remains useful at 500M and 1B on C4 English. Full-stack results are shown against loop count for the completed TPP settings. The improvement is generally clearest at the larger loop counts, although the absolute differences are smaller than in the 70M–300M sweep. Blank or disconnected cells are unavailable and are not imputed. Lower BPB is better.
Figure 100: The larger-scale specialization gains require relatively few additional stored parameters on C4 English. BPB is plotted against additional non-embedding parameters relative to the identical-weight full-stack model. Lower BPB is better.
Figure 101: Specialization continues to improve Full-Stack Looping after accounting for its additional training cost on C4 English. Training FLOPs include repeated execution, loop-specific corrections, attention-score computation, and realized training tokens. Lower BPB is better.
Figure 102: Loop-specific specialization remains useful at 500M and 1B on WikiText-103. Full-stack results are shown against loop count for the completed TPP settings. The improvement is generally clearest at the larger loop counts, although the absolute differences are smaller than in the 70M–300M sweep. Blank or disconnected cells are unavailable and are not imputed. Lower BPB is better.
Figure 103: The larger-scale specialization gains require relatively few additional stored parameters on WikiText-103. BPB is plotted against additional non-embedding parameters relative to the identical-weight full-stack model. Lower BPB is better.
Figure 104: Specialization continues to improve Full-Stack Looping after accounting for its additional training cost on WikiText-103. Training FLOPs include repeated execution, loop-specific corrections, attention-score computation, and realized training tokens. Lower BPB is better.
Figure 105: Loop-specific specialization remains useful at 500M and 1B on Dolma v1.5. Full-stack results are shown against loop count for the completed TPP settings. The improvement is generally clearest at the larger loop counts, although the absolute differences are smaller than in the 70M–300M sweep. Blank or disconnected cells are unavailable and are not imputed. Lower BPB is better.
Figure 106: The larger-scale specialization gains require relatively few additional stored parameters on Dolma v1.5. BPB is plotted against additional non-embedding parameters relative to the identical-weight full-stack model. Lower BPB is better.
Figure 107: Specialization continues to improve Full-Stack Looping after accounting for its additional training cost on Dolma v1.5. Training FLOPs include repeated execution, loop-specific corrections, attention-score computation, and realized training tokens. Lower BPB is better.
Figure 108: Loop-specific specialization remains useful at 500M and 1B on RedPajama. Full-stack results are shown against loop count for the completed TPP settings. The improvement is generally clearest at the larger loop counts, although the absolute differences are smaller than in the 70M–300M sweep. Blank or disconnected cells are unavailable and are not imputed. Lower BPB is better.
Figure 109: The larger-scale specialization gains require relatively few additional stored parameters on RedPajama. BPB is plotted against additional non-embedding parameters relative to the identical-weight full-stack model. Lower BPB is better.
Figure 110: Specialization continues to improve Full-Stack Looping after accounting for its additional training cost on RedPajama. Training FLOPs include repeated execution, loop-specific corrections, attention-score computation, and realized training tokens. Lower BPB is better.
Figure 111: Loop-specific specialization remains useful at 500M and 1B on Dolma Programming Languages. Full-stack results are shown against loop count for the completed TPP settings. The improvement is generally clearest at the larger loop counts, although the absolute differences are smaller than in the 70M–300M sweep. Blank or disconnected cells are unavailable and are not imputed. Lower BPB is better.
Figure 112: The larger-scale specialization gains require relatively few additional stored parameters on Dolma Programming Languages. BPB is plotted against additional non-embedding parameters relative to the identical-weight full-stack model. Lower BPB is better.
Figure 113: Specialization continues to improve Full-Stack Looping after accounting for its additional training cost on Dolma Programming Languages. Training FLOPs include repeated execution, loop-specific corrections, attention-score computation, and realized training tokens. Lower BPB is better.
Figure 114: The shifted re-injection interface is the main source of PLT’s quality loss at 200M. DCLM BPB at TPP 40 (top) and TPP 80 (bottom) for . The left panels isolate the unnormalized PLT designs. The unshifted control closely tracks sequential Full-Stack Looping, whereas raw Shifted CLP can fail to match even Vanilla. The right panels show that separately normalizing and then controlling the mixture of the two inputs provides the largest recovery. Lower BPB is better.
Figure 115: The same design trade-off persists at 300M. DCLM BPB at TPP 40 (top) and TPP 80 (bottom) for . The normalized unshifted control is strongest but remains sequential. Among the shifted designs, a fixed 50/50 gate with a separate KV cache for every loop generally gives the strongest quality. KV sharing alone remains weak, while adding G-SWA closes much of the gap and can approach Full-Stack quality at larger loop counts. Lower BPB is better.
Figure 116: The scale difference changes with training budget and loop depth. Hidden-state RMS divided by token-embedding RMS for final 100M, checkpoints evaluated on the same DCLM sample. Each point is the ratio of the mean hidden-state RMS across tokens to the mean embedding RMS across tokens. The figure compares Full-Stack Looping, raw shifted PLT, and normalized shifted PLT. The raw shifted model uses neither normalization nor a gate before combining its inputs. The normalized reference uses the rank-8 linear gate defined in Subsec. G.3 . The -axis uses logarithmic spacing, but the tick labels show the underlying RMS ratios, not their logarithms.
Figure 117: KV sharing and G-SWA do not remove the scale difference. Hidden-state RMS divided by token-embedding RMS at TPP 40 and for final 100M and 200M checkpoints evaluated on the same DCLM token batch. The -axis uses logarithmic spacing, but the tick labels show the underlying RMS ratios, not their logarithms. Original PLT is the unnormalized design with KV sharing and G-SWA. For the normalized model, the plotted quantity is the hidden state produced by the Transformer. Normalization is applied only when this state is combined with the embedding at the next loop, so it does not force the plotted ratio to equal one. The normalized reference uses the rank-8 linear gate defined in Subsec. G.3 .
Figure 118: Gate choice has little effect once the inputs are normalized. DCLM BPB for fixed 50/50, scalar, rank-8 linear, and rank-8 MLP mixing at 100M and 200M. Rows show model size and columns show training budget; each panel compares loop counts. Insets report matched BPB differences from the normalized fixed 50/50 baseline, with negative values favoring the learned gate. The rank-8 linear gate improves over fixed 50/50 in most matched settings, but the differences are small and should be treated as directional because most settings use one training seed. Lower BPB is better.
Figure 119: The gate comparison remains close at larger model sizes. DCLM BPB for the same four mixing rules at 300M and 500M. Rows show model size and columns show training budget; insets report matched differences from the normalized fixed 50/50 baseline. The rank-8 linear gate is again generally favorable, but only by a small margin. Lower BPB and negative inset values are better.
Figure 120: Shifted looping requires the current-token embedding. DCLM BPB across training budgets for the four inputs in Subsec. G.4 . The current-token representation without the embedding closely matches the independent Full-Stack implementation, providing a check of the implementation. Adding the embedding improves both the previous-token and current-token variants. When the representation comes from the previous token, removing the embedding produces results worse than Vanilla (No Looping), and several runs diverge. The two conditions with both inputs use separate RMS normalization and the rank-8 linear gate. Missing or diverged results are not filled in. Lower BPB is better.
Figure 121: Normalized shifted PLT preserves the benefit of looping at 1B across six evaluation datasets. At , shifted PLT closely matches unshifted PLT and performs better than ordinary Full-Stack Looping, which in turn performs better than Vanilla (No Looping). Shifted PLT retains the parallel decoding schedule, while unshifted PLT and Full-Stack Looping remain sequential. The panels show BPB across TPP 10–160 on DCLM, C4 English, WikiText-103, Dolma v1.5, RedPajama, and Dolma programming languages. The same trained checkpoints are evaluated on every dataset, and each panel uses its own y-axis. Lower BPB is better.
Figure 122: DCLM quality across loop counts at 100M and 200M. The top row shows 100M and the bottom row shows 200M. Each panel fixes TPP and varies . The comparison includes Vanilla (No Looping), ordinary Full-Stack Looping, shifted PLT, and unshifted PLT where available. The relative ordering varies more at these smaller model sizes. Lower BPB is better.
Figure 123: DCLM quality across loop counts at 300M and 500M. The top row shows 300M and the bottom row shows 500M. Each panel fixes TPP and varies . Shifted PLT is more consistently competitive with ordinary Full-Stack Looping at these larger model sizes, while unshifted PLT provides the stronger sequential reference. Lower BPB is better.
Figure 124: DCLM quality across training budgets at 100M and 200M. The top row shows 100M and the bottom row shows 200M. Each panel fixes the loop count and varies TPP from 10 to 160 for Full-Stack Looping, shifted PLT, and unshifted PLT. Unshifted PLT achieves the lowest BPB, while shifted PLT usually matches or improves upon Full-Stack Looping. Lower BPB is better.
Figure 125: DCLM quality across training budgets at 300M and 500M. The top row shows 300M and the bottom row shows 500M. Each panel fixes the loop count and varies TPP from 10 to 160 for Full-Stack Looping, shifted PLT, and unshifted PLT. At , shifted PLT closely tracks unshifted PLT across the evaluated TPP range; at , a clearer gap remains. Shifted PLT usually matches or improves upon Full-Stack Looping while retaining its parallel decoding schedule. Lower BPB is better.
Figure 126: Stored parameters and training FLOPs favor different loop counts. DCLM BPB versus stored non-embedding parameters (top) and total training FLOPs (bottom). Columns fix TPP, and marker shape denotes loop count. Looping can improve quality with little additional parameter storage. When training FLOPs are the constraint and additional parameters can be stored, models made deeper with distinct Transformer layers generally give the stronger quality–compute trade-off. Lower BPB is better.
Figure 127: Measured quality versus training FLOPs for shifted PLT at larger model sizes. DCLM BPB for Vanilla (No Looping) and shifted PLT at 300M, 500M, and 1B. Panels fix or , and points along each curve increase from TPP 10 to 160. With , the curves cross at some measured operating points: shifted PLT trained on fewer tokens can outperform Vanilla trained on more tokens. With , the additional loop computation generally gives a weaker quality–compute trade-off. We do not extrapolate beyond the measured points. Lower BPB is better.
Figure 128: Unshifted PLT provides a stronger quality reference but not parallel decoding. DCLM BPB versus total training FLOPs for Vanilla (No Looping) and unshifted PLT at 300M, 500M, and 1B. Unshifted PLT passes each token’s representation directly from one loop to the next, so its loops remain sequential during decoding. Lower BPB is better.
Figure 129: DCLM quality versus total training FLOPs at 100M and 200M. The left plot shows 100M and the right plot shows 200M. Insets report BPB differences from the corresponding Full-Stack model. Lower BPB and negative inset values are better.
Figure 130: DCLM quality versus total training FLOPs at 300M and 500M. The left plot shows 300M and the right plot shows 500M. Insets report BPB differences from the corresponding Full-Stack model. Lower BPB and negative inset values are better.
Figure 131: DCLM quality versus total training FLOPs at 1B. The and panels show all five measured training budgets. Insets report BPB differences from the corresponding Full-Stack model. Lower BPB and negative inset values are better.
Figure 132: Loop-specific specialization preserves the ordering among ordinary and parallel looping designs. DCLM BPB for 300M and 500M models. The figure compares ordinary Full-Stack Looping, shifted PLT, and unshifted PLT with shared weights and with Loop Matrix. Unshifted PLT gives the strongest quality but remains sequential during decoding. Shifted PLT generally improves over Full-Stack Looping while retaining parallel decoding. Lower BPB is better.
Figure 133: Parallel looping with Loop Matrix on C4 English. Absolute BPB for ordinary Full-Stack Looping, shifted PLT, and unshifted PLT with shared weights and with Loop Matrix. The comparison uses 300M and 500M models, TPP 10–80, and . The ordering broadly matches DCLM: unshifted PLT is strongest, while shifted PLT generally improves over Full-Stack Looping and retains parallel decoding. Lower BPB is better.
Figure 134: Parallel looping with Loop Matrix on WikiText-103 and Dolma v1.5. The top figure shows WikiText-103 and the bottom figure shows Dolma v1.5. Both compare ordinary Full-Stack Looping, shifted PLT, and unshifted PLT with shared weights and with Loop Matrix. Specialization is generally more useful at than at . Lower BPB is better.
Figure 135: Parallel looping with Loop Matrix on RedPajama and Dolma programming languages. The top figure shows RedPajama and the bottom figure shows the code-heavy Dolma programming-languages split. The benefit of specialization is generally clearer at ; at , the differences are smaller and less uniform, particularly on the programming-languages split. Lower BPB is better.
More in this topic
Recurrent Transformers
Recurrent Transformers reuse layers or state across steps to extend context, reasoning depth, and sequence modeling with less parameter growth.
See all 183 papers