Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA
Authors: Sangmin Bae
Adam Fisch
Hrayr Harutyunyan
Ziwei Ji
Seungyeon Kim
Tal Schuster
Venue: ICLR 2025 Poster (2025)
Keywords: Efficient LLM
Parameter Sharing
Relaxed Recursive Transformer
Continuous Depth-wise Batching
Early-Exiting
Primary area: applications to computer vision, audio, language, and other modalities

Abstract: Large language models (LLMs) are expensive to deploy. Parameter sharing offers a possible path towards reducing their size and cost, but its effectiveness in modern LLMs remains fairly limited. In this work, we revisit "layer tying" as form of parameter sharing in Transformers, and introduce novel methods for converting existing LLMs into smaller "Recursive Transformers" that share parameters across layers, with minimal loss of performance. Here, our Recursive Transformers are efficiently initialized from standard pretrained Transformers, but only use a single block of unique layers that is then repeated multiple times in a loop. We further improve  performance by introducing Relaxed Recursive Transformers that add flexibility to the layer tying constraint via depth-wise low-rank adaptation (LoRA) modules, yet still preserve the compactness of the overall model. We show that our recursive models (e.g., recursive Gemma 1B) outperform both similar-sized vanilla pretrained models (such as TinyLlama 1.1B and Pythia 1B) and knowledge distillation baselines---and can even recover most of the performance of the original "full-size" model (e.g., Gemma 2B with no shared parameters). Finally, we propose Continuous Depth-wise Batching, a promising new inference paradigm enabled by the Recursive Transformer when paired with early exiting. In a theoretical analysis, we show  that this has the potential to lead to significant (2-3$\times$) gains in inference throughput.

Published as a conference paper at ICLR 2025
RELAXED RECURSIVE TRANSFORMERS: EFFECTIVE PARAMETER SHARING WITH LAYER-WIS LORA
Sangmin Bae', Adam Fisch², H Harutyunyan', Zi Seungyeon Kim², Tal Schuster))
'KAIST AI Google DeepMind *Google Research
2smn0223@kaise ac kr, talschuster com
ABSTRACT
Large language models (LLMs) are expensive to deploy. Parameter sharing offers possible path towards reducing their size and cost, but its effectiveness in modern LLMs remains fairly limited. In this work, we revisit "layer tying" as form of param- eter sharing in Transformers, and introduce novel methods for converting existing LLMs into smaller "Recursive Transformers" that share parameters across layers, with minimal loss of performance. Here, our Recursive Transformers are efficiently initialized from standard pretrained Transformers, but only use single block of unique layers that is then repeated multiple times loop. We further improve per- formance by introducing Relaxed Recursive Transformers that add flexibility to the layer rying constraint via depth-wise low-rank adaptation (LoRA) modules, yet still preserve the compactness of the overall model. We show that our recursive models (e.g., recursive Gemma IB) outperform both similar-sized vanilla pretrained models (such as TinyLlama 1.1B and Pythia IB) and knowledge distillation baselines- and can even recover most of the performance of the original "full-size" model (e.g., Gemma 2B with shared parameters). Finally, we propose Continuous Depth- wise Batching, promising new inference paradigm enabled by the Recursive Transformer when paired with early exiting. In a theoretical analysis, we show that this has the potential to lead to significant (2-3x) gains in inference throughput.
INTRODUCTION
Efficient deployment of large language models (LLMs) demands a balance between performance and resources (Raposo al., 2024; Leviathan al., 2023; Rivière al., 2024; Wan et al., 2024; Zhou et al., 2024). While larger models with more parameters consistently demonstrate superior performance (Rosenfeld et al., 2020; Rae et al., 2021; Hoffmann et al., 2022), their substantial memory and computational demands are expensive (Pope et al., 2023). Parameter sharing approaches (e.g. Dehghani et al., 2019; Xia et al., 2019; Lan et al., 2020; Takase & Kiyono, 2023), wherein weights are reused across model layers, can lower these costs by reducing memory footprint, and thereby allow for the use of fewer (or lower-grade) accelerators, or larger batch sizes for better throughput. While parameter sharing has shown encouraging capabilities in previous work (Lan et al., 2020: Giannou et al., 2023), its application to modern LLMs has yielded limited reported success.
In this work, we revisit parameter sharing for LLMs, and propose novel methodologies convert existing, unshared models into smaller, and more efficient, Recursive Transformers. These models use single block of unique layers that are recursively reused across multiple loops, yet still achieve im- pressive performance relative to their reduced size. To mitigate the potential performance degradation associated with parameter sharing, first initialize the shared block of layers based on the original model's 's pre-trained parameters, and then finetune the resulting recursive model for limited number of "uptraining" steps. Importantly, we show that our initialization strategies allow us achieve strong performance with minimal training time. This is aligned with observations that model compression techniques such as layer skipping (Zhang et al., 2024a; Zeng al., 2023; Fan et al., 2020; Elhoushi et al., 2024), pruning (Frankle & Carbin, 2019; Ramanujan et al., 2020) or nesting (Devvrit et al., 2023) can preserve surprisingly high performance furthee motivating our approach of compressing models to more compact yet performant architectures (here, repeated layers with low-rank adapters).
Work done during an intersship Google DeepMind. 'Corresponding author.

---

Published as a conference paper at ICLR 2025
(N/K)
X (N/K)
Layer N
Layer
Layer
Layer
(1)
(2)
[N/K]
convert
Layer 2
Layer
Layer
Layer
Layer 1
Layer 1
Vanilla Transformer
Recursive Transformer
Relaxed Recursive Transformer
Overview of the conversion from a vanilla N-layer Transformer to a Recursive Transformer with N/K blocks of K shared layers. The Recursive Transformer obtained by repeating a single block of K layers multiple times, resulting in a looped architecture. The Recursive Transformer can also be converted into a Relaxed Recursive Transformer by adding layer-specific LoRA modules. This preserves many of the advantages of weight sharing, but also allows for better performance.
Figure I:
As depicted in Figure 1, we further propose the Relaxed Recursive Transformer, extension of the Recursive Transformer in which the weight tying across repeated layer blocks is slightly relaxed through the incorporation of multiple layer-specific, low-rank adaptation (LoRA) modules (Hu et al., 2022). Despite its simplicity, this strategy offers several non-trivial advantages. First, it allows for low- rank deltas between shared layers, while only adding minimal overhead. Second, the rank of the LoRA matrices can be adjusted to control the degree of relaxation, which directly influences model capacity. Furthermore, since the relaxed model has the same overall shape as the original Transformer, we can efficiently initialize LoRA modules via truncated Singular Value Decomposition (Hansen, 1987) on the residual matrices hetween the original layer weights and the shared layer weights. Hence, the rank values serve as a pivotal hyperparameter, enabling the Relaxed Recursive Transformer to seamlessly transition between the two extremes of the vanilla and Recursive Transformer architectures.
While the primary focus of this paper lies in how to formulate and train Recursive Transformers, we also highlight their potential to achieve significant throughput gains via a new batched inference paradigm, Continuous Depth-wise continuous sequence-win wise batching (Yu et al., 2022; Kwon et al., 2023), which leverages the fact that the computation performed compute new token is functionally the same (and uses the same model parameters) regardless of the token position within the sequence. This allows new requests to be continuously scheduled when slots within batch become available. For example, when one response is completed, the start of the next response to be formed can immediately take the finished response's ' place in the batch, without waiting for the rest of the batch responses that might be longer. In our Recursive Transformer, parameter sharing occurs not only across different timesteps, but also across different depths (loop iterations). This enables an extra dimension of dynamic grouping: jointly com- puting different iterations of the looped layer blocks per individual responses within the same batch.
wise Batching, that their recursive nature enables. Prior work introduced
Our key contributions are as follows:
We introduce framework for initializing and training Relaxed Recursive Transformers and demon- strate strong performance compared to non-recursive models of comparable size. For example, when we uptrained recursive Gemma IB model converted from a pretrained Gemma 2B (Team et al., 2024), we observed up to 13.5 absolute accuracy improvement (22% error reduction) on few-shot tasks compared to a non-recursive Gemma 1B model (pretrained from scratch). Further- more, we show that by incorporating knowledge distillation (Hinton et al., 2015; Kim & Rush, 2016), our recursive Gemma model, uptrained on 60 billion tokens, achieves performance on par with the full-size Gemma model trained on a massive 3 trillion token corpus (see $3.3 for details).
Based on our Relaxed Recursive Transformer, we also evaluate key use case for continuous wise batching with early-exiting (Bae et al., 2023; Schuster et al., 2022; Elbayad et al., 2020; Graves, 2016a), which opportunistically makes predictions for samples with high confidence at earlier stages. From our simulation, Early Exits reveal a substantial throughput improvement of up 2-3x compared to vanilla Transformer with the same architecture. Notably, the recursive Gemma model, which outperforms the vanilla Pythia model, can theoretically achieve a nearly 4x increase in throughput (see $3.8 for details).
2

---

Published as a conference paper at ICLR 2025
x
X
x
Layer
Layer
Layer 6
Layer 3,
Layer 3
Layer 3,
Layer
Layer 3
Layer 2,
Layer 2
Layer 2,5
Layer
Layer 1
Layer 1,
Layer 1
Layer
Layer
Stepwise
Average
Lower
T-SVD (Layer 1 Layer 1, (V,, =
Layer
Three proposed methodologies for initializing looped layers Recursive Transformer. Each layer number indicates the source layer in the full-size model used for initialization. Right: Example of a Relaxed Recursive Transformer initialized by SVD method. Here, looped layers are initialized using Average method.
Figure 2: Left:
An example of unshared, full-size model with 6 layers. Middle:
EFFECTIVE MODEL COMPRESSION WITH RECURSIVE PATTERNS
In this section, we present the main details our method for converting a vanilla Transformer model into parameter-shared model that outperforms models of equivalent size. We first provide short overview of the Transformer architecture ($2.1). Then, we introduce the Recursive Transformer and present effective techniques to initialize its looped layers by leveraging the weights of the original pretrained model ($2.2). In $2.3, we relax the parameter-sharing constraint in the model design, and add limited of layer-specific parameters to further improve the model's accuracy while maintaining compact representations. Finally, we show how, beyond reduced memory, Recursive Transformers readily support further throughput optimizations via a novel inference paradigm ($2.4).
2.1 BASIC TRANSFORMER ARCHITECTURE
Large language models (Rivière et al., 2024; Reid et al., 2024; OpenAL, 2023; Dubey et al., 2024) typically leverage the Transformer architecture (Vaswani et al., 2017). A Transformer consists layers, where the hidden states at each time step are computed by running through the series of layers:
h = fbb--;; Φ2). L LE L].
(1)
with hς representing the embedding of the token 1/-1 from the previous time step, and t denoting the trainable parameters of the f-th layer. Each layer has two core components: multi-head attention (MHA) mechanism and feed-forward network (FFN), MHA employs multiple attention heads to capture diverse relationships within the input sequence via linear attention weights and scaled dot-product attention mechanisms. The FFN structure typically consists of two linear transformations, but different models exhibits distinct structural variations. See Appendix for further details.
2.2 RECURSIVE TRANSFORMER: LOOPED LAYER TYING
In this work, revisit parameter sharing in the context of LLMs and propose the Recursive Transformer architecture. Among various looping strategies (refer to Appendix D), we specifically adopt the CYCLE strategy (Takase & Kiyono, 2023) for Recursive Transformers, wherein a single block of unique layers is recursively reused. This inherent design aligns seamlessly with early-exiting mechanisms, potentially offering substantial speedup. The model's hidden states are computed as:
h = fhh l((e-1) mod L/B)+1), i LI,
(2)
where the parameter-shared model is parameterized by &', and B denotes the number of looping blocks (we restrict B be a factor of L). For example, Gemma 2B (Team al., 2024) with 18 layers be converted to a recursive variant with 2 blocks by storing weights for only the first 9 layers. The forward pass will loop twice through these 9 layers. We tie all trainable parameters. including the weights of the linear layers in the Transformer blocks and the weights of the RMSNorm (Zhang & Sennrich, 2019).
Initializatinn techniques for looped layers To mitigate the potential performance drop associated with reduced capacity in parameter- shared models, we propose several novel initialization method- ologies to facilitate effective knowledge transfer from unshared, pretrained models to Recursive
3

---

Published as a conference paper at ICLR 2025
Transformers. Figure illustrates three such techniques. The Stepwise method selects intermediate layers at specific intervals while keeping the first and last layer fixed. This is motivated by prior work (Liu et al., 2023; Zhang et al., 2024a; Zeng al., 2023; Fan et al., 2020) showing minimal impact on generation quality when skipping a few layers in L.MMs. The Average method initializes the shared weights among tied layers by averaging their weight matrices, whereas the Lower method directly uses weights from the first K layers of the unshared model. We conducted a brief uptraining on 15 billion tokens to investigate the extent of performance recovery in these initialized models ($3.4) and found the Stepwise approach to perform best for Recursive Transformers. However, we found the Average method to perform best for Relaxed Recursive Transformers, discussed next.
2.3 RELAXED RECURSIVE TRANSFORMER: MULTI-LORA LAYERS
While full layer-tying is effective for compressing the model's s size while maintaining strong capa- bilities, it has two noticeable limitations: (1) the set of possible model sizes limited scaling the number of layers, and (2) each model layer ends up having to serve multiple roles associated with different depths of the model. To address this, we introduce Relaxed Recursive Transformers in which we incorporate independent adapter modules (Hu et al., 2022; Houlsby et al., 2019) for each layer, relaxing the strict parameter sharing. While we experiment with various approaches like layer-specific prefixes (Liu et al., 2021) (see Appendix K), we find low-rank adaptation (LoRA) modules (Hu et al., 2022) to efficiently capture the subtle variations between tied layers. Specifically, we modify Eq. 2 to:
h = fb! (((6-1) mod L/B)+1 A0'), LE [1,
where A is the (small) set of parameters for the LoRA modules.
In this relaxed model, each looped layer augmented with multiple LoRA modules. For example, recursive model with two loop iterations has single block of shared layers, and two different LoRA modules are attached to each layer within this block. The first and second LoRA modules per layer are used during the first and second loop iterations, respectively. Functionally, these LoRA modules introduce low-rank deltas to all of the shared, linear weight matrices. More concretely, for a base transformation h = W'x, our modified forward pass yields h W'x + AW'x = W'x + BAx, where A E R (rx*)
and B E gidxr) denote the weight matrices of LoRA with rank r.
LoRA initialization via truncated SVD Unlike typical LoRA finetuning setups that train only the LoRA parameters, here we train all model parameters to let the shared parameters learn an optimal centroid for all of the layer depths that they support. Therefore, instead of following standard zero initialization for adaptation to the frozen base model, we propose novel initialization methods, especially designed for Relaxed Recursive Transformers. To effectively match the performance of the original full-size model after initializing the tied weights as described in $2.2, we aim for the sum of the tied weights (t') and LoRA weights (Ad') to approximately recover the full-size model's weights (b) We exploit truncated Singular Value Decomposition (SVD) (Hansen, 1987) on residual matrices between original weights and tied weights:
(4) where outputs retain the first columns corresponding to the largest singular values. W denotes the weight matrices of the full-size model, and W' denotes those of the Recursive Transformer. We initialize the LoRA's weights with principal components in Eq. 4: B as the product of U, and Σ,, and A as the transpose of the right singular vectors V, (see Figure 2). With sufficiently large ranks, our Relaxed Recursive Transformer (Eq. 3) approximates the full-size vanilla model (Eq. I):
U D V Truncated SVD(W, - W' ((É-1) mod L/B)+15 r), E [1,
Wx W'x + (U,E,)(V,) W'x + BAx = AW'x,
(5)
Mcanwhile, setting the rank to zero reduces the model to a Recursive Transformer, as the LoRA modules contribute additional parameters, highlighting the flexibility of this relaxation approach.
2.4 CONTINUOUS DEPTH-WISE BATCHING AND EARLY-
In real-world deployments, user requests arrive sequentially and asynchronously. Recent research has introduced continuous sequence-win batching (Yu et al., 2022; Kwon et al., 2023), serving strategy that allows new requests to immediately replace completed (terminated) sequence within a batch. This approach exploits the fact that the computation performed for new token is functionally
4

---

Published as a conference paper at ICLR 2025
Time Step
Time Step
Time Step
1
2
3
4
5
6
9
1
2
3
4
5
8
9
2
4
5
6
7
8
9
Model Stage
3
16
32
16
16
8
2
16
32
16
16
16
16
16
8
1
16
32
16
16
16
8
16
16
16
32
8
16
32
8
16
32
Input Stream
(a) Vanilla Batching
(b) Depth-wise Batching
(c) With Early-exiting
Figure 3: An illustrative example of continuous depth-wise batching strategy together with early- exiting. We assume a maximum batch size of 32, three model "stages" (e.g.. layer blocks), and a stream of batched inputs that arrive sequentially in time. In (a), all three model stages must complete for the first (non-maximal) batch of 16 before the second batch of 32 examples that arrives next can be started. In (b), however, half of second batch of 32 examples can share computation with the first batch of 16 that is still finishing. Finally, (c) demonstrates a situation where some examples within each batch can early-exit after stage 2; their vacant slots in the batch are then immediately filled.
the same and utilize the same model parameters. By continuously scheduling requests in this manner, models can operate their maximum batch capacity, thereby enhancing serving efficiency.
The repetitive structure of Recursive Transformers allows for the same function to be applied not just across sequences, but also across depths (loop iterations). This introduces new dimension for continuous batching. which we call Continuous Depth-wise Batching. This technique enables the simultaneous computation of different iterations of the looped layer block for different samples (See Figure 3 for an example with a single forward pass; this easily extends to multiple decode iterations per request.) With maximum batch size of 32, standard Transformer must wait for all model stages to complete before processing new requests. In contrast, our Recursive Transformer, because it shares layer functions across all stages, can immediately schedule new incoming requests at timestep 2, maximizing batch size utilization. This strategy can yield a substantial speedup in generation and reduce the time to first token (Fu et al., 2024; Míao et al., 2023) through faster scheduling.
Throughput improvements from depth-wise batching are further amplified when combined with early- exiting (Bae et al., 2023; Schuster et al., 2022; Elbayad al., 2020). As depicted in Figure 3c, once some samples exit after certain looping iterations, queued requests can then be immediately scheduled. While Recursive Transformers leverage the speedup from carly-cxiting, they also inherently address key challenge of batched inference in early-exiting approaches: the synchronization issue when serving large batches, as early-exited tokens might wait for others to complete processing through the entire model. We demonstrate that Recursive Transformers, equipped with this dynamic sample scheduling at various depths, can theoretically allow up to 2-3 x speedup on evaluated LLMs.
EXPERIMENTS
3.1 EXPERIMENTAL SETUP
We evaluate our method on three popular pretrained LLMs: Gemma 2B (Team et al., 2024), TinyL- lama I.IB (Zhang et al., 2024b), and Pythia IB (Biderman et al., 2023). Table 2 summarizes each model's 's architecture and pretraining recipes, and their few-shot performance is summarized in Ap- pendix F. After converting to Recursive Transformers, we uptrained models on the SlimPajama dataset (Soboleva et al., 2023). We used the Language Model Evaluation Harness framework (Gao et al., 2023) to evaluate accuracy on seven few-shot tasks, and averaged them for performance comparison. Detailed experimental setup for uptraining or evaluation can be found in Appendix G.
3.2 NON-RECURSIVE MODEL BASELINES
Full-size model Our ultimate goal is for the Recursive Transformer to achieve performance compa- rable to the original, full-size pretrained model, without much uptraining. However, we observed that the distribution divergence between the pretraining and uptraining datasets can hinder achieving the desired performance. In particular, uptraining on new datasets, particularly those of comparatively lower quality, sometimes led to performance degradation certain benchmarks. Table 4 summarizes
5

---

Published as a conference paper at ICLR 2025
Few-shot Accuracy
Few-shoff Accuracy
46
Few-shot Accuraoy
55
44-
50
45
Reducad-size
Full-size
42-
sull-size
Reduced-size Recursive Rolaxed
Full-sze Reduced-size
Recursive Relaxed
Recursive
Relaxod
1.08
2.08
0.68
0.88
1.08
0.4B
0.68
0.8B
Model Sizes
Model Sizes
Model Sizes
(a) Gemma
(b) TinyLlama
(c) Pythia
Figure Recursive and Relaxed Recursive Transformers achieve comparable performance to full-size models, and significantly outperform reduced-size models. Recursive models were initialized using the Stepwise method, while relaxed models utilized Average and SVD methods for looped layers and LoRA modules. We show the performance of four different rank values: 64, 128, 256, and 512. Recursive and reduced-size models were either uptrained (recursive model) and pretrained from scratch (reduced-sio model) on 60 billion tokens using knowledge distillation objective.
the evaluation results of full-size models based on the number of uptraining tokens. For instance, in the case of Gemma, where the pretraining dataset unrelcased but potentially well-curated (Team al., 2024), all few-shot performance metrics gradually decreased after uptraining the SlimPajama dataset. This suggests that the achievable upper bound performance with the SlimPajama dataset might be considerably lower than the original model performance. Therefore, set the target performance for Gemma and Pythia models as the performance achieved by uptraining a full-size pretrained model with an equivalent number of tokens. Since Tiaya was already pretrained on SlimPajama-whie is the same dataset we use for uptraining (eliminating any distribution shift)-for slightly longer than our runs, we use the performance of the original checkpoint as reference.
Reduced-size model To demonstrate the performance advantages of Recursive Transformers compared to models with an equivalent number of parameters, we introduce another baseline: reduced-size models. These models have either half or one-third the parameters of their full-sized counterparts, matching the parameter count of our recursive models. However, these reduced models pretrained from scratch on the same training recipe (number of training tokens and distillation from full-size model), but without the benefits the pretrained weights and the looping mechanism. This comparison serves to highlight the efficacy of our initialization techniques and the recursive function itself attaining strong performance, even with a constrained model size.
3.3
MAIN RESULTS
Figure 4 presents the few-shot performance of Recursive Transformers with two blocks and their relaxed variants. Recursive Transformers, even without relaxation, demonstrate remarkably high performance despite having only half the parameters of the full-size model. The Gemma model achieved a 10%p performance gain compared to the reduced-size model, which was also trained on 60 billion tokens using distillation loss. Remarkably, the recursive TinyLlama model even surpassed the vanilla model's performance, even though the latter was pretrained on a larger corpus of 105 billion tokens. Our initialization techniques proved highly effective in achieving this superior result, along with the benefit of the uptraining dataset (SlimPajama) being the same as its pretraining dataset.
The relaxed models effectively interpolate between the full-size model and the Recursive Transformer, depending on the LoRA rank. As the model size increases with larger LoRA modules, SVD initialization methods allow for a more precise approximation of full-rank matrices, resulting in improved performance. Notably, the relaxed Gemma model with a rank of 512 achieves performance on par with the original model pretrained on 3 trillion tokens (58.44 558.4% vs. 58.6%). despite using fewer parameters and uptraining only 60 billion tokens. This trade-off provides flexibility in selecting the best configuration for various deployment scenarios. We believe that additional uptraining and higher-quality datasets could yield better performance with even
streamlined models.
In the subsequent sections, we provide comprehensive overview of extensive ablation studies conducted prior to achieving this final performance. In $3.4, delve into the analysis of various initialization methodologies for Recursive Transformers. Insights into the relaxation model are detailed in $3.5. Finally, explore enhanced training strategies like knowledge distillation ($3.6).

---

Published as a conference paper at ICLR 2025
Full-sizze
Slepwise Average Lower
60
Stepwise Average
Few si--shot Accuracy
Few-sho: shot Accuracy
Training Loss
Lower Random
50-
Random
800
Gemma
40
TinyLiama Pythia
size Reduced-- size
2.6
30-
Recursive (B2)
Recursive (B3)
2.2
08
5B
10B
15B
20
1.0B
1.58
of Tokens
Gemma
TinyLlama
Pythia
Model Sizes
(a) Loss curves for Gemma
(b) Average few-shot performance
(c) Recursive model performance
Figure 5: (a) Among the proposed methods, the Stepwise method obtains the lowest training loss on the SlimPajama dataset. (b) The Stepwise method consistently demonstrate the highest average few-shot accuracy across three architectures. (c) Recursive Transformers initialized with the Stepwise method demonstrated significant performance gains compared to non-recursive model baselines.
3.4 INITIALIZATION TECHNIQUES FOR LOOPED LAYERS
Stepwise initialization serves as the best initial point for Recursive Transformers We present the training loss of Gemma models initialized using three different methods in Figure 5a, and their few-shot performance in Figure 5b. Our proposed methods significantly outperformed random initial- ization, which simply adds recursion to reduced-size model, suggesting that leveraging pretrained weights any manner is beneficial for performance boost. Moreover, the Stepwise methodology consistently demonstrated best performance, aligning with insights that LLMs can preserve perfor- mance even with a few layers skipped (Raposo et al., 2024; Zhang et al., 2024a; Elhoushi et al., 2024). Interestingly, summarized in Table 5, the recursive TinyL.ama model, uptrained only 15 billion tokens, yields few-shot performance comparable to the original model pretrained on 105 billion tokens. This suggests that with sufficient training, performance of full-size pretrained model (Dehghani et al., 2019; Takase & Kiyono, 2023).
recursive architecture can match the
Recursive Gemma 1B outperforms both pretrained TinyLlama 1.1B and Pythia 1B The looped Gemma IB model, utilizing our proposed Stepwise method, outperformed reduced-size baselines with equivalent parameter counts by up 13.5 percentage points (51.7% vs. 38.2%). Furthermore, it even outperformed the full-size TinyLlama 1.IB and Pythia 1B models (see Figure 5c). This is noteworthy achievement given that Pythia was pretrained on 300 billion tokens, whereas the recursive Gemma was uptrained on only 15 billion tokens. Consequently, high-performing LLMs serve as promising starting point, as their recursive counterparts readily outperform other ordinary vanilla models of similar size. Further details can be found in Appendix I.
Takeaways for the Recursive Transformer
We find that converting well- vell-pretrained models into Recursive Transformers leads to high- performing models with minimal uptraining. Notably, initializing looped layers via the Stepwise method yields the best results. With just 15 billion tokens of uptraining, recursive Gemma model outperforms even the full-size pretrained TinyLlama and Pythia models.
3.5 RELAXATION OF STRICT PARAMETER SHARING VIA LORA MODULES
Average initialization for looped layers optimally suits Relaxed Recursive Transformer Fig- ures 6a and 6b illustrate the effect of relaxing parameter sharing via layer-wise LoRA modules. No- tably, initializing tied layers in relaxed models with Average method yielded substantial performance improvements, even outperforming the non-relaxed model initialized with Stepwise. Approximating residual matrices between averaged weights and their individual weights appears readily achievable using truncated SVD with low ranks. In contrast, we observed an intriguing phenomenon where our models initialized with Stepwise occasionally showed performance degradation after relaxation. This likely because capturing the nuances between entirely distinct layer weights challenging with an insufficient rank, leading suboptimal solution. Further details provided in Appendix J.
7

---

Published as a conference paper at ICLR 2025
LORA LoRA
LaRA LaRA x
LoRA o LoRA
Zero Init SVD Init
Few-shat Accuracy
Few-shot Accuracy
Training Loss
Stepwise Average Lower
Slepwise Average Lower
56
50-
2.7
50
45
40
2.3
0B
5B
10B
15B
40
# of Tokens
Gemma
TinyLlama
Pythia
Stepwise
Average
Lower
(a) Loss changes in Gemma model
(b) Accuracy gains from relaxation
(c) Effects of SVD initialization
Figure 6: The Relaxed Recursive Transformer, with its looped layer initialized using Average method, achieved the best performance in terms of both (a) training loss and (b) few-shot accuracy. The models utilize two blocks, with the LoRA modules initialized using the SVD method rank of 512. (c) SVD initialization method significantly enhanced performance compared to zero initialization.
SVD initialization to approximate pretrained weights outperforms zero initialization LoRA modules initialized with zero values guarantee that the model begins training from the same point as the non-relaxed model. Conversely, SVD initialization positions the model closer to either the full-size model (with full-rank) or the non-relaxed model (with small rank). To emphasize the effectiveness of initializing near full-size model weights, we compared these two methods at a moderately large rank of 512, shown in Figure 6c. Our proposed SVD) strategy demonstrated impressive performance boost of up to 6.5 points, facilitating faster convergence by updating the principal low-rank matrices (aligned with findings in Meng al. (2024)). For results across other architectures, refer to Figure 15,
Higher rank enhances recovery of original pretrained weights At full rank, relaxed models can perfectly match full-size pretrained models. Consequently, as illustrated in Figure 7a, performance generally improves with increasing rank, resulting in a clear Pareto frontier between model size and performance. However, only Stepwise initialization showed U-shaped performance trend: middle- range rank resulted in poor approximation, whereas very low ranks (akin to random initialization for LoRA modules) yielded better performance. The overall results are summarized in Table
Takeaways for the Relaxed Recursive Transformer
Adjusting the LoRA rank in the Relaxed Recursive Transformer, together with our SVD-based initialization technique, allows for smoother trade-off between a fully weight-tied recursive model and vanilla model. Furthermore, find that initializing the shared weights in the looped layers with the Average method leads to the best performance in this setting.
3.6 EXTENDED UPTRAINING AND KNOWLEDGE DISTILLATION
We further enhanced the performance of our low-rank models by introducing two techniques: up- training on an extended corpus and knowledge distillation from the full-sized model. Specifically, we increased the number of uptraining tokens from 0.5% to 2% of the total 3 trillion tokens used for pretraining Gemma models, resulting in a total of 60 billion tokens. Additionally, we regularized the losses using forward Kullback-L Leibler divergence (Hinton et al., 2015; Kim & Rush, 2016), which exhibited the best performance gains among the examined distillation losses. Table 11 summarizes the results of various ablation studies conducted investigate the impact of these two techniques.
The combined effect of these techniques is presented in Figure 7b, demonstrating an improvement of up to 4.1 percentage points few-shot accuracy compared the previous 15 billion token uptraining results. Notably, the relaxed Gemma model with rank of 512 nearly matched the performance of the full-size model. We also expect that further performance gains can be achieved with much lighter recursive model by utilizing superior teacher model or conducting more extensive training on high-quality data. Figure 7c illustrates the Pareto frontier achieved by the final models. All models exhibit competitive performance compared to the full-size model. Moreover, the superior performance of the recursive Gemma model strongly highlights the advantages of converting high-performing LLMs recursive architecture. Additional details can be found in Appendix L.
8

---

Published as a conference paper at ICLR 2025
Stepwise Average
Few-shal Accuracy
Few-sshof Accuracy
55
Few-shot Accuracy
50-
Mcr=150 B Mcy=008 KD
50
Gemma TinyLiama Pythia
50
3: Full-size
40
Redaced-size Recursive
Full-sip
Recuraive
0 Rolaxed
Relaaed
45
35
64 128
256
512
1.0B
1.58
2.0B
0.5B
1.0B
1.5B
2.0B
LoRA Rank
Model Sizes
Model Sizes
(a) Performance by LoRA rank
(b) Gains from longer KD training
(c) Overall performance
Figure 7: (a) Increasing the LoRA rank typically leads improved performance in relaxed Gemma models, attributed to the use of SVD initialization.
(b) Extended uptraining and knowledge distillation yielded substantial accuracy improvements for Gemma models. Note that the full-size model is pretrained model that is further uptrained on 60 billion tokens. (c) Recursive and Relaxed Recursive Transformers achieve a compelling Pareto frontier with respect model size and performance. Recur- sive and relaxed models used Stepwise and Average method initialize looped layers, respectively.
Table I: small loss coefficient to the first loop output (intermediate output) can significantly improve intermediate performance without compromising the final performance. Performance was evaluated under static-exiting scenario (Schuster al., 2022), where all tokens exit either first or second loop. We further trained the previously uptrained Gemma models 15 billion tokens (post-training Delta (A) denotes the performance changes in the final outputs after carly-exit training.
Uptrain
Looping
Early-Exit Train
Few-shot Accuracy
N-emb
PT
Nock
Block
Init
Noo
CE
KD
LD
HS
PQ
WG
ARC-e
ARC-c
OB
Avg
A
0.99B
/
15B
Step
53.0
57.3
73.2
56.2
56.1
29.2
36.6
51.7
0.99B
15B
Step
15B
Weighted
x
48.9
55.5
72.7
55.3
54.9
30.1
36.0
50.5
-1,2
49.5
54.8
72.0
53.4
54.1
29.1
35.6
49.8
0.99B
15B
Step
15B
Agg (0.1)
x
53.0
59.1
73.9
55.4
57.4
30.6
37.8
52.5
+0.8
45.9
51.2
71.4
54.5
48.1
26.8
32.0
47.1
0.99B
Step
15B
Weighted
47.7
48.3
54.9
55.1
72.1
73.2
55.6
55.9
54.3
54.5
29.1
28.4
37.2
35.4
49.9
50.4
1.3
0.99B
15B
Step
15B
Agg (0.1)
/
52.9
46.3
52.1
58.9
71.6
73.7
55.3
55.7
49.2
57.5
28.5
31.1
32.6
38.2
48.0
52.6
+0.9
3.7 EXILING AND RECURSIVE TRANSFORMER
The throughput of Recursive Transformers can be amplified by an early-exiting framework. Hence, we further train intermediate representations from fewer looping iterations enable token prediction. conducted an ablation study on various strategies, as summarized in Table (more detailed results are presented in Table 13). Directly applying the weighted CE loss (L = DB agLi where a; = E/E, i) commonly used prior works (Schuster al., 2022; Bae al., 2023) led an overemphasis on the training of intermediate representations. To address this, we employ an aggressive coeflicient strategy that aggressively reduces the loss coefficient for intermediate outputs while maintaining coefficient of for the final output. Our experiments demonstrated that an aggressive coefficient of 0.1, utilizing knowledge distillation from the detached final outputs (Bae al., 2023), effectively preserves final performance while enhancing intermediate performance. Notably, the first loop output yielded only difference of 4.6 percentage points in accuracy compared to the final output. This underscores the potential to maximize the benefits of carly-exiting in parameter-share LLMs.
We applied this post-training strategy for early-exiting to our final uptrained models (shown in $3.3), with all experimental results detailed in Appendix M. The aggressive coefficient strategy, combined with deff-distillation, consistently achieved the best performance for intermediate outputs while maintaining strong performance for the final loop output across all models. While the optimal strategy from non-relaxed models was applied to relaxed models, a tailored training approach could further improve intermediate loop output performance in Relaxed Recursive Transformers.

---

Published as a conference paper at ICLR 2025
60
Few shot Accuracy
x2.86
N-cmb
Loop LoRA
Batch
Exit
Acc.
Thr.
Av
Ass
1.99B
x
57.3
1080
x1.00
x0.71
1.99B
CSB
1528
x1.41
x1.00
50
x1.71
TinyLlama
Germma
0,99B
CDB /
54.0
2877
x2.66
x1.88
Pythia
1.07B
2
64
CDB
/
54.0
2157
X2.00
x1.41
45-
Vanilla (CSB)
Vanilaa
1.15B
1,30B
2
256
128
CDR
CDB
/
54.6
55.2
2149
1921
x1.99
11.78
x1.26
x1 66
Relaxed
Recursive
1.60B
2
512
CDB
/ 56.2
1719
x1.59
x1.13
1.5
3.5
Normalized Throughput
Figure 8: Continuous depth-wise batching (CDB) with early exiting enables Recursive Transformers to theoretically achieve significant throughput improvements. Throughput (tokens/sec) was averaged across SlimPajama. RedPajama, and PG19, and then normalized to the throughput of the vanilla Pythia model. The accompanying table gives detailed throughout and performance measurements for Gemma. Ay measures throughput relative to the vanilla Gemma model, while Aseg measures throughput relative to the vanilla Gemma model with continuous sequence-wise batching (CSB).
3.8 HYPOTHETICAL GENERATION SPEEDUP VIA CONTINUOUS DEPTII-WISE BATCHING
How we theoretically approximate actual throughput As developing practical early-exiting algorithms is beyond the scope of this work, we present hypothetical throughput improvements based on an oracle-exiting approach (Schuster et al., 2022; Bae et al., 2023). This assumes that tokens exit at the earliest looping block where their prediction aligns with the final loop's 's prediction. We simulated the generation of language modeling datasets as if they were generated by our models, to obtaín the exit trajectory for each token. Then, we measured the average per-token generation time under specific constraints, such as different memory limit or context lengths. Using these measurements and the exit trajectory data, conducted simulations to estimate theoretical throughput. Detailed explanations and limitations are discussed in Appendix N.
Continuous depth- wise batching with early-exiting can substantially boost throughput Figure 8 illustrates the throughput of our proposed models and the vanilla Transformer across three architec- tures. We consistently achieve higher speeds than the vanilla models by combining continuous depth- wise batching with early-exiting, even surpassing those with continuous sequence-wise batching (Yu al., 2022; Kwon et al., 2023). In particular, Recursive models demonstrate up to 2.66x speedup in generation compared to vanilla counterparts. Additionally, the recursive Gemma model significantly outperforms the vanilla pretrained Pythia model, with nearly 4x improvement throughput. Relaxed recursive models show a clear trade-off between achievable performance and throughput, modulated by the degree of relaxation through LoRA ranks. This characteristic enables flexible model selection tailored to specific deployment scenarios. Comprehensive results are presented in Tables 17 and 19.
Takeaways for Continuous Depth-wise Batching
We analyze the potential for throughput improvement in the Recursive Transformer via contin- uous depth-wise batching, a novel inference paradigm. In theory, find that we can achieve up to 2-3x x speedup compared to a vanilla Transformer. This even outperforms the throughput gain achieved by existing continuous sequence-wise batching methods in vanilla models.
4 CONCLUSION
In this work, we introduced Recursive Transformers, in which we compress LLMs via parameter shar- ing across recursively looped blocks of layers. Additionally, presented novel relaxation strategy that allows for low-rank deltas between shared layers by integrating layer-specific LoRA modules into the fully-tied structure. Through novel initialization techniques for looped layers and LoRA modules, we achieved significant performance improvements that closely approximate the original pretrained model. Finally, by exploiting the recursive patterns and an early-exiting approach, we propose continuous depth-wise wiie batching paradigm tailored for efficient serving systems of Recursive Trans- formers. We theoretically demonstrated that an oracle-exiting strategy can yield substantial throughput gains, reaching up to 2-3 x speedup. The limitations and future works are discussed in Appendix A.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
We thank Jacob Eisenstein for valuable feedback on an earlier version the paper. We thank Jiyoun Ha, Alfred Piccioni, Dayeong Lee for the support with setting up the experimental environment. We also thank Donald Metzler, Ivan Korotkov, Jai Gupta, Sanket V. Mehta, Vinh Q- Tran, Brennan Saeta, Jean-François Kagy, Zhen Qin, and Jing Lu for helpful conversations. Finally, thank the Google Cloud Platform for awarding Google Cloud credits for this project.
REPRODUCIBILITY STATEMENT
To ensure the reproducibility of work, we provide comprehensive description of our model architectures in Appendix F, and details of experimental settings can be found in Appendix G. We utilized the open-source HuggingFace framework and followed established open-source frameworks for evaluation, further enhancing reproducibility. We plan to release the source codes upon publication facilitate future research.
REFERENCES
Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. On-policy distillation language models: Leamning from self- generated mistakes. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024, 2024. URL https: /opp/reveeeeeeeeeeeeeeeeee http:/neee/ooooreeo forum? Id=3zXtaqxLhix
Joshua Ainslie, James Lee- Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit training generalized multi-query transformer models from multi-head checkpoints. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 10, 2023, pp. 4895-4901 2023. doi: 10.18653/V 1/2023.. EMMLP-MAIN.299
Sanghai. GQA:
Sangmin Bae, Jongwoo Ko, Hwanjun Song, and Se- Young Yun. Fast and robust early-exiting framework for autoregressive language models with synchronized parallel decoding. In Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, pp. 5910-5924 2023. doi: 10.18653/V 1/2222/.2222.PPP. 362.
Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. Deep equilibrium models. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NewrIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, neurips cc/ acc/paper/20/ 2ttp//2299/hhha/ //hash/ - .tmll.
pp. 2019. URL wtt.ee/e60000111eeeeeeeeeeaeeae0
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar der Wal. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, ICMI, 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202, pp. 2397-2430, 2023. URL https: http:////oceeeeeneeeeeee www.ang.mm.......m.............. malr.press. 21202//bderraa222. 3a. html.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense natural language. In Proceedings of the AAAI conference artificial intelligence, volume 34, pp. 7432-7439, 2020.
William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, and Jonathan Ragan- Kelley. Reducing transformer key-value cache size with cross-layer attention. CoRR, abs/2405. 12981, 2024. doi: 10.48550/ARXIV24005..
Lequn Chen, Zihao Ye, Yongji Wu, Danyang Zhuo, Luis Ceze, and Arvind Krishnamurthy. Punica: Multi-tenant lora serving. Proceedings of Machine Learning and Systems, 6:1-13, 2024a.
11

---

Published as a conference paper at ICLR 2025
Yanxi Chen, Xuchen Pan, Yaliang Li, Bolin Ding, and Jingren Zhou. EE-LLM: large- training and inference of early-exit large language models with 3d parallelism. In Forty-first International Conference on Machine Learning. ICML 2024, Vienna, Austria, July 21-27, 2024, 2024b. URL httpps: http:///peeeevveeeeeeeeeeeeeeeee
iew. ew. net/ forum? d=xFk Cw9
Jeffrey Cheng and Benjamin Van Durme. Compressed chain of thought: Efficient reasoning through dense representations. arXiv preprint arXiv:2412. 13171, 2024.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprini arXiv: 1.Xxvv/8003055577 2018.
Together Computer. Redpajama: An open source recipe to reproduce Ilama training dataset, 2023.
URL httph:tpp////. /ttp:////tttubbtt com www/toge toge St comput /Rede RedF jmaa-Data
Róbert Csordás, Kazuki Irie, Jürgen Schmidhuber, Christopher Potts, and Christopher D Manning.
Moeut: -fxture-of-experts universal transformers. arXiv preprint arXiv:2405.16033 2024.
Raj Dabre and Atsushi Fujita. Recurrent stacking of layers for compact neural machine translation models. In The Thirty- Third AAAI Conference Artificial Intelligence, AAAI 2019, The Thirty- First Innoaatiee Applications of Artificial Intelligence Conference, IAAI 2019, The Ninth AAAI
Symposium on Educational Advances in Artificial Intelligence, EAAI 2019, Honolulu, Hawaii, USA, January 27- February 2019, pp. 6292-6299, 2019. doi: 10.1609/AAIL, A1300A.31L11000.33316020999
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Advances in Neural In- formation Processing Systems 35: Anmual Conference on Neural Information Process- ing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December 2022, 2022. URL http: /papers/pppprsseee/ee nips cctpaper_afiieeeeeeeeeee fwwess/ppeee/eeeeeeeeeeeaaaeeaat (e177722222E00077330 E00.000200000003330000002266665
/hash/ html.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019. URL https://///........... http:://pppeeeeeeeee...eee /trr//uruteeeeeeeeeo forumfirrmeo Hyz6RiR9Y7.
Mostafa Dehghani, Yi Tay, Anurag Amab, Lucas Beyer, and Ashish Vaswani. The efficiency misnomer. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022, 2022. URL https: /wpporpeeevveeeeeeee new. net/forum? YR....
Luciano Del Corro, Allie Del Giorno, Sahaj Agarwal, Bin Yu, Ahmed Awadallah, and Subhabrata Mukherjee. Skipdecode: Autoregressive skip decoding with batching and caching for efficient 11m inference. arXiv preprint arXiv:2307.0266 2023.
Devvrit, Sneha Kudugunta, Aditya Kusupati, Tim Dettmers, Kaifeng Chen, Inderjit Dhillon, Yulia Tsvetkov, Hannaneh Hajishirzi, Sham Kakade, Ali Farhadi, and Prateek Jain. Matformer: Nested transformer for elastic inference, 2023. URL https: : www.i.www..............oo.o..oor abbs/////////// 2210 07707.
Abhimanyu Dubey, Abhinav Jauhri, Abhinay Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshomn Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurélien Rodriguez, Austen Gregerson. Ava Spataru, Baptiste Rozière, Bethany Biron, Binh Tang. Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong. Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Graeme Nail, Grégoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel M. Kloumann, Ishan Misra, Ivan Eytimoy, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet
12

---

Published as a conference paper at ICLR 2025
Shah, Jelmer van der Linde, Jennifer Billock. Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang. Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, and et al. The llama 3 herd models. CoRR, abs/2407.2178.. 2024. doi: (048500ARXIV.247..2111
David Eigen, Jason Tyler Rolfe, Rob Fergus, and Yann LeCun. Understanding deep architectures using recursive convolutional network. In 2nd International Conference on Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Workshop Track Proceedings, 2014. URL http: att/vvivvvvv////////ee/////eevva org/ abs/ /1////112222 1847.
Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth-adaptive transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, 2020. URL https: wwwpo/ppennn..ee/ee//eeee net forum? iinn Sd= K.
Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, Ahmed A Aly, Beidi Chen, and Carole-Jean Wu. Layerskip: Enabling early exit inference and self-speculative decoding. In Proceedings of the 62nd Annual Meeting of the Association for Compulational Linguisties (Volume 1: Long Pupers), ACI. 2024, Bangkok, Thailand, August / 11-16, 2024, pp. 12622-12642, 2024.
URL https: /aaaccaannn aaw.aaannnnnooooocaaaaaa org/o/////.... 2024 acl- -oog-loog.... 681.
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, 2020. https::: http:///ppnneeeeeeeeeeeeeee net. forum?
/
id-Syyo2yStDr.
Ying Fan, Yilun Du Kannan Ramchandran, and Kangwook Lee. Looped transformers for length generalization. arXiv preprint 3XXi::2409.156477 2024.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 23:120:1-120:3, 2022. URL https: //ttp://jnet //jmlr r. org/ haagg//aprss//// v/3//223///222 00021-0-9-..0 html.
Wenfeng Feng, Chuzhan Hao, Yuewei Zhang, Yu Han, and Hao Wang. Mixture-of-loras. An efficient multitask tuning method for large language models. In Proceedings of the 2024 Joint International Conference on Computational Linguistics, Language Resources and Evaluation, LREC/COLING 2024, 20-25 May, 2024, Torino, Italy, pp. //aclantho ogy- org/ 2024. 1 rerc-maa n. 994.
2024. URL https:
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019. URL https: https:://penneview..ee wwtp://reeeee/eiooooomoe rJ1- b3Rc7.
Qichen Fu, Minsik Cho, Thomas Merth, Sachin Mehta, Mohammad Rastegari, and Mahyar Najibi. Lazylim: Dynamic token pruning for efficient long context LLM inference. CoRR, Abs/2407.14457, 2024. doi: 10.48550/ARXIV.2407.11
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. Megablocks: Efficient sparse training with mixturee-ff-expere Proceedings of Machine Learning and Systems,
2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00 000070 2020.
Leo Gao. Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang. Kevin Wang. and Andy Zou. framework for few-shot language model evaluation, 12 2023.
13

---

Published as a conference paper at ICLR 2025
Khashayar Gatmiry, Nikunj Saunshi, Sashank J. Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn to implement multi-step gradient descent for in-context learning?, 2024.
Tao Ge, Si-Qing Chen, and Furu Wei. Edgeformer: parameter-efficient transformer for on-device seq2seq generation. In Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pp. 10786-10798, 2022. doi:
Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein. Scaling up test-time compute with latent reasoning: A recurrent depth approach. arXiv preprint arXiv:2502.05171, 2025.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202, pp. 11398- 11442, 2023. URL https: httpppp://ppooeeeeeee pttpppedoocengsseeeee mlr. malr.prest www.ppees..ee... /v022/ http////////////nnn//uu giw.nou2nn..a..mmmmma mm..
Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. Zamba: A compact 7b SSM hybrid model. CoRR, abs/2405.16712. 2024. doi: 10445500ARXXV.2405..655
Sachin Goyal, Ziwei Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024, 2024. URL hettps://openrevieee
Alex Graves. Adaptive computation time for recurrent neural networks. CoRR, 0bs//603.0899 2016a. URL http: /aw.ivvvvvvv./..a///a///avv.a..a wwtp:/org/b....../... 603 08983.
Alex Graves. Adaptive computation time for recurrent neural networks, 2016b.
Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. Minillm: Knowledge distillation of large
language models. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024, 2024. URL https: wpwnrevieneeeeeeeeeeeeeeeeeeeeee net forum? id=5h0qf7IBZZ.
/
Qiushan Guo, Zhipeng Yu, Yichao Wu, Ding Liang, Haoyu Qin, and Junjie Yan. Dynamic recursive neural network. IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pp. 5147-5156, 2019. doi: 10.1109/CVPRR.219.00022
Song Han, Jeff Pool, John Tran, and William J. Dally. Learning both weights and connections for efficient neural networks. CoRR, 0bs/1506.02622 2015. URL http: httt://arrxiv org/abgg 1506. 6.
Per Christian Hansen. The truncated svd as a method for regularization. BIT Numerical Mathematics, 27:534-553, 1987.
Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu. Jason Weston, and Yuandong Tian. Training large language models to reason in a continuous latent space. arXiv preprint arXiv:2412.06769, 2024.
Geoffrey E. Hinton, Oriol Vinyals, and Jeffire Dean. Distilling the knowledge in a neural network. CoRR, htd//5003.023... 2015. URL http:://arae wwtg/rsgg////oooooooo..oosssssse abs /1503.153.... 02531.
Jordan Hoffimann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimat large language models, 2022. URL https: : / aaxiv.. org abs/ 2203. 15556.
14

---

Published as a conference paper at ICLR 2025
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin de Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for NLP. In Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97, pp. 2790-2799, 2019. URL http:/ http://pproceee ppoceedinggs ngs. .mlr w.r...rsss//.. /tt/////////// /007//hulsse 1oyyl 9a html.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022, 2022. URL. https: wpwnreview.eeeeeeeee/eeee net/ http://wwwwee/t/uttttteeeee www.mmidd-neeeeeeeeeeeeeeeee
Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew G. Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for inference. In 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, pp. 2704-2713, 2018. doi: 10. 109//199 2018.00288..
efficient pnieger-arithmetic-one
Sanghoo Kim, Dabyun Kim, Chanjun Park, Wonsung Lee, Wonho Song, Yunsu Kim, Hyeonwoo Kim, Yungi Kim, Hyeonju Lee, Jihoo Kim. Changbac Ahn, Sconghoon Yang, Sukyung Lee, Hyunbyung Park, Gyoungjin Gim, Mikyoung Cha, Hwalsuk Lee, and Sunghun Kim. SOLAR 10.7b: Scaling large language models with simple yet effective depth up-scaling. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies: Industry Track, NAACL 2024, Mexico City, Mexico, June 16-21, 2024, pp. 23- 2024. doi: 10.18653/V 10B33/1/2224NNAACC CADDO-NNAALLINDSSRRY....
Yoon Kim and Alexander M. Rush. Sequence-level knowledge distillation. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas, USA, November 1-4, 2016, pp. 1317-1327, 2016. doi: 10.18653/V 1139.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, SOSP 2023, Koblenz, Germany, October 23-26, 2023, pp. 611-62216 2023. doi: 10.1145/3600006.661333
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. ALBERT: A lite BERT for self-supervised learning of language representations. In 8th International Conference on Learming Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, 2020. URL https: wwwnoeeieeeeeeeeee net/ forum? d=H1 eA et.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202, pp. 19274-19286, 2023. URL https:/ /proceedings..... mlr. pl.sspress/vee v2tsss//222//// 2ttees///222//viithaaa222222 htt//eleeaaaanaaaaaeeeeaaea hhml..
Fangcheng Liu, Yehui Tang, Zhenhua Liu, Yunsheng Ni, Kai Han, and Yunhe Wang. Kangaroo: Lossless self-speculative decoding via double early exiting. arXiv preprini arXiv:2404. 18911, 2024a.
Xiao Liu, Kaixuan Ji, Yicheng Fu, Zhengxiao Du, Zhilin Yang, and Jie Tang. P-tuning v2: Prompt tuning can be comparable to fine-tuning universally across scales and tasks. CoRR, abs/2l 10..07602 2021. URL https:/ / arxiv. otggabs/////eoooooooe//.///aaaar
abs/ 1abbs/211 22100.100 07602.
Zechun Liu, Changsheng Zhao, Forrest N. landola, Chen Lai, Yuandong Tian, Igor Fedorov, Yunyang Xiong, Ernie Chang, Yangyang Shi, Raghuraman Krishnamoorthi, Liangzhen Lai, and Vikas Chandra. Mobilellm: Optimizing sub- billion parameter language models for on-device use cases. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024, 2024b. URL https: // opwnrevieeeeeeeeeoe eew.wwwww net./ fo ruomu i G=x.GUXX
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang. Yuandong Tian, Christopher Ré, and Beidi Chen. Deja vu: Contextual sparsity for efficient llms at inference time. In International Conference on Machine Learning, ICML
15

---

Published as a conference paper at ICLR 2025
2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202, pp. 22137 22176, 2023, URL https: pptp:////ccceeengg.ccoe.. www.mm.....ss.... 102002/1iu23e am..aam.mmmmmmmm..mmmmmmmmm...a. am.html.
Ilya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings, 2017. URL https: /ttp:////ppe/eeeeeee/em net forum? ?id= Skq89Scxx.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019. URL https:://// http://ppeneeeeeeeeeeeeeeeeeeeee net for Bttp://oooorrunmaadeeeeee
Sean Michael McLeish and Long Tran-Thanh. [re] end-to-end algorithm synthesis with recurrent networks: Logical extrapolation without overthinking. In ML Reproducibility Challenge 2022, 2022.
Fanxu Meng, Zhaohui Wang, and Muhan Zhang. Pissa: Principal singular values and singular vectors adaptation of large language models. CoRR, Abs//244.002488, 2024. doi: 10.48550/ARXIV.240044 02948.
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Hongyi Jin, Tianqi Chen, and Zhihao Jia. Towards efficient generative large language model serving: A survey from algorithms to systems. CoRR, abs/2312.15234, 2023. doi: 10.48550/ARXI.2311..5555
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? A new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods Natural Language Processing, Brussels, Belgium, October 31 - -Nveemeer 4, 2018, pp.
2018. doi: 10. 1//53/V1/D18-111
Jeremiah Milbauer, Annie Louis, Mohammad Javad Hosseini, Alex Fabrikant, Donald Metzler, and Tal Schuster. LAIT: Efficient multi-segment encoding in transformers with layer-adjustable interaction. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 10251-10269, July 2023. doi: 10.18653/v1/2023. acl-long.571.
Amirkeivan Mohtashami, Matteo Pagliardini, and Martin Jaggi. Cotformer: More tokens with attention make up for less depth. arXiv preprint arXiv:2310. 10845, 2023.
OpenAl. GPT-4 technical report. CoRR, abs/2303.00877, 2023. doi: 10.48500AARXIV.2233.0777
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. arXiv preprint arXiv: 1606.60063. 2016.
Jacob Pfau, William Merrill, and Samuel Bowman. Let's think dot by dot: Hidden computation in transformer language models. arXiv preprint 2rXiv:2404.15755 2024.
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. In Proceedings of Machine Learning and Systems, volume 5. pp. 606- 624, 2023. URL file/ Aile/c4be11 2et pttp:/aa27000--papp----aees----- mlssss222sss222323mmm pdf.
mwsys htwp://g...nlssss.oo/aaaeeeeeeee filesfffeesss wtw.ffepateeeee/eeeeeeeeeaaaeeee htppp/ppppr/////////
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling, arXiv preprint arXiv:1911.0550) 2019.
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young. Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George van den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang, Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang. Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen. Siddhant Jayakumar, Elena Buchatskaya. David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, Laurent Sifre, Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro,
16

---

Published as a conference paper at ICLR 2025
Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, Jean-Bapy cun--Bptiste Lespiau, Maria Tsimpoukelli, Nikolai Grigorev, Doug Fritz, Thibault Sottiaux, Mantas Pajarskas, Toby Pohlen, Zhitao Gong, Daniel Toyama, Cyprien de Masson d' d'Attmme, Yujia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury, Matthew Johnson, Blake Hechtman, Laura Weidinger, lason Gabriel, William Isaac, Ed Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals, Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Koray Kavukcuoglu, and Geoffrey Irving. Scaling language models: Methods, analysis & insights from training gopher, 2021.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Vivek Ramanujan, Mitchell Wortsman, Aniruddha Kembhavi, Ali Farhadi, and Mohammad Rastegari. What's 's hidden in randomly weighted neural network? In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, pp. 11890-11899, 2020. doi: 10. 0B09CCCVPR22000220000033
David Raposo, Samuel Ritter, Blake A. Richards, Timothy P. Lillicrap, Peter Conway Humphreys, and Adam Santoro. Mixture-of-depths: Dynamically allocating compute in transformer-based language models. CoRR, abs/2404.02258, 2024. doi: 1..4550UARRXIV2244002222
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System optimiza- tions enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 3505-3506, 2020.
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy P. Lillicrap, Jean- Baptiste Alayrae, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, loannis Antonoglou, Rohan Anil, Sebastian Borgeaud, Andrew M. Dai, Katie Millican, Ethan Dyer, Mia Glaese, Thibault Sottiaux, Benjamin Lee, Fabio Viola, Malcolm Reynolds, Yuanzhong Xu, James Molloy, Jilin Chen, Michael Isard, Paul Barham, Tom Hennigan, Ross Mcllroy, Melvin Johnson, Johan Schalkwyk, Eli Collins, Eliza Rutherford. Erica Moreira, Kareem Ayoub, Megha Goel, Clemens Meyer, Gregory Thornton, Zhen Yang, Hemryk Michalewski, Zaheer Abbas, Nathan Schucher, Ankesh Anand, Richard Ives, James Keeling, Karel Lenc, Salem Haykal, Siamak Shakeri, Pranav Shyam, Aakanksha Chowdhery, Roman Ring, Stephen Spencer, Eren Sezener, and et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. doi: 1044550/ARXX.22400055555
CoRR, abs/2403.05530, 2024.
Morgane Rivière, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, Johan Ferret, Peter Liu, Pouya Tafti, Abe Fricsen, Michelle Casbon, Sabela Ramos, Ravin Kumar, Charline Le Lan, Sammy Jerome, Anton Tsitsulin, Nino Vieillard, Piotr Stanczyk, Sertan Girgin, Nikola Momchev, Matt Hoffman, Shantanu Thakoor, Jean Bastien Grill, Behnam Neyshabur, Olivier Bachem, Alanna Walton, Aliaksei Severyn, Alicia Parrish, Aliya Ahmad, Allen Hutchison, Alvin Abdagic, Amanda Carl, Amy Shen, Andy Brock, Andy Coenen, Anthony Laforge, Antonia Paterson, Ben Bastian, Bilal Piot, Bo Wu, Brandon Royal, Charlie Chen, Chintu Kumar, Chris Perry, Chris Welty, Christopher A. Choquette-Choo, Danila Sinopalnikov, David Weinberger, Dimple Vijaykumar, Dominika Rogozinska, Dustin Herbison, Elisa Bandy, Emma Wang. Eric Noland, Erica Moreira, Evan Senter, Evgeni Eltyshev, Francesco Visin, Gabriel Rasskin, Gary Wei, Glenn Cameron, Gus Martins, Hadi Hashemi, Hanna Kllimezak-Plucinsee Jacinda Mein, Jack Zhou, James Svensson, Jeff Stanway, Jetha Chan, Jin Peng Zhou, Joana Carrasqueira, Joana Iljazi, Jocelyn Becker, Joe Fernandez, Joost van Amersfoort, Josh Gordon, Josh Lipschultz, Josh Newlan, Ju-yeong Ji, Kareem Mohamed, Kartikeya Badola, Kat Black, Katie Millican, Keelin McDonell, Kelvin Nguyen. Kiranbir Sodhia, Kish Greene, Lars Lowe Sjösund, Lauren Usui, Laurent Sifre, Lena Heuermann, Leticia Lago, and Lilly McNealus. Gemma 2: Improving open language models at practical size. CoRR, abs/2408. 20112000000011111 18, 2024. doi: 10.48550/AR AA..550AARXI..248...... XTV.2408..01111111 18.
Harleen Batra, Harsh Dhand, Ivan Nardini,
17

---

Published as a conference paper at ICLR 2025
Jonathan S. Rosenfeld, Amir Rosenfeld, Yonatan Belinkov, and Nir Shavit. A constructive prediction of the generalization error across scales. In International Conference on Learning Representations, 2020. URL https::///.......... http://pppppeeeeeeeeeee/ee http:s///ppeneeeeeee..cooooo net/ http://wwww...tttonntneeeeeeae wtt.:rr......y.d.e.E.......
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge at scale. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirly-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial In- telligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 8732 8740, 2020. doi: 10.1609/AAAL. V34105.6399.
Nikunj Saunshi, Stefani Karp, Shankar Krishnan, Sobhan Miryoosefi, Sashank Reddi, and Sanjiv Kumar. On the inductive bias of stacking towards improving reasoning, 2024.
Pedro Savarese and Michael Maire. Learning implicitly recurrent cnns through parameter sharing. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019, URL https: opw.reviewwwwweeeeepeeeee.roeewm net Forum? id weg.-rngynnn 9Fm.
Tal Schuster, Adam Fisch, Tommi Jaakkola, and Regina Barzilay. Consistent accelerated inference via confident adaptive transformers. In Proceedings of the 202] Conference Empirical Methods in Natural Language Processing, pp. 4962-4979, November 2021. doi: 10. emnlp-main.406.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 December 9, pacc/ppppr_filllc heiilles/apppee / / 2022, 2022. URL http: httpp/ppppp./e///ee/eeeeeeeee..p nips. tt/ 6fac9e31 6a4 444444577724444 -orpaaeeeeeeect-aaaaaaa----o--ee wow.raacttttcoooeeeeeeeeeeetaaaa html.
Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Furong Huang, Uzi Vishkin, Micah Goldblum, and Tom Goldstein. Can you learn algorithm? generalizing from easy to hard problems with recurrent networks. In Advances Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 6695-6706, 2021. URL htt.ttppp//pppccceeeeiiee www.:iiinnnnnnnunne.eeeeeeeee.ee haapp..cc/peee cctppppee//21......e..e../ / 3501
html.
Noam Shazeer. Fast transformer decoding: One write-head is all you need. CoRR, abs/1911.02150, 2019. URL
02150.
Noam Shazeer. GLU variants improve transformer. CoRR, abs/2002.05202, 2020. URL https: // //arxiv. org/ abs/2002. 05202.
Zhiqiang Shen, Zechun Liu, and Eric P. Xing. Sliced recursive transformer. In Computer Vision ECCV 2022: 17th European Conference, Tel Aviy, Israel, October 23-27. 2022, Proceedings, Part XXIV, volume 13684, pp. 727-744, 2022. doi:
Ying Sheng. Shiyi Cao, Dacheng Li, Coleman Hooper, Nicholas Lee, Shuo Yang, Christopher Chou, Banghua Zhu, Lianmin Zheng, Kurt Keutzer, Joseph E. Gonzalez, and Ion Stoica. S-lora: Serving thousands concurrent lora adapters. CoRR, 0ts//311..32253 2023. doi: 10.48550/ARXIV.3111. 03285.
Alex Sherstinsky, Fundamentals of recurrent neural network (RNN) and long short-term memory (LSTM) network. CORR, 0tb//888.00333 2018. URL httt:://axxi..err abs///////// 18008. 03314.
Kyuhong Shim, Jinkyu Lee, and Hyunjae Kim. Leveraging adapter for parameter-efficient asr encoder. In Proc. Interspeech 2024, pp. 2380- 2024.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hes- tness, and Nolan Dey. plicated version of RedPajama. s slimpajarre a- http://gnnnn..ccee.eca/////////e /
SlimPajama: A 627B token cleaned and dedu-
http://we..bbbaa...eeee.oo.... net/blog/ http://ewweeeeecddeeeeeeeeeeeeee of-redaajaaaaaaaaaaaaaaapaaa-aae 667B.
2023. URL https: httpss:///ggggggggggeeee c
18

---

Published as a conference paper at ICLR 2025
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang. Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. CoRR, abs/2405.05254, 2024. doi: 1044850/ARXXI.2440.00055
Sho Takase and Shun Kiyono. Lessons on parameter sharing across layers in transformers. In Proceedings of The Fourth Workshop on Simple and Efficient Natural Language Processing, SustaiNLP 2023, Toronto, Canada (Hybrid), July 13, 2023, pp. 78-90, 2023. doi: 10.18/53/V1/ 4023.SUSTAINLP-.5.
Peng Tang, Pengkai Zhu, Tian Li, Srikar Appalaraju, Vijay Mahadevan, and R. Manmatha. DEED: dynamic early exit on decoder for accelerating encoder-decoder transformer models. In Findings of the Association for Computational Linguistics: NAACL 2024, Mexico City, Mexico, June 16-21, 2024, pp. 116-131, 2024, doi: 10.18653/V 1/2/33//1/222..1//20 FINTIO2AAIIDIINNNNAAAAAILLL
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based gemini research and technology. arXiv preprint arXiv:2403.08229 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Zhongwei Wan, Xin Wang, Che Liu, Samiul Alam, Yu Zheng, Jiachen Liu, Zhongnan Qu, Shen Yan, Yi Zhu, Quanlu Zhang, Mosharaf Chowdhury. and Mi Zhang. Efficient large language models: A survey. Trans. Mach. Learn. Res., 2024, 2024. URL https: /ttp//pppeeeeeeeeeeee id=bsCCJJb08A,
/ ew net forum?
Haowen Wang, Tao Sun, Congyun Jin, Yingho Wang, Yibo Fan, Yunqi Xu, Yuliang Du, and Cong Fan. Customizable combination of parameter-efficient modules for multi-task learning. In The Twelfth International Conference on Learning Representations, 2023.
Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou. f-divergence minimization for sequence- level knowledge distillation. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 10817-10834, July 2023. doi: 10. 111666////20223acl--ooog
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: State-of -oee----thherrr art natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations, pp. 38 45, 2020.
Yingce Xia, Tianyu He, Xu Tan, Fei Tian, Di He, and Tao Qin. Tied transformers: Neural machine translation with shared encoder and decoder. In The Thirly-Third AAAI Conference Artificial Intelligence, AAAI 2019, The Thirty-First Innovative Applications of Artificial Intelligence Confer- ence, IAAI 2019, The Ninth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2019, Honolulu, Hawaii, USA, January 27 February 1, 2019, pp. 1466,, 2019. doi: 00.1609/AAAALV333111.0 3331101.3301446.
Liu Yang, Kangwook Lee, Robert D. Nowak, and Dimitris Papailiopoolos.. Looped transform- are better at learning learning algorithms. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024, 2024. https: //openreview net wwt./ett//ounn forum? dd-HHHbR<DDDD...
Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung- Chun. Orea: A distributed serving system for transformer-based generative models. In 16th USENIX Symposium Operating Systems Design and Implementation, OSDI 2022, Carlsbad, CA, USA, July 11- 13, 2022, pp. 521-538, 2022. URL httph:h///// /www://www.. usenix http//ooeeeeeee/eeeeeeeee presentation ietennatioon/oo yuu.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv: 1905.0780 2019.
19

---

Published as a conference paper at ICLR 2025
Dewen Zeng, Nan Du, Tao Wang, Yuanzhong Xu, Tao Lei, Zhifeng Chen, and Claire Cui. Learning to skip for language modeling. CoRR, abs/2311.15436, 2023. doi: (0-4850/ARXIV.311.55552
Biao Zhang and Rico Sennrich. Root mean square layer normalization. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Pro- cessing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 12360- 12371, 2019. URL https:: neuriinnnnnn..un.n.o.e.e.m....ee newrips cctppppe///11g.//// hash// 18223338224883057777773------100 html.
Jun Zhang, Jue Wang, Huan Li, Lidan Shou, Ke Chen, Gang Chen, and Sharad Mehrotra. Draft& verify: Lossless large language model acceleration via self-speculative decoding. Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), ACL 2024, Bangkok, Thailand, August 11-16, 2024, pp. 11263-11282, 2024a. URL https: / acl antho ogy org/ 2024 wwl.0244cc11..0 -long- 607.
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model, 2024b.
Zixuan Zhou, Xuefei Ning, Ke Hong. Tianyu Fu, Jiaming Xu, Shiyao Li, Yuming Lou, Luning Wang, Zhihang Yuan, Xiuhong Li, Shengen Yan, Guohao Dai, Xiao-Ping Zhang. Yuhan Dong, and Yu Wang. A survey efficient inference for large language models. CoRR, 0as/2404.142999 2024. doi: (044550/ARXIV.244.1422
20

---

Published as a conference paper at ICLR 2025
A LIMITATION AND FUTURE WORK
Compatibility with sparse designs Sparsity-based approaches, such as pruning (Han et al., 2015), quantization (Jacob et al., 2018), or layer-skipping mechanisms (Raposo et al., 2024), recently also give promising model compression results. In fact, many of these techniques are complementary to our approach: for example, we can seamlessly have a recursive, sparse architecture. In this work, we rather choose to focus on recursive dense designs (a domain that remains relatively unexplored) that also have very promising, practical performance traits (i.e., allowing for continuous depth-wise batching for faster throughput). That said, while in this work we take the first step at studying Relaxed Recursive Transformer with dense Transformer layers, we do believe that incorporating Mixture-of- expert (Fedus al., 2022), activation-skipping (Liu al., 2023) and SSM components (Glorioso et al., 2024) within the looped blocks are promising directions for future research.
Latent Reasoning via Recurrent Depth Beyond efficiency gains through down- down-scaling materialized parameters with recursive patterns, an alternative research direction lies in scaling-up recurrent depth to facilitate latent reasoning. Specifically, recurrent computation can manifest thinking vertically by processing internal hidden states at each depth. One promising approach involves leveraging contem- plation tokens (Pfau et al., 2024; Goyal et al., 2024) or latent (continuous) space representations (Hao et al., 2024; Cheng & Van Durme, 2024) to enhance reasoning in mathematical and code generation tasks. Another valuable direction focuses on enhancing the efficiency and training stability of ap- proaches that recursively scale-up depth, building upon concepts of deep thinking (Schwarzschild et al., 2021; Geiping et al., 2025).
Scaling up Recursive Transformers Scaling our approach to larger LLMs (7B and beyond) is a promising avenue for future research. While our methodology is expected to remain effective, achieving comparable performance may require significantly higher uptraining costs. Increased model size offers the potential for a reduced memory footprint from recursive patterns; however, it is unclear whether this translates to larger batch sizes, given the corresponding increase in hidden dimensions. Nevertheless, our continuous depth-wiss wise batching will yield considerable gains in serving efficiency.
Beyond hypothetical gemeration speedup Our oracle-exiting approach assumes any intermediate prediction matching the final output can be exited. However, accurate throughput measurement requires confidence-based early-exiting algorithms (Schuster et al., 2022; Bae et al., 2023). Moreover, practical deployment needs to address decoding bottlenecks like key-value cache computation for exited tokens remaining loops. Nevertheless, there are potential solutions: for example, the missing KV cache computations can be addressed by leveraging continuous depth-wise batching, allowing the KV cache for exited positions in subsequent loops to be performed parallel with the computations for the next sequence sample. Moreover, we can explore key-value cache sharing strategies (Sun et al., 2024; Brandon et al., 2024) for future work.
Efficient serving of multi-LoRA layers Relaxed models require the computation of distinct LoRA modules during batched inference, akin to multi-task learning (Feng et al., 2024; Wang et al., 2023), hindering parallel computation. We concatenated LoRA weights into single weight improve efficiency over sequential computation, yet it introduces redundancy. To mitigate this we can explore optimized CUDA kernels for LoRA serving (Sheng et al., 2023; Chen al., 2024a) and parallelization across accelerators, inspired by distributed training for Mixture of Experts (Fedus al., 2022; Gale al., 2023).
B RELATED WORK
Cross-layer parameter sharing has proven to be an effective method for achieving parameter efficiency in deep learning models such as RNNs (Sherstinsky, 2018; Graves, 2016b), CNNs (Eigen et al., 2014; Savarese & Maire, 2019; Guo et al., 2019; Shen et al., 2022), and the popular Transformer architecture. The Universal Transformer (Dehghani al., 2019), a recurrent self-attentive model, demonstrated superior performance to non-recursive counterparts with significantly fewer parameters. This cross- layer parameter sharing approach has subsequently been explored in various tasks, including language understanding (Lan et al., 2020), language modeling (Bai et al., 2019; Mohtashami et al., 2023; Liu et al., 2024b; Csordás et al., 2024; Glorioso et al., 2024), and machine translation (Dabre & Fujita, 2019; Milbauer et al., 2023; Xia et al., 2019; Takase & Kiyono, 2023; Ge et al., 2022). These methods often claim to achieve comparable performance with more compact models and increased
21

---

Published as a conference paper at ICLR 2025
computational speed, while also setting the ground for effective adaptive compute solutions (Dehghani et al., 2019; Graves, 2016b; Schuster et al., 2021).
Concurrently, there has been growing interest in exploiting recurrent architectures for algorithmic or logical reasoning tasks (Saunshi et al., 2024). Prior research (Schwarzschild et al., 2021; McLeish & Tran-Thanh 2022) has shown that recurrent networks can extrapolate reasoning strategies learned on simple problems to harder, larger problems through additional recurrences during inference. The looped Transformer structure has also been employed to emulate basic computing blocks for program simulation (Giannou et al., 2023), to learn iterative algorithms for data-fitting problems (Yang al., 2024), to achieve length generalization in algorithmic tasks (Fan et al.. 2024), and promising theoretical potential for few-shot learning (Gatmiry et al., 2024).
However, previous work has predominantly focused on relatively small Transformer models, trained from scratch without leveraging pretrained model weights. Our work distinguishes itself by investi- gating parameter sharing in the context of LLMs and proposing effective initialization strategies that leverage the knowledge embedded within existing LLMs. To the best of our knowledge, we are the first to propose a generalized framework for parameter-shared models, enabling relaxation weight tying constraints through layer-specific modules.
In this paper, we also discuss how Recursive Transformers can be well suited for early-exiting techniques accelerate decoding in LLMs. The inherent recursive structure readily enables early- exiting for individual responses within a large serving batch, which is often a practical limitation of such techniques. Vanilla Transformers encounter synchronization issue with early-exiting, where the model must forward all layers even a single token a batch requires full processing (exited tokens must wait for them). Several approaches attempt exploit this idle time by computing missing KV caches for exited tokens in later layers, which are essential for subsequent sequence generation. These techniques include state propagation (Schuster al., 2022: Elbayad et al., 2020), SkipDecode (Del Corro et al., 2023), and parallel decoding (which can be combined with Speculative Decoding) (Bae et al., 2023; Elhoushi et al., 2024; Liu et al., 2024a; Chen et al., 2024b; Tang et al., 2024). Nevertheless, the heterogeneous parameters across varying model depths still hinder the efficient progression of exited tokens subsequent sequences. Transformers enable parallel computation for tokens at different depths and sequences (in continuous depth-wise batching paradigm)-also allow for parallel computation of missing KV caches with minimal overhead during the memory-boundee decoding phase.
contrast, our Recursive
C COMPONENTS IN TRANSFORMER ARCHITECTURE
The Transformer block consists of two core components: a multi-head attention (MHA) mecha- nism and feed-forward network (FFN). MHA utilizes multiple attention heads to capture diverse relationships within the input sequence. The computation within each attention head is formulated as:
QK"
Attention(4 Q, K, V) softmax
Vdx
where Q, K, and V are linear projections of the input, parameterized by learned weight matrices w wk, and W/. respectively. The outputs from each head of the multi-head attention are concatenated and then projected back to the original hidden size using learned weight matrix Waut
While the FFN structure typically consists of two linear transformations, in the Gemma model, deviates from this standard architecture follows:
FFN(x) wilowm (GLLU(XWo0 *
with three learned linear weight matrices and a GeGLU activation (Shazeer, 2020).
D PARAMETER SHARING STRATEGY
Takase & Kiyono (2023) discuss three strategies for partial layer tying in Transformer models, as depicted in Figure 9. The SEQUENE strategy is the simplest, assigning the same parameters to consecutive layers. The CYCLE strategy repeatedly stacks single block of unique layers to achieve the desired depth. Meanwhile, the CYCLE (REV) strategy stacks the lower layers in reverse order for the remaining layers.
22

---

Published as a conference paper at ICLR 2025
In the comparative analysis of SEQUENCE and CYCLE strategies (Liu et al., 2024b), CYCLE demonstrated marginally superior zero-shot performance. Although the SEQUENCE approach, which caches shared weights (the capacity of SRAM is typically sufficient to hold a single transformer block) and computes them iteratively, has the potential to mitigate the weight transfer bottleneck between SRAM and DRAM, prioritized compatibility with early-exiting. Consequently, specifically employed the CYCLE strategy, which enables continuous depth-wise batching and thereby maximizes the throughput of Recursive Transformers.
Layer
Layer
Layer
Layer 5
Layer 5
Layer 5
Layer
Layer
Layer 4
Layer
Layer
Layer 3
Layer
Layer
Layer
Layer
Layer
Layer 1
(a) SEQUENCE
(b) CYCLE
(c) CYCLE
Figure 9: Three strategies for parameter sharing (Takase Kiyono, 2023). The examples utilize models with six layers, where identical colors represent shared weights.
E ILLUSTRATIVE EXAMPLES OF SVD INITIALIZATION IN RELAXED RECURSIVE TRANSFORMER
We propose SVD initialization approach for LoRA modules within Relaxed Recursive Transformer, effectively steering the summation of base and LoRA weights towards the pretrained weights of their corresponding depth. Figure 10 illustrates initialized under three different initialization techniques (Stepwise, Average, and Lower) for looped layers. One crucial point that if the initialized looped layer's weights match those of the original pretrained model, its corresponding LoRA module undergoes standard zero initialization: random Gaussian for matrix and for B. For example, with the Stepwise method, the first loop's 's LoRA module receives standard zero initialization, while the second loop's 's LoRA is initialized using our proposed initialization.
overview of how the LoRA module is
Layer
x 2
x
Layer
Layer 6
Layer 3,
Layer
Layer
X
Layer 3
Layer 2,
Layer 2
Layer
Layer
Layer 1,
Layer
Layer
Layer
Stepwise-based
Average-based
Lower-based
Figure 10: Overview of the proposed SVD initialization method for the Relaxed Recursive Trans- former. We visualize how LoRA modules initialized under three different looping initialization methods, assuming full-size model with six layers and two looping blocks. A matrices are colored according the corresponding full-size model weights, while B matrices are colored based on the looped layer weights. White B matrices indicate cases where the full-size model and recursive model weights are identical, resulting in standard zero initialization.
23

---

Published as a conference paper at ICLR 2025
F OVERVIEW OF THREE PRETRAINED LLMS
We utilized three pretrained models Gemma 2B (Team et al., 2024), TinyLlama 1.1B (Zhang et al., 2024b), and Pythia IB (Biderman al., 2023) and converted them into Recursive Transformers. De- tailed model configurations are summarized Table 2, and their corresponding few-shot performance results presented in 3.
Table 2: Key parameters and pretraining details of three models. The sizes of each model refer to the number of embedding parameters (embedding matrices and classifier heads), and all other non- embedding parameters. Gemma and TinyLlama utilize Multi-Query (Shazeer, 2019) and Grouped- Query (Ainslie et al., 2023) attention mechanisms, which leads to reduced number of key-value heads. Especially, we take an early TinyLlama checkpoint study an trained model.
Model Architecture
Pretraining
Models
N-emb
Emb
N₂
dowadet
Nicod
NKv
ducus
Vocab
Datuset
No.k
Lete
Gemma 2B
1.98B
0.52B
18
2048
8
256
256K
Unreleased
3T
8K
TinyLlama I.IB
0.97B
0.13B
2048
32
32K
SlimPujama
Starcoderdata
73B
32B
2K
Pythia
0.81B
0.21B
2048
8
Pile
2K
Table Few-shot performance of pretrained models. Few-shot accuracy measured the LAM- BADA, HellaSwag, PIQA, WinoGrande, ARC-easy, ARC-challenge, and OpenBook QA benchmarks. We evaluated intermediate checkpoints up to the fully trained checkpoint for TinyLlama 1.1B. Among these, we utilized the 105B intermediate checkpoint to study an mnder-trained model.
Few-shot Accuracy
Models
N-emb
Dataset
Nuoken
LD
HS
PQ
WG
ARC-e
ARC-c
OB
Gemma 2B
1.99B
Unreleased
3T
63.13
71.38
78.13
65.04
72.26
41.89
40.20
61.72
105B
43.26
42.23
66.81
53.35
44.74
23.21
29.20
43.26
503B
48.92
49.56
69.42
55.80
48.32
26.54
31.40
47.14
SlimPajama
TinyLlama I.IB
Starcoderdata
53.00
52.52
69.91
55.96
52.36
27.82
33.40
2T
53.33
54.63
70.67
56.83
54.67
28.07
33.40
50.23
3T
58.82
59.20
73.29
59.12
55.35
30.12
36.00
53.13
Pythia IB
0.81B
Pile
300B
$7.52
49.10
70.40
52.80
$1.89
26.71
33.40
48.83
compare This diversity offers several benefits. First, with three versions of recursive models, their performance based on the number of trainable parameters. Notably, the comparison between the recursive Gemma and the pretrained TinyLlama and Pythia models highlights that leveraging well-trained model weights can lead to a superior Recursive Transformer of equivalent size, even with substantially lower uptraining costs. Second, by utilizing models ranging from under-trained (e.g., TinyLlama) to significantly over-trained (e.g., Gemma), we can gain insights into the uptraining costs required for Recursive Transformers closely match the performance pretrained models. Finally, the diversity in pretraining datasets allows us to observe how Recursive Transformers perform when faced with distribution shifts in the uptraining dataset. Table Section 3.2 presents the evaluation results obtained after uptraining each of the pretrained models. While Tinyl lama readily improves its performance due to uptraining on the same dataset, Gemma and Pythia show decline in few-shot performance with SlimPajama uptraining, which can be attributed to the differences in data distribution and the lower quality the uptraining dataset.
24

---

Published as a conference paper at ICLR 2025
G EXPERIMENTAL SETUP
Uptraining setting To convert vanilla Transformers into Recursive Transformers, we conducted further uptraining on either 15 billion or 60 billion tokens from the SlimPajama dataset (Soboleva et al., 2023). SlimPajama is an open-source dataset designed for training large language models, which is created by cleaning and deduplicating the RedPajama dataset (Computer, 2023). The source data primarily consists of web-crawled data, along with data from Github, books, Arxiv, Wikipedia, and Stack Exchange We employed the HuggingFace training framework (Wolf et al., 2020) and enhanced memory efficiency through the Zero Redundancy Optimizer (ZeRO) (Rajbhandari et al., 2020) from the DeepSpeed library (Rasley et al., 2020), along with mixed precision training. The context length was set to 2048, and the batch size was approximately 2 million tokens. We used the Adam W optimizer (Loshchilov & Hutter, 2019) with a learning rate of 2e-4, utilizing cosine annealing learning rate scheduler (Loshchilov & Hutter, 2017). Additionally, we set warmup steps to 200 for 15 billion token training and 800 for 60 billion token training. Eight H100 GPUs were used for the training.
Early-exit training setting Similar to the uptraining process, we used the SlimPajama dataset to enable models to predict next tokens at intermediate loops. Models with two looping blocks total of two exit points, whereas models with three blocks were trained on three exit points. We explored various strategies, but by default, we continued training on an additional 15 billion tokens, starting from the uptrained Recursive Transformers. We also utilized eight H100 GPUs and maintained consistent configurations with the uptraining settings, including batch size, context length, and learning rates.
underwent additional training
Evaluation setting We evaluated perplexity on test sets from three language modeling datasets: SlimPajama, RedPajama, and PG19 (Rae et al., 2019). Additionally, we used the Language Model Evaluation Harness framework (Gao et al., 2023) to evaluate accuracy on seven few-shot tasks: LAMBADA (Paperno al., 2016), HellaSwag (Zellers al., 2019), PIQA (Bisk et al., 2020), Wino- Grande (Sakaguchi et al., 2020). ARC-easy and ARC-challenge (Clark et al., 2018), and Open- BookQA (Mihaylov et al., 2018). We adhered to the standard number of shots specified by the evaluation framework for each dataset. For few-shot datasets, excluding LAMBADA and Wino- Grande, we normalized accuracy by the byte length of the target string. All evaluation performance measurements were conducted using single H100 GPU.
Throughput measurement settings To present the hypothetical generation speeds of our Recursive Transformers, we prepared two key elements: per-token generation time and exit trajectory datasets. Firstly, we measured the generation time under various model configurations using dummy weights and inputs. We measured the time for each component, such as embedding matrices, Transformer blocks, and the classifier head. Note that, for simplicity, throughput comparisons were based solely on the time spent within the Transformer block components. We tested two settings of prefix and decoding lengths (512/200) 2048 and 64/ / 256), calculating the per-token time by dividing the total elapsed time by the decoding length. Using single A100 40GiB GPU, we measured these decoding times across different batch sizes, until an out-of-memory error occurred or under a specific memory constraint was reached. To obtain exit trajectory data, we assumed an oracle-exiting approach, where all tokens could exit at intermediate loops if intermediate predictions matched the final loop's s prediction. Since our models are not finetuned on any specific downstream tasks, we simulated the generation of language modeling datasets as if they were generated by our models. For simplicity, we assumed a queue of 20K samples with varying context lengths, rather than considering their arrival in static dynamic time intervals. With these two datasets, we present the hypothetical throughput of Recursive Transformers under various simulation scenarios.
25

---

Published as a conference paper at ICLR 2025
H PERFORMANCE OF FULL-SIZE MODEL BASELINES
Our ultimate goal is for the Recursive Transformer to achieve performance comparable to the original, full-size pretrained model, but using the least amount of uptraining tokens possible. This is challenging because our recursive models have substantially fewer parameters, and model capacity is primarily determined by model size. However, prior works have suggested that FL.PPs also play role in influencing model performance (Dehghani et al., 2019; 2022; Goyal et al.. 2024). Consequently, by recursively applying the function, we anticipate that with effective initialization techniques or training strategies, it might be possible to attain performance that closely approaches that of the full-size model.
However, the uptraining dataset itself can hinder this goal. Specifically, poor quality of the uptraining dataset or a significant distribution shift from the pretraining dataset can negatively impact perfor- mance. Indeed, as shown in Table 4, the Gemma model exhibited a performance decrease across all few-shot benchmarks after uptraining on SlimPajama. Conversely, Tinyl.lama, where the uptraining and pretraining datasets are both SlimPajama, consistently showed performance improvements.
Considering these results and our original goal, we adopted the following full-size model baselines: the original pretrained model for Liaaa and vanilla models uptrained with the same cost their recursive counterparts for Gemma and Pythia.
Table 4: Uptraining pretrained models on datasets that differ significantly in quality or distribution from their pretraining data can lead decreased performance. We evaluated models after uptraining on the SlimPajama dataset. We measured perplexity on test sets of SlimPajama, RedPajama, and PG19, and shotshot accuracy on LAMBADA, HellaSwag, PIQA, WinoGrande, ARC-easy, ARC- challenge, and OpenBookQA benchmarks.
Uptrain
Perplexity
+
Few-shot Accuracy
Models
N-emb
PT
Nesk
SlimP
RedP
PG19
LD
HS
PQ
WG
ARC-e
ARC-c
OB
Avg
1.99B
11.46
8.18
13.52
63.1
71.4
78.1
65.0
72.3
41.9
40.2
61.7
Gemma
1.99B
15B
10.76
8.47
13.08
63.5
68.5
77.0
63.5
67.6
38.1
42.6
60.1
1.99B
60B
10.58
8.44
12.71
60.3
67.9
76.9
63.5
64.9
37.2
39.6
58.6
0.97B
12.26
9.37
11.94
43.3
42.2
66.8
53.4
44.7
23.2
29.2
43.3
TinyLlama
0.97B
/
15B
9.87
8.24
10.73
49.2
46.3
68.8
54.0
48.2
26.0
32.2
46.4
0.97B
60B
9.59
8.12
10.42
51.6
48.8
68.6
54.1
49.9
26.2
32.8
47.4
0.81B
15.68
9.90
12.05
57.5
49.1
70.4
52.8
51.9
26.7
33.4
48.8
Pythia
0.81B
15B
13.46
9.95
13.38
55.0
49.0
71.0
53.6
51.8
28.2
32.8
48.8
0.81B
60B
12.83
9.76
13.57
53.0
50.2
71.1
54.8
51.9
27.7
31.6
48.6
I EXPANDED RESULTS OF INITIALIZATION METHODS FOR LOOPED LAYERS
Ablation study of Stepwise method We initially hypothesized that the Stepwise method's 's per- formance could be significantly influenced by the specific rule used for layer selection from the pretrained model. To investigate this, we conducted controlled experiment (illustrated in Figure 1a), where layers were selected certain intervals starting from the first layer. We then varied whether the final layer of the pretrained model was included in the initialization or not. While a Pythia model showed no discernible differences in training loss few-shot performance, other models like Gemma exhibited markedly superior results when both the first and last layers were preserved. This observation aligns well with prior work suggesting that maintaining the weights of the first and last layers during depth up-scaling for L.MSs can yield performance benefits (Kim et al., 2024).
Ablation study of Average method The Average initialization method exhibited notably poor perfor- mance, particularly when applied to the Gemma model. We hypothesized that this could be attributed to instability the model's learned distribution, potentially arising from averaging normalization layer weights. Relatedly, several studies (Csordás et al., 2024; Shim et al., 2024; Mohtashami et al., models. To investi-
2023) have explored the careful design of layer normalization in parameter-shane gate this further, experimented with three different methods for initializing normalization weights, as outlined in Figure Ilb: averaging weights (Norm-avg), selecting weights from a single layer (Norm- choice), and zero initialization (Norm-zero). The performance trend observed among these methods
26

---

Published as a conference paper at ICLR 2025
varied across different model architectures. However, zero initialization of normalization layers resulted huge performance drop in certain architectures like TinyLlama and Pythia. Conversely, observed big difference between averaging and single-layer selection, suggesting that any form of distillation of the normalization weights appears to be sufficient for maintaining performance.
Gemma
TinyLiama
Pythia
2.9
Fixed-ands
3.9
Nommavg
Fixed-start
Training Loss
Nomm-zer
3.1
2.5
2.7
2.3 0B
2.3
SB
10B
15B
08
5B
108
15B
of Tokens
of Tokens
(a) Stepwise ablations
(b) Average ablations
Figure 1I: Training loss curves stepwise and average initialization variants across three models with two blocks. (a) "Fixed-start" indicates that the first layer of the pretrained model is selected initially, and subsequent layers are repeatedly chosen at a fixed interval. "Fixed-ends" means that the first and last layers are included, and intermediate layers selected at specific step intervals. (b) When tying LayerNorm weights, we consider whether to average the weights (LN-avg), select single weight (LN-choice), or use zero initialization (LN-zero).
Overall comparison of training perplexity Figure 12 presents comparative analysis of training loss across three model architectures and varying looping blocks, incorporating our proposed initial- ization methodologies. To set upper bound performance, uptrained on SlimPajama, accounting for the distribution shift between uptraining and pretraining data. Additionally, we trained Recursive Transformer from random initialization, ensuring its exclusive reliance on the recursive architecture without leveraging any pretrained weights. While some variance was observed across architectures, all proposed methods utilizing pretrained model weights demonstrated significantly superior performance compared to random initialization. Notably, the Stepwise method consistently achieved the best performance across diverse settings. Although the full-size model's performance was considerably higher, bridging this gap with only 15 billion tokens of uptraining represents a remarkable achievement.
utilized full-size model further
Full-size
Stepwise
Average
Lower
Random
Training Loss
3.0
3.0
2.6
2.6
2.2
2.2
2.2
22-
08
58
10B
10B
158
08
58
108
10B
158
of Tokens
of Tokens
#
Tokens
of Tokens
(a) Gemma (2 blocks)
(b) Gemma (3 blocks)
(c) Tamy (2 blocks)
(d) Pythia (2 blocks)
Training loss for Recursive Transformers using various initialization. We omitted a separate curve for the full-size TaayL model, used the original pretrained model as the full-size model since both pretraining and uptraining datasets are same as the SlimPajama dataset. Refer Section 3.2 for more details.
Figure 12:
27

---

Published as a conference paper at ICLR 2025
Overall comparison of few-shot performance Few-shot performance exhibited consistent trend with training perplexity. Table provides a comparative summary of the proposed looping initializa- tion methods against the full-size model, the reduced-size model, and Recursive Transformers utilizing random initialization. Moreover, Figure 13 visually illustrates the performance differences across different datasets. Notably, the Stepwise method consistently demonstrated the best performance, showing a performance improvement of up to 14.1 %6 compared to random initialization.
Table 5: Evaluation results of various initialization methods for looped layers. We indicate whether pretrained weights are used and the number of uptraining tokens. Perplexity is evaluated on test sets of three language modeling datasets, and accuracy is evaluated on seven few-shot benchmarks. Delta values (A) show improvements over random initialization. We highlight the configurations that demonstrate the best performance.
Uptrain
Looping
Perplexity
Few-shot Accurncy
Models
N-emb
PT
Noa
Block
Init
SlimP
RedP
PG19
LD
HS
PQ
WG
ARC-e
ARC-c
OB
Avg
A
1.998
15B
10.76
8.47
13.08
68.5
77.0
67.6
38.1
42.6
60.1
0.99B
X
15B
22,63
20.03
32.60
28.9
31.6
63.1
52.3
41.2
23.5
27.8
38.2
0.66B
24.44
21.69
27.2
30.6
63.8
40.6
22.0
27.0
37.4
0.99B
15B
Step
12.85
10.29
16.21
53.0
57.3
73.2
56.2
56.1
29.2
36.6
51.7
+14.1
Gemma
0,991
15B
Avg
15.15
12.57
19.86
43.6
47.4
70.4
52.6
50.5
27.8
34.4
46.7
+9.1
0.99B
15B
Lower
15.03
12:46
19.63
42.5
48.0
71.0
54.6
52.2
27.7
33.8
47.1
+9.5
0.99B
15B
Rand
22,66
20.06
32.86
27.4
31.6
63.4
50.5
39.7
21.9
28.8
37.6
0.668
Step
14.75
12.10
19.32
69.8
55.8
52.7
27.9
33.6
47.8
+9.9
0.66B
15B
Avg
17.45
14.65
23.63
39.4
39.0
66.6
48.7
46-5
24.7
31.8
42.4
+4.5
0.66B
15B
Lower
15.96
13.24
20.90
41.9
43.2
70.0
52.6
49.5
26.6
31.6
45.0
+7.1
0.66B
15B
Rand
22.67
20.09
32.77
28.1
31.4
63.8
51.1
41.0
23.0
26.6
37.9
12.26
23.2
29.2
0.48B
x 15B
16.61
15.66
20.27
22.3
30.0
60.9
50.6
37.0
23.0
28.0
36.0
TamaLlama
0.48B
Step
9.89
39.8
24.9
42.7
+6.2
0.48B
15B
Avg
11.86
10.29
13:42
38.6
39
66.1
52.8
42.7
25.4
30.6
42.2
+5.7
22.1
27.8
0.48B
x 15B
Rand
16.14
15.11
19.55
24.7
30.7
61.2
50.6
36.4
22.6
29.2
36.5
0.8
15B
13.46
9.95
13.38
55.0
49.0
71.0
53.6
51.8
28.2
32.8
48.8
0.40B
x 15B
25.69
20,00
32.08
24.3
30.0
61.9
50.7
38.3
23.3
26.0
36.2
Pythia
0.40B
15B
Step
16.38
12.37
17.74
43.4
40.5
67.4
50.8
46.3
25.7
30.0
43.5
+7.3
0.40B
Arg
12.76
18.63
68.2
25.1
29.8
43.3
+7.1
0.40B
1SB
Lower
17.04
12.62
18.44
43.9
39.2
66.3
53.4
45.4
25.8
31.2
43.6
+7.4
0.40B
x 15B
Rand
24.45
18.93
29.63
25.2
30.2
62.1
51.1
39.2
22.4
23.6
36.2
28

---

Published as a conference paper at ICLR 2025
Stepwise Avarage Random
Accuracy
0
Average
LAMBADA
HellaSwag
PIQA
WinoGrande
ARCC-asy
ARC-challenge
OpenBookQA
(a) Recursive Gemma with 2 blocks
Steeswize
Average
80
Random
Apcuracy
Average
LAMBADA
HellaSwag
WinoGrande
ARC-easy
ARC-challenge
OpenBookOA
(b) Recursive Gemma with 3 blocks
Sieprwize
Lower
Handom
Accuracy
20
Average
LAMBADA
HellaSwag
WinoGrande
ARC-easy
ARC-challenge
OpenBookOA
(c) Recursive TinyL with blocks
80
Stepwise
Average
Accuracy
40
Average
LAMBADA
HellaSwag
PIOA
WinoGrande
ARC-easy
ARC-challenge
OpenBookQA
(d) Recursive Pythia with blocks
Figure 13: Few-shot performance on seven benchmarks and their average accuracy based on four looping initialization methods. Full-size model performance is represented by gray dotted line.
29

---

Published as a conference paper at ICLR 2025
Comparison across various base model sizes We observed consistent superiority of Stepwise initialization strategy for recursive conversion across both 1B and 2B model scales. To further evaluate wide range of base model sizes, we additionally experimented with two smaller model sizes, Pythia 410M and 160M. Since we uptrained models on the Pile dataset (Gao et al., 2020), the pretraining corpus of the original Pythia model, we the original model performance as the baseline for comparison. The results in Table 6 further validate the superior performance of the Stepwise method for looped layer initialization. These findings reinforce the robustness of our key observations regarding initialization methods for recursive conversion, complementing our original extensive experiments.
Table Comparison between initialization methods for looped layers on Pythia 4 410M and 160M. Uptraining was performed using the Pile dataset, which was also used for pretraining the original Pythia model. In light of the inherent randomness in few-shot accuracy, a comparison based on the perplexity (PPL) would provide more stable measure of performance.
Uptrain
Looping
P
Few-shot Accuracy
N-emb
PT
Niok
Block
Init
Pile
LD
HS
PQ
WG
ARC-e
ARC-c
Avg
300M
/
44.96
40.97
66.97
53.28
44.40
25.51
30.20
43.76
150M
/
15B
Step
11.03
43.41
35.59
64.58
53.04
41.58
23.81
28.80
41.54
150M
/
15B
Lower
11.47
42.98
34.32
63.93
52.41
42.34
24.15
25.00
40.73
150M
/
15B
2
Avg
11,55
39.84
34.17
64.31
52.25
41,04
24.66
26.60
40.41
13.53
30.67
58.22
48.62
36.62
25.00
28.60
34.47
43M
/
15B
Step Lower
15.93
21.02
29.28
60.01
48.93
37.92
23.98
28.00
35.59
43M
/
15B
16.19
21.46
29.61
59.90
50.67
38.52
22.95
28.00
35.87
43M
/
15B
Avg
16.12
22.36
29.07
60.17
49.96
37.24
23.29
26.60
35.53
Individual contributions of leveraging pretrained weights and recursive patterns To understand the performance of our Recursive Transformer, we established two non-recursive baselines: full-size model and reduced-size model. The reduced size model performance is meant to serve as lower bound which techniques that are made possible by our approach and (2) leveraging pretrained layers. To further ablate the effect of each two components, we conducted experiments using the Pythia 410M model presented in Table 7. Intuitively, we observed significant performance gains by leveraging pretrained layers, with further improvement achieved through recursion. We believe this additional experiment provides valuable insight into the performance contributions of the two approaches proposed for constructing Recursive Transformers.
use to better judge the efficacy of (1) unique looping and parameter sharing
Table 7: Performance of recursive and baseline models with Pythia 410M to investigate the individual contributions of pretraining layers and looping strategy. Uptraining performed using the Pile dataset (Gao et al., 2020), which was also used for pretraining the original Pythia model.
Uptrain
Looping
PPL.]
Few-shot Accuracy
N-emh
PT
No.k
Block
Init
Pile
L.D
HS
PQ
WG
ARC-e
ARC-c
OB
Avg
300M
/
44.96
40.97
66.97
53.28
44.40
25,51
30.20
43.76
150M
X
15B
14.11
31.48
29.53
61.37
52.49
39.14
22.44
27.00
37.63
150M
X
15B
2
13.81
31,55
29.94
62.30
50.88
40.28
23,98
28.20
38.02
150M
Step
11.48
40.48
34.19
63.42
50.99
41.84
23.12
28.40
40.35
150M
V
15B
2
Step
11.03
43.41
35.59
64.58
53.04
41.58
23.81
28.80
41.54
30

---

Published as a conference paper at ICLR 2025
J EXPANDED RESULTS OF RELAXED RECURSIVE TRANSFORMERS
Training perplexity changes with LoRA modules Figure 14 illustrates the changes in training loss after incorporating the layer-wise LoRA modules. The Average and Lower initialization methods, when coupled with our proposed SVD-based initialization of the LoRA modules, demonstrated significantly enhanced benefits. In particular, the Relaxed Recursive Transformer employing the Average method consistently outperformed the others. This suggests that it is considerably easier to leam the difference between the original pretrained weights and the averaged looped weights using low-rank matrices.
Stepwise
Average
Lower
3.1
3.0
3.1
LoRA v LORAX
LoRA LORAX
LORA LORAX
Training Loss
2.8
5B
10B
15B
0B
10B
15B
of Tokens
# of Tokens
of Tokens
(a) Gemma
(b) TinyLlama
(c) Pythia
Figure 14: Comparison of training loss for recursive and relaxed recursive models. All recursive models utilize two looping blocks, and the LoRA rank is set to 512. The SVD initialization method is used for LoRA modules.
Comparison between SVD and zero initialization The utilization of layer-wise LoRA modules enhances model capacity by introducing additional parameters and relaxation, thereby potentially improving performance. As depicted in Figure 15, SVD initialization significantly amplified these performance gains compared to standard zero initialization. However, an interesting exception was observed with the Stepwise method, where the SVD initialized LoRA module surprisingly led to a performance degradation in Gemma and TinyLlama. This appears to be attributed to the LoRA rank being insufficient adequately approximate the low-rank deltas across layers, resulting in initialization at sub-optimal point.
LoRA
Zero Init
Init
LoRA
SVD
Init
Few-shot Accuracy
55
45
46
44
40
42
35
Stepwise
Average
Stepwise
Average
Stepwise
Average
(a) Gemma
(b) TinyLlama
(c) Pythia
Figure 15: Comparison of average few-shot accuracy between zero and SVD initialization methods for layer-wise LoRA across three models. Performance gains due to LoRA relaxation are indicated by hatched bars, while cases where performance is lower than the recursive model without LoRAs are represented by dotted lines.
31

---

Published as a conference paper at ICLR 2025
Ablation study the LoRA rank values Our proposed SVD initialization ensures that the Relaxed Recursive Transformer can function as an interpolation between vanilla and Recursive Transformers. The approximation accuracy of SVD is directly influenced by the LoRA rank value; a higher rank leads to improved restoration of the pretrained model weights. In Figure 16, we present a summary of the performance changes observed in the relaxed models by varying the LoRA ranks. As expected, for the Average and Lower looping initialization methods, larger rank value results in enhanced performance. The Stepwise method, consistent with previous experimental findings, exhibited U-shaped trend: with lower rank, it behaves similarly to random initialization, resulting in slight performance increase. However, with higher, the approximation becomes more accurate, leading to further increase in performance.
Stepwise
Average
Lower
55
Faw-shol Accuracy
42
50
45
43
64
128
256
64 128
256
512
LoRA Rank
LoRA Rank
LoRA Rank
(a) Gemma
(b) TinyL.ama
(c) Pythia
Figure 16: Performance comparison with varying LoRA ranks under different initialization methods
for looped layers. All LoRA weights initialized using our proposed SVD initialization method.
We further experimented with assigning different ranks to LoRA modules associated with each linear layer. Given the computational overhead inherent to LoRA modules, allocating varying ranks to each module can offer optimal balance between performance and computational efficiency. The experimental results Table reveal a strong correlation between performance and overall model sizes. Due to the substantial hidden dimension of the linear weight within the FFN layer, reducing its rank led to the most significant performance drop. Conversely, the relatively smaller size of other attention weights resulted in less performance drops. An intriguing observation is the comparable performance maintained even with minimal relaxation of key-value weights (achieved through small ranks), despite the inherent strong sharing of key-value caches in the Multi-Query attention structure (Ainslie et al., 2023).
Table 8: Evaluation results of relaxed recursive Gemma models with varying LoRA ranks applied Transformer components. We adjusted the LoRA ranks attached to query, key-value, out, and FFN linear weights. Non-embedding parameter sizes include both the base layer parameters and the attached LoRA weights.
Uptrain Looping
LoRA
Perplexity L
Few-shot Accuracy 1
Nick
Slimm
RedP
PQ
ARC-e
ARC-e
Avg
1.99R
15B
10.76
8.47
13.08
63.5
68.5
77.0
63.5
67.6
38.1
42.6
60.1
0.99B
15B
22.63
20.03
32.60
28.9
31.6
63.1
52.3
41.2
22.5
27.8
38.2
Avg
58.2
60.7
73.7
59.0
57.6
32.1
54.2
1,28B
15B
Avg
128
256
128
256
SVD
12.27
9.81
15.10
57.4
60.2
72.5
58.9
58.1
32.6
37.8
53.9
1.29B
15B
Avg
256
128
256
256
SVD
12.33
9.90
15.25
58.5
59.7
73.3
58.3
56.6
32.0
40.0
54.1
1.18B
15B
Avg
256
256
256
128
SVD
12.56
10.12
15.59
57.0
58.7
73.0
57.4
57.0
31.6
38.2
53.3
1.27B
15B
Avg
128
128
128
256
SVD
12.36
9.92
15.31
57.2
59.2
73.1
57.3
58.0
32.2
38.6
53.7
Avg
10.07
56.1
57.1
30.7
$2.5
1.14B
15B
Avg
64
128
64
128
SVD
12.61
10.14
15.69
55.0
57.8
73.0
57.5
56.7
30,9
38.8
52.8
Avg
10.18
55.5
72.7
57.0
56.9
30.1
1.08B
15B
Avg
128
128
128
64
SVD
12.80
10.33
15.97
55.3
56.7
72.9
57.7
55.0
29.6
36.0
51.9
Avg
73.0
56.5
30.5
52.3
32

---

Published as a conference paper at ICLR 2025
Overall performance comparison of Relaxed Recursive Transformers A comprehensive per- formance comparison is presented in Table 9. This encompasses an evaluation of the performance across three models and various looping initialization methods, considering the degree of relaxation induced by the layer-wise LoRA module. The configuration utilizing the Average method for looped layer initialization, paired with SVD) initialization for the LoRA module, consistently outperformed all other baselines, Furthermore, performance clearly improved with increasing rank.
Table 9: Evaluation results of relaxed recursive models varying LoRA ranks. Delta (A) represent the accuracy differences between relaxed and non-relaxed models using the same looping initialization.
Uptrain
Looping
LoRA
Perplesity
Few-shot Accurucy
Models
N-emb
PT
Nos
Block
Init
Rank
Init
SlimP
RedP
PG19
LD
HS
PQ
WG
ARC-c
ARC-c
OB
Avg
A
1.99B
/ 15B
10.76
8.47
13.08
63.5
68.5
77.0
63.5
67.6
38.1
42.6
60.1
0.998
x 15B
22.63
20.03
32.60
28.9
31,6
63,1
52.3
41.2
22.5
27.8
38.2
0.998
15B
Rand
22.66
20.06
32.86
27.4
63.4
21.9
37.6
0.998
/ 18B
2
Step
12.85
10.29
16.21
53.0
57.3
73.2
56.2
56.1
29.2
36.6
51.7
1,078
15B
Step
SVD
12.76
10.21
15.99
52.1
73.0
56.9
28.9
51.8 +0.1
1.15H
18B
2
Step
128
SVD
13.44
10.80
16.98
50.5
53.0
71.5
54.4
55,9
29.3
34.8
49.9
1..8
1.30B
15B
Step
256
SVD
14.02
11.44
18.09
46.1
49.1
71.8
53.2
52.8
27.8
33.4
47.8
3,99
1.608
15B
Step
512
SVD
13.13
10.66
16.63
53.0
54.3
72.1
54.9
54.8 8
28.8
35.4
50.5
11.
1660B
/ 15B
2
Step
512
Zero
12:46
9.97
15.58
54.9
58.8
74.0
58.1
58.8
30.6
36.6
53.1
+1.4
Gemma
0.998
/
15B
Avg
15.15
12.57
19.86
43.6
47.4
70.4
50.5
27.8
46.7
1,07B
15B
2
Avg
64
SVD
12.83
10.35
16.02
55.9
56.8
72.5
56.8
55.7
30.6
36.2
52.1 +5.4
1.158
15B
Avg
128
SVD
12.52
10.07
15.51
56.1
58.2
72.3
55.8
57.1
30.7
37.2
52.5
+5.8
1.300
15B
Avg
256
SVD
12.10
9.71
14.89
58.2
60.7
73.7
59.0
57.6
32.1
38,0
54.2
+7,5
1,608
15B
Avg
512
SVD
11.83
9.46
14.57
59.3
62.8
74.0
61.6
60.1
32,9
37.6
55.5
+8.8
15B
Avg
Zero
13.78
11.31
17.71
49.8
52.4
71.7
51.2
29.4
49.0
0.998
/
15B
2
Lower
15.03
12.46
19.63
42.5
48.0
71.0
54.6
52.2
27.7
33.8
47.1
1.07B
15B
Lower
SVD
14.21
11.77
18:40
47.5
70.9
29.2
48.9
+1.8
1.150
15B
Lower
128
SVD
14.23
11.83
18.49
48.0
50.5
72.0
56.8
54.4
27.5
33.4
48.9
+1.8
1.30B
/
15B
Lower
256
SVD
13.51
11.06
17.30
53.1
53.7
71.8
57.4
52.5
28:7
35.2
50.3
+3.2
15B
Lower
SVD
12.54
10:22
15.90
57.1
58.2
73.7
58.6
57.6
31.5
53.2
.6.1
1660
/ 15B
2
Lower
512
Zero
13.95
11.59
18.02
48.4
52.1
71.9
55.7
54.9
28.8
34.6
49.5
+2.4
0.97B
12:26
9:37
11.94
43.3
42.2
66.8
53.4
44.7
23.2
29:2
43.3
0:18B
X 15B
16.61
15.66
20.27
22.3
30.0
60.9
50.6
37.0
23.0
28.0
36.0
0.488
15B
Rand
16.14
15.11
19.55
24.7
61.2
36.4
22.6
36.5
0.48B
/ 15B
2
Step
11.61
9.89
13,000
39.6
39.8
66.5
$2.9
44.3
24.9
30.6
42.7
0.538
/ 15B
Step
SVD
12:10
10.40
13:75
38.9
38.3
65.2
51.5
42.0
26.0
31.0
41.9
0..8
0.580
15B
Step
128
SVD
12.41
10.72
14.10
36.8
37.4
64.7
53.4
42.2
24.7
30.4
41.4
-1.
0,688
15B
Step
256
SVD
11.96
10.35
13.48
38.9
38.3
65.8
51.9
43.1
25.4
29.8
41.9
0..8
0.860
15B
Step
SVD
11.33
12.61
42.2
40.9
67.7
45.0
25.3
43.2
+0.5
0.86H
15B
Step
512
7em
11.24
9.60
12.56
42.0
41.0
67.4
52.2
44.5
25.9
31.2
43.4
+0.7
TinyLlama
0.18B
/
15B
2
Avg
11.86
10.29
13:42
38.6
66.1
42.7
25.4
42.2
0.530
15B
Avg
64
SVD
11.22
9.66
12.51
41.8
41.6
67.0
53.3
43.9
24.7
31.2
43.4
+1.2
0,58B
15B
Avg
128
SVD
10:99
9.45
12.21
43.2
42.1
68.3
53.2
44.8
25.9
30.4
44.0
+1.8
0.680
15B
Avg
10.71
11.82
68.1
53.5
44.4
25.7
44.2
+2.0
0,86B
15B
Avg
512
SVD
10.46
8.92
11.50
46.0
44.1
68.2
$3.0
45.8
25,1
31.2
44,8
+2.6
0.868
15B
Avg
512
Zeno
11.28
9.75
12.69
41.5
41.0
66.8
53.2
44.8
25.5
31.2
43.4
+1.
0.488
V
15B
2
Lower
14:67
12.67
16.68
31.9
32.3
62.6
$2.0
39.1
22.1
27.8
38.5
0.538
/
158
Lower
SVD
13.68
11:77
15.48
35.5
34.0
63.8
24.6
39.5
+1,2
0,58B
15B
Lower
128
SVD
13.00
11.14
14.61
37.6
35.4
65.3
51.5
40.4
24.5
27.6
40.3
+2.0
0.68B
15B
Lower
256
SVD
12.14
10.39
13:59
40.0
37.7
66.1
52.6
42.5
24.8
30.0
42.0
+3.7
0,86B
15B
Lower
11.31
12.19
43,2
40.5
66.0
50.8
43.9
24.8
42.8
+4.5
0.868
15B
2
Lower
512
Zero
14.56
12.69
16.57
21.2
32,9
63.9
52.6
39.5
22,9
27.8
37.3
-1.0
0.810
15B
13.46
13.38
55.0
49.0
71.0
51.8
28.2
48.8
0.40H
X
15B
25.69
20.00
32.08
24.3
30.0
61.9
50.7
38..3
22.3
26.0
36.2
0.40B
15B
Rand
24.45
18.93
29.63
25.2
30.2
62.1
51.1
39.2
22.4
23.6
36.2
0.408
/
15B
2
Step
16.38
12:37
17:74
43.4
40.5
67
50.8
46.3
25.7
30.01
43.5
0.44B
15B
Stop
64
SVD
16.44
12.44
17.89
43.7
40.4
66.5
52.9
46.5
26.2
28.8
43.6 +0,1
0.488
15B
Step
128
SVD
16.63
12.61
18.35
42.4
39.3
68.0
51.5
46.3
26.7
30.6
43.5 +0.0
0.55B
15B
Step
256
SVD
16.54
12.61
18.39
42.8
39.1
67.2
53.7
46.4
25.9
27.8
43.3
-0,2
0.708
15B
Step
SVD
15.68
11.88
17:25
45.4
68.5
46.7
25.44
44.4
+0.9
0.70B
15B
Step
512
Zero
15.88
12.01
17.16
45.5
41.8
68.0
52.6
46.6
26.3
30.0
44.4
+0.9
Pythia
0.40H
/
15B
Avg
16.76
12.76
18.63
43.6
39.1
68.2
$1.9
45.4
25.1
29.8
43.3
0.44B
15B
Avg
64
SVD
16.03
12.19
17.59
45.8
40.9
67.3
50.0
45.8
25.5
31.8
43.9
+0.6
0.48B
15B
Avg
128
SVD
15.67
11.93
17.10
46.9
41.9
67.4
51.2
45.4
24.8
31.2
44.1
+0.8
0.55B
15B
Avg
15.22
48.5
51.4
25555
0.708
15B
Avg
512
SVD
14.70
11.07
15.71
50.2
44.7
68.2
51.6
47.6
25.4
31.2
45.6
+2.3
0.70B
/ 15B
2
Avg
512
Zero
15.97
12.14
17.65
45.7
41.5
68.1
51.7
46.5
25.7
30.0
44.2
+0,9
0.408
/
15B
2
Lower
17:04
12.62
18:44
43.9
39.2
66.3
53.4
45.4
25.8
31.2
43.6
0.448
15B
2
Lower
64
SVD
17:03
12.78
18:73
44.1
38.3
66.9
51.9
45.4
24.7
30.8
43.2
-0.4
0.488
15B
Lower
SVD
16.63
12.49
18.17
45.2
39.2
66.8
51.0
45.6
24.9
43.2
0,558
15B
2
Lowar
256
SVD
13.93
11.99
17.30
47.6
41.4
68.3
53.2
46.0
25.8
31.0
44.8
+1.2
0.708
158
Lower
SVD
15.11
11:34
16:07
50.2
43.5
67.8
47.2
25.2
32.0
45.4
+1.8
0.70B
15B
2
Lower
512
Zero
16.45
12:25
17.76
45.2
40.4
66.4
54.5
45.8
25.9
32.6
44.4
+0.8
33

---

Published as a conference paper at ICLR 2025
K ALTERNATIVE APPROACHES TO RELAXATION OF PARAMETER SHARING
To mitigate the restrictive weight-tying inherent in parameter sharing, we employed LoRA modules as discussed in $2.3, similar to prior work (Ge et al., 2022; Shim et al., 2024). However, efficiently computing multiple LoRA modules requires specialized kernels and sequential computations of the base layers and LoRA modules, incurring computational overhead. Consequently, explored layer-specific prompts (Liu et al., 2021) as an alternative. This approach integrates prompts specific to each layer as prefix tokens, generating layer-wise key and value states for self-attention, and is significantly more amenable to parallel computation.
Table 10 summarizes performance of the prefix tuning method. While offering computational advan- tages, its reliance small, learnable prompts resulted in limited performance gains. Additionally, without leveraging the original pretrained weights, performance was significantly lower (52.1%% % vs. 47.6% with the Average method in 1.07B model size). Future work will explore enhancing the effec- tiveness of this parallel approach, as well as other strategies such as bias term-based relaxation (Ge et al., 2022).
Table 10: Evaluation results of relaxation through prefix tuning methods. Prefix length denotes the sequence length of trainable vectors used to generate key-value prompts in each self-attention layer. Non-embedding parameter sizes include the sizes of these trainable prefixes. Delta (A) represent the accuracy differences between non-relaxed models and their corresponding prefix-tuned models using the same looping initialization.
Uptrain|
Looping
Prefix
Perplesity
Few-shot Accuracy
Models
N-emb
PT
Noa
Black
Init
Len
Size
SlimP
RedP
PG19
1.1)
HS
PQ
WG
ARC-e
ARC-c
OB
Avg
A
1.99B
10.76
13.08
68.5
38.1
0.998
X
15B
22.63
20.03
32,60
28.9
31.6
63.1
52.3
41.2
22.5
27.8
38.2
0,998
2
12.85
10.29
16.21
53.0
57.3
56.1
29.2
36.6
51.7
1.008
158
2
Step
256
9.4M
12.62
10.06
15.80
53.5
58.3
73.9
57.6
$7.5
29.3
35.6
52.2
+0.5
1.018
158
2
Step
512
18 98.
12.67
10.10
15.85
54.1
57.8
73.8
58.4
57.2
28.7
35.8
52.3
+0.6
Step
1024
12.89
10.34
53.5
57.1
72.4
56.9
1.07B
15B
2
Step
2048
75.5M
12.75
10.21
16.09
55.0
57.3
73.3
58.2
56.8
29.2
37.8
52.5
+0.8
Gemma
0.998
2
Avg
15.15
47.4
50.5
46.7
1.008
15H
2
Avg
256
9.4M
14.85
12.31
19.41
46.9
48.3
70.4
52.7
51 4
27.2
34.0
47.3
0..6
1.01B
15B
2
Avg
512
18.9M
15.23
12.64
19.98
44.5
47.2
70.7
54.5
-49.5
28.0
33.2
46.8
+0,1
1.098
15B
2
Avg
1024
37.7M
14.60
12.02
18.
46.9
49.7
71.1
51.00
28.6
47.7
+1.0
1.07B
15B
2
Avg
2048
75.5M
14.63
12.07
19.03
47.3
49.5
70.8
53.1
50.7
28.2
33.4
+0,9 47.6
0,99B
2
12.46
19.63
48.0
27.7
1.008
158
2
Lower
256
9.4M
14.59
12.12
19.02
46.3
49.7
71.5
55.1
52.9
29.0
34.0
48.4
+1.3
1.01B
15B
2
Lower
512
18.9M
14.53
12.03
18.88
45.7
49.8
71.9
56.4
53.6
29.4
33.2
48.6
+1.5
1,098
158
2
Lower
1024
37.7M
14.43
11.98
18.74
46.3
50.0
71.9
55.1
54.3
29.7
33.8
48.7
+1.6
1.07B
/
15B
2
Lower
2048
75.5M
14.79
12.26
19.23
46.1
48.7
71.4
55.4
51.3
28.2
34.0
47.9
+0.8
0.97B
12.26
11.94
43.3
42.2
44.7
29.2
43.3
0.48B
X 150
16.61
15.66
20.27
22.3
30.0
60.9
50.6
37.0
23.0
28.0
36.0
0.48B
2
Step
11.61
39.8
24.9
0.19B
15B
2
Step
256
11 11.5M
11.61
9.89
13.00
39.6
39,9
66.5
53.9
44.4
25.3
30.6
42.9
+0.2
0.50B
150
2
Step
512
23, 1M
11.61
9.89
13.01
39,5
39.9
66.7
53.4
44.1
25.3
29.8
42.7
+0.0
0,538
15H
2
Step
1024
11.60
9.88
13:00
39.7
39,9
66.7
25.1
42.8
+0.1
0.57B
150
2
Step
2048
92.3M
11.58
9.87
13.01
40.1
39.9
66.8
53.4
44.4
24.9
30.0
42.8
+0.1
Lamaa
0.48B
2
Avg
11.86
10.29
39.4
42.7
0.49B
15B
2
Avg
256
11.5M
11.86
10.28
13.41
38.5
39.4
66.2
52.5
42.8
25.9
30.8
42.3
+0.1
0,50B
15B
2
512
23, IM
11,86
10.28
13.41
38.1
39.3
66.3
52.2
42.6
25.6
30.8
42,1
00.1
0,538
158
2
Avg
1024
46 IM
11.86
10.28
13.42
38.4
39,2
65.7
52.7
42.5
25.5
31.0
42.1
-0.1
0.57B
/
15B
2
Avg
2048
92.3M
11886
10.28
13.42
38.5
39,5
65.9
52.7
42.4
25.7
31.0
42.2
00.0
0.48B
15B
2
14.67
12.67
16.68
31,9
32,3
39.1
22.1
27.8
38.3
0.49B
15B
2
Lower
256
11.5M
14.67
12.67
16.69
31.9
32,4
62.7
51.5
38.9
22.5
27.8
38.2
0..1
0,508
158
2
512
1M
14.67
12.67
16.69
32,3
51.7
38.9
22.2
18.2
0.53B
15B
2
Lower
1024
46. IM
14.67
12.67
16.68
31.6
32.3
63.0
51.9
38.9
22.1
28.0
38.3
+0.0
0.578
/
15B
2
Lower
2048
92,
14.67
12.67
16.67
34,1
32,5
62.8
52.4
38.5
23.0
27.6
38.7
+0.4
34

---

Published as a conference paper at ICLR 2025
L EXPANDED RESULTS OF EXTENDED UPTRAINING AND DISTILLATION
Ablation study on individual techniques To further enhance performance through uptraining, we increased the number of uptraining tokens and employed knowledge distillation loss (Hinton et al., 2015; Kim & Rush, 2016). Specifically, expanded the token number from 15 billion to 60 billion. Furthermore, we designated the teacher model as the full-size model for cach architecture, uptrained 15 billion tokens from the SlimPajama dataset. Given the huge number of uptraining tokens, we adopted an online approach to extract logits from the teacher model. Four loss functions were utilized: forward KL (FKL; Kim & Rush (2016)), reverse KL (RKL; Gu et al. (2024)), Jensen-Shannon divergence (JSD; Agarwal et al. (2024)), and total variation distance (TVD; Wen al. (2023)). Table 11 summarizes the controlled experimental results for each method. We observed performance improvement of 1.7% attributed to the extended uptraining and up to 1.7% from the KD loss. We finally selected forward KL as the loss function due its simplicity and superior performance. These significant gains suggest that combining both techniques could yield even greater gains.
Table 11: Evaluation results of ablation studies related to longer uptraining and knowledge distillation loss. Performance improvements, represented by Delta, were measured for each technique. For the knowledge distillation loss function, we experimented with four options: FKL, RKL, JSD, and TVD. Forward KL was selected the final configuration due its simplicity and superior performance.
Uptrain
Looping
LoRA
Perplexity
Few-shot Accuracy
N-emb
PT
Newt
KD
Func
Black
Init
Rank
Init
SSimP
RedP
PG19
I.D
HS
PQ
WG
ARC-e
ARC-c
OB
Avg
A
0,991
15B
x
Step
12.85
10.29
16.21
53.0
57.3
73.2
56.2
56.1
29.2
36.6
51.7
0,998
60B
Step
12.00
9:70
14.84
52.5
59.9
74.7
58.5
58.0
30.3
40.2
53.4
+1.7
0.99B
/ I5B
X
Step
12.85
10:29
16.21
53.0
57.3
73.2
56.2
56.1
29.2
36.6
51.7
FKL
Step
74.8
+1.7
0.99B
/ 15B
RKL
Step
12.56
10.09
15.80
55.6
S8.3
74,3
58.6
58.3
30.4
37.4
53.3
+1.6
0.998
JSD
Step
12.60
10.06
15.77
56.1
73.4
58.4
29. 8
37.2
52.9
+1.2
0.99B
/ 15B
TVD
Step
12:47
9.92
15.52
55.1
58.5
74.0
58.2
58.9
29.5
36.8
53.0
+1.3
1.30B
15B
Avg
256
SVD
12.10
9.71
14.89
58.2
60.7
73.7
59.0
57.6
32.1
38.0
54.2
1,308
/ 15B
FKL
Avg.
256
SVD
11.90
9.52
14.63
59.9
62.0
74,1
60.0
58.6
33.2
38.0
55.1
+0.9
Avg
74.5
58.1
1.30B
/
15B
JSD
Avg
256
SVD
12.09
9.65
14.81
58.1
61.1
73.1
60.8
59.0
33.2
38.6
54.8
+0.6
TVD
Avg
59.3
73.9
38.2
Overall performance after longer training with distillation loss Figure 17 and Table 12 summarize the performance gains achieved by incorporating advanced training techniques: extended uptraining and knowledge distillation loss. We consistently observed substantial improvements in few-shot performance across all architectures and with varying numbers of looping blocks. We anticipate that further performance enhancements can be achieved by utilizing a superior teacher model and increasing the uptraining cost.
35

---

Published as a conference paper at ICLR 2025
Nus 15B
Nok 60B KD Loss
Few-shot Accuracy
+11.5%
+3.8%
44.1%
*
50
+1.6%
43
++20.0% 10000%
Reduced-size
39
Recursive
Full-size
39
Ful-size
40
Recursive
Reducad-slze
Reducad-size Recursive (B2) Relaxed (B2)
Relaxed
Recuratve (B2)
Relaxed
(B3)
Relaxed
(B2)
35
35
0.5B
1.0B
1.5B
2.0B
0.5B
1.0B
0.5B
0.75B
Model Sizes
Model Sizes
Model Sizes
(a) Gemma
(b) TinyLlama
(c) Pythia
Figure 17: Few-shot performance three models with extended uptraining and knowledge distillation. Optimal configurations are used for each model size. The full-size model is the pretrained model itself for Tinyllama, but for other models, is further uptrained on 60 billion tokens. Reduced-size models non-recursive and pretrained from scratch their corresponding sizes. Dotted lines represent the Pareto frontier, showing the optimal trade-offs between model size and performance for each setting.
36

---

Published as a conference paper at ICLR 2025
Table 12: Evaluation results of our Recursive Transformers with 60 billion token uptraining and knowledge distillation loss. We utilized the forward KL loss as the knowledge distillation loss function. Full-size model baselines for Gemma and Pythia are the pretrained models further uptrained on 60 billion tokens, accounting for distribution shifts between Slimapajama and their pretraining datasets. Delta (A) represents the accuracy differences between the longer uptrained models with KD) and their 15 billion uptrained counterparts. We omit the Delta values for relaxed recursive Gemma models with three blocks as they lack 15 billion uptrained counterparts. Results with extended uptraining and knowledge distillation are highlighted.
Uptrain
Looping
LoRA
Perplexity
J
Accurucy 1
Models
N-emb
PT
Nus
KD
Block
Init
Rank
Init
SlimP
RedP
PG19
LD
HS
PQ
WG
ARC-e
ARC-c
OB
Avg
A
1,998
608
8.44
12:21
60.3
67.9
7699
63.5
37.2
39.6
58.6
0.998
60B
/
15.33
13.04
20.37
42.3
43.0
68.8
53.1
49.4
26.3
31.0
44,9
0.99B
15B
22.63
20.03
32,60
28.9
31.6
63.1
52.3
41.2
22.5
27.8
38.2
0.668
16.79
22.85
37.5
38.4
68.7
50.4
24.6
31.6
42.5
0.668
15B
X
24.44
21.69
36.03
27.2
30.6
63.8
50.5
40.6
22.0
27.0
37.4
Step
73.2
56.1
36.6
0.66B
15B
3
Step
14.75
12.10
19.32
45.0
49.9
69.8
55.8
52.7
27.9
33.6
47.8
1,07B
15B
Avg
64
SVD
12.83
10.35
16.02
55.9
56.8
72.5
56.8
55.7
30.6
36.2
52.1
1.15B
Avg
56.1
72.3
37.2
1.308
1SB
Avg
256
SVD
12.10
9.71
14.89
58.2
60.7
73.7
59.0
57.6
32.
38.0
54.2
Gemma
1,608
15B
512
SVD
11.83
9:46
14.57
59.3
62.8
74.0
61.6
60.1
32,9
37.6
55,5
0,998
608
Step
11.44
9.14
13.98
56.5
62.1
75.2
59.4
59.8
32.5
38.6
54.9
43.2
L,07B
60B
2
Avg
64
SVD
11.36
9.14
13.82
58.9
62.8
75.1
61.5
61.2
33.7
37.6
55.8
+3.7
1,15B
60B
2
128
SVD
11.25
9.04
13.64
58.7
63.6
76.5
61.2
62.6
34.6
39.0
56.6
+4.1
1,308
608
2
Avg
256
SVD
11.05
8.88
13.35
60.6
64.7
75.3
62.5
61.6
35 3
38.8
57.0
42.8
1660B
60B
Avg
512
SVD
10.81
8.63
12.94
61.4
65.8
76.3
63.5
65.1
37.1
39.4
58.4
+2.9
Step
9.90
58.
73.1
60.2
37.2
53.3
0.748
608
Avg
64
SVD
12.13
9.80
14.95
55.5
58.
75.5
60.1
58.0
31
36.8
53.3
0.82B
60B
Ave
128
SVD
11.83
9.53
14.51
56.7
60.2
74.2
59.8
59.1
33.0
35.1
54,1
Avg
SVD
9.17
59.3
40.2
56.1
1,278
60H
Avg
512
SVD
11.01
8.80
13:25
61.5
64.9
76.2
62.0
64.3
35.6
39.2
57.7
0,978
12.26
9.37
11.94
43.3
42.2
66.8
53.4
44.7
23.2
29.2
43.3
0.48B
60B
11.93
10.86
13:93
33.3
37.3
66.8
50.1
41.7
23.9
30.2
40.5
0,488
158
16.61
15.66
20.27
22.3
30.0
60.9
50.6
57.0
23.0
28.0
36.0
0,488
Step
11.61
9.89
13:00
39.8
66.5
$2.9
24.9
30.6
42.7
0.538
15B
2
Avg
64
SVD
11.22
9.66
12.51
41.8
41.6
67.0
$3.3
43.9
24.7
31.2
43.4
0.58B
15B
2
128
SVD
10.99
9.45
12.21
43.2
42.1
68.3
53.2
44.8
25.9
30.4
44.0
Tinyl.lama
0.688
2
Avg
SVD
10.21
9.18
43.2
68.1
53.5
30.4
44.22
0.868
15B
X
2
Avg
512
SVD
10.46
8.92
11.50
46.0
44.1
68.2
$3.0
45.8
25.1
31.2
44.8
0,488
Step
44.2
44.3
0.53B
60B
2
Avg
64
SVD
10.14
8.77
11.19
44.3
44.9
69.5
52.5
46.5
26.
31.6
45.1
+1.6
0.58B
60B
2
Avg
128
SVD
10.07
8.68
11.07
45.9
45.1
69.4
50.5
46.8
25.4
31.6
45.0
+1.0
2
Avg
9.96
8.56
31.6
0.868
60B
Avg
512
SVD
9.85
8.44
10.76
47.4
46.3
69.7
52.8
47.5
26.3
31.4
45.9
+1.1
0.81B
60B
12.83
9.76
13.57
53.0
50.2
71.1
54.8
$1.9
27.7
31.6
48.6
0.408
608
18.27
14.39
21.93
32,1
35.0
66.1
49.6
42.9
24.2
27.0
39.5
0.40B
15B
X
25.69
20.00
32.08
24.3
30.0
61.9
50.7
38.3
22.3
26.0
36.2
0.40B
Step
50.8
25.7
30.0
43.5
0.448
158
Avg
64
SVD
16.03
12.19
17.59
45.8
40.9
67.3
50.0
45.8
25.5
51.8
43.9
0.48B
15B
128
SVD
15.67
11.93
17.10
46.9
41.9
67.4
51.2
45.4
24.5
31.2
44.1
Pythia
2
Avg
SVD
48.5
25.5
32.0
0.708
158
X
2
Avg
512
SVD
14,70
11.07
15:71
50.2
44.7
68.2
51.6
47.6
25.4
31:2
45.6
0,400
Step
47.8
43.
09.3
30.4
45.2
0.44B
60B
2
Avg
64
SVD
14.24
10.89
15.52
50.0
41.5
68.9
54.1
48.0
26.5
31.2
46.2
+2.3
0.48B
60B
2
Avg
128
SVD
14,10
10.79
15.27
50.1
45.5
69.0
52.6
48.3
25.8
32.0
46.2
+2.1
0.558
Avg
256
SVD)
10.61
14.91
50.5
45.6
68.7
51.2
25.7
32.8
46.1
0.708
608
2
Avg
512
SVD
13.59
10.38
14.43
52,0
47.0
69.6
53.4
48.9
26.9
31.2
47.0
+1.4
37

---

Published as a conference paper at ICLR 2025
M EXRLY-EXIT TRAINING
Ablation study on early-exit training strategy To enable early-exiting capabilities, all models require additional training to align intermediate representations with classifier heads. In this study, we conduct ablation studies various strategies, demonstrating Recursive Transformers can be trans- formed into carly-exiting models without compromising final loop output's 's performance. Table 13 presents a comprehensive summary of our findings across various categories, including training procedures, loss functions, and early-exit training data. Our key findings are as follows:
Post-training after uptraining is essential for preserving final loop performance. Jointly training intermediate loop output during the uptraining phase, even with an aggressive loss coefficient strategy, significantly degraded the final output performance.
Training solely the early loops with learnable LoRA modules, while freezing other parameters, hindered effective intermediate representation learning. We attempted to fine-tune intermediate outputs by attaching LoRA modules to classifier heads, but this proved ineffective.
The aggressive coefficient strategy successfully maintained final loop output performance while enhancing intermediate layer performance. Moreover, incorporating knowledge distillation from detached final outputs further enhanced intermediate layer performance.
No significant performance differences were observed when using the same uptraining data versus new SlimPajama tokens for post-training.
Finally, we opted to utilize the uptrained model and perform post-training with new tokens sourced from the same SlimPajama dataset. Moreover, we incorporated distillation loss from the final loop output, while using a strategy that aggressively reduces the loss coefficients of intermediate outputs.
Table 13: Ablation studies on early-exit training for recursive Gemma models. We evaluated performance in a static-exiting scenario (Schuster et al., 2022; Bae et al., 2023), where all tokens exit at either 9th or 18th layers. We explored post-training (after uptraining) and co-training (during uptraining) approaches. We experimented with freezing uptrained weights and adding LoRA with the rank of 128 to the classifier head, and we used weighted CE and aggressive CE loss functions. Early-exit training utilized 15 billion tokens, either overlapping with uptraining data entirely new. Delta (A) indicates the performance changes of the final layer. We highlight the final configuration: post-training with aggressive CE and KD loss on 15 billion new tokens.
Uptrainn
Looping
Earty-Exir Teain
Perplesity I
Fow-shot Aceuracy
N-omb
N...
Rlock Train Freeze Nu.
Tatn
SlimP
RedP
PG19
PQ
ABC=
ARC-
Avg
1.998
158
10.76
8.47
13,08
63.5
68.5
77.0
63.5
67.6
18.1
42.6
60.1
0.99B
x 150
22.63
20.03
32.60
28.9
31.6
63.1
52.3
41.2
22.5
27.8
38.2
0,998
15B
2
Step
12:85
10.29
16.21
53.0
57.3
73.2
56.2
56.1
29.2
36.6
51.7
0.99B
Step
Weighted
Ovlp
12:97
13:11
10.51
10.59
16.71
54.8
72.7
72.0
53.4
54.1
29.1
30.1
36.0
35.6
19.8
-1.2
0.99B
Step
Age (0.3)
Ovlp
12.00
13.03
10.21
11.02
17.35
47.5
53.0
71.2
73.7
56.9
50.2
28.2
29.9
34.8
1:5
Post-
Apo (0.1)
Ovlp
12:37
14.55
9,94
73.9
46.1
26.5
30.6
32.0
0.998
15R
Post-
Agg (0.05)
Ovlp
15.70
12:33
12.93
9.90
15.31
52.8
59.2
73.6
57.5
57.7
26.9
30.5
37.2
52.6
+0,9
0,998
< 15B
2
Step Posi
15B
Agg (0.01)
Ovlp
22.76
12.28
9.80
15.23
52.9
59.5
61.5
73.3
56.5
40..5
57.2
26.5
30.1
37.2
52.4
+0.7
0.99B
< 15D
2
Step
Post-
15B
Weighted
Ovlp
1304
13:04
10.57
10.54
16:66
16.56
47.7
48.3
54.9
55.1
72.1
73.2
55.6
55.9
54.5
54.3
29.1
28.4
35.4
49.9
11.3
0.998
<
Step
158
Age (0.1)
Oxlp
12.40
14.11
11.47
9.97
18.32
46.3
52.1
73.7
71.6
55.3
49.2
28.5
32.6
48.0
0.99B
/ 1:8
2
Step Post-
15B
Standard
Ovip
43.74
12.85
41.63
10.29
56.78
16.21
53.0
5.3
57.3
57.9
61.4
73.2
52.6
$6.2
56.1
35.3
29.2
24.0
29.0
36.6
51.7
35.0
+0.0
0.99B
< 15B
2
Step Post-
/
15B
Standard
Ovlp
12.85
10.29
16.21
53.0
57.3
73.2
56.2
56.1
29.2
36.66
51.7
+0.0
43.09
39.97
55.37
5.6
57.7
62.5
52.7
34.5
24.7
29.3
35.3
0.998
1.5
Step
Agu (0.1)
Oxlp
13:24
13.59
10.67
10.89
71.2
72.2
54.4
$4.7
27.5
28.9
35.00
37.4
0.99R
<
15R
2.
Step
Poot-
x
15B
Agg (0.1)
x
New
14.49
12:44
11.86
9.92
15.31
18.89
52.3
43.9
59.0
51.3
71.0
73.8
57.6
54.9
48.1
55.5
27.5
30.4
37.2
31.4
52.3
46.9
+0.6
38

---

Published as a conference paper at ICLR 2025
Early-exit training results final models We applied the aggressive coefficient strategy with distillation loss to the models uptrained on 60 billion tokens. Tables 14 and 15 summarize the performance of intermediate loops and the final loop actoss three models. For fair comparison, the full-size models (Gemma and Pythia) were also uptrained with 60 billion tokens and then post-trained with 15 billion tokens. As the optimal strategy derived from the non-relaxed models was directly applied to the relaxed models, further exploration of optimal strategies specifically for relaxed models left for future work.
Table 14: Evaluation results of Gemma models after early-exit training. For relaxed models, we also experimented with increasing the coefficient to 0.3 because of the lower performance of the intermediate layer. The relaxed model with three blocks shows more significant performance drop because KD loss could not be utilized due to out-of-memory issues. Delta (A) represent the accuracy changes of the original last layer after early-exit post-training. These changes should be compared in reference to the performance drops observed in 75B and 60B uptraining for the full-size model.
Uptrain
Looping
LoRA
Early-Esit Train
Perplesity
Few-shot Accuracy
N-emb
PT
Neca KD
Block
Init
Rank
Init
Naw
CE
KD
|SlumP
RedP
PG19
L.
HS
PQ
WG
ARC-e
ARC-c
OB
Ang
1.998
/
CUB
10:58
8.44
12.71
60.3
67.9
76.9
63.5
64.9
37.2
39.6
58.6
1,98
75B
11.03
8.88
13.33
57.0
65.9
76.2
63.9
63.0
35.9
38.8
57.3
-1.3
0.99R
Step
11.44
13.08
56.5
62.1
75.2
59.4
59.8
32.5
54.9
1,078
(0B
Avg
64
SVD
11.36
9.14
13.82
58.9
62.8
75.1
61.5
61.2
33.7
37.6
55.8
L.158
60B
Avg
128
SVD
11.25
9.04
13.64
58.7
63.6
76.5
61.2
62.6
34.6
39.0
56.6
1,30B
60B
Arg
256 SVD
11.05
8.88
13.35
60.5
64.7
75.3
62.5
61.6
35.3
38.8
57.0
L600
Aug
SVD
10.81
12.94
61.4
65.8
76.3
63.5
0.99R
60B
2
Step
158
Agg (0.1)
11.71
13.68
11.39
9.56
14.16
17.60
45.0
51.0
54,1
61.7
75.1
71.9
$8.9
58.5
49.8
6
31.9
28.8
37.6
33.4
51.0
48.8
-0.9
1.078
/
60B
2
Avg.
64
svn
158
Agg (0.1)
11.79
19.45
16.46
9.70
26.10
14.52
53.7
30.7
60.8
37.9
73.6
66.5
55.3
61.1
58.7
42.2
25.3
52.9
57.2
27.6
54.0
40.8
-1.8
L.15B
/
60B
2.
Aog
128
SVD
15B
(0.1)
v
11.66
19.65
16.77
9.59
26.44
14.32
29.7
53.3
62.1
37.7
74.9
66.8
60.0
52.6
41.4
59.9
25.3
53.4
28.0
38.8
54.6
40.2
-20
L,30B
60B
2.
Aog
256
SVD
15D
(0.1)
11.47
19.67
16.82
9.39
26.40
14.03
29.7
51.9
63.0
38.3
74,5
66.4
61.7
53.1
43.8
60.5
24.7
33.1
27.6
58.4
400.5
55.2
-1.5
1.60R
60B
2.
Aug.
512
SVD
ISR
(0.1)
11.20
19.29
16.47
9.14
25.73
13.58
57.2
32.0
39.6
61.1
67.6
75.2
53.3
61.7
43.2
62.1
34.6
25.8
30.2
58.2
41.7
56.2
-2.2
1.078
Arg
64
SVD
158
Agg (0.3)
16.09
12.11
13.54
21.19
14.97
35.4
59.8
42.8
74.4
69.8
59.4
52.8
57.6
45.8
31.1
25.8
31.0
43.3
1.158
60B
Avg
128 SVD
15B
Agu (0.3)
16.28
11.96
13.77
21.45
14.16
52.3
35,2
60.5
42.1
74.2
69.8
$3.5
58.9
46.5
33.0
25.8
31.2
43.4
L,30B
60B
Avg 256 SVD
15B
Agu (0.3)
11.73
16.41
13,89
21.68
14.43
54.3
35.6
61.4
42.3
75.0
69.0
60.7
$2.7
46.8
58.8
26.4
33.1
20.8
38.6
43.2
54.6
L6608
GOB
2
Arg
512 SVD
15B
Agg (0.3)
16.24
11:47
13.72
9.36
13.93
21.42
56.2
37.8
62.7
43.6
75.4
69.0
60.9
54.4
60.4
45.5
26.4
34.0
31.2
37.0
55.2
44.0
-3.2
0.66B
60B
Step
12.27
9.90
15:24
55.6
58.1
73.1
60.2
58.5
30.2
37.2
53.3
0,74B
Avg
12.13
58.3
0.878
60B
Avg
128
SVD
11.83
9.53
14.51
56.7
60.2
74.2
59.8
59. 1
33.0
54,1
0.978
60B
Avg.
256
SVD
11.49
9.17
13,87
59.3
63.6
74.7
61.2
61.6 66
32.9
40.2
56.1
1.278
60B
Avg
512
SVD
11.01
8.80
13:25
61.5
61.9
76.2
62.0
61.3
35.6
39.2
57.7
12.75
10.48
16.01
50.2
$7.0
72.7
58.6
56.7
30.0
38.2
-1.4
0.668
Step
15B Agg (0.1)
13.81
11.47
17.50
53.0
72.4
55.6
51.6
27.2
49.0
16.72
14,23
22.97
37.7
44.2
69.8
$3.6
44.2
24.6
30.2
43.5
12.64
10:43
15.81
51,4
56.3
72.2
$7.9
56.7
30.4
35.0
51.4
1.9
0.74B
60B
Arg
64 SVD
15B Age (0.1)
19.90
16.88
26.26
30.4
39.3
66.3
54.1
41.2
24.8
29.2
40.8
26.31
22.19
36.10
20.9
31.2
62.6
50.8
37.2
22.0
23.0
36.1
12.37
10:21
15.38
52.0
58.0
72.0
56.5
55.4
30.0
35.2
51.7
-2.4
0.828
/ 60B
Avg
128 SVD
15B
Agg (0.1)
x
20:07
17.09
26.47
30.9
40.5
66.3
55.4
40.8
24.4
29.6
41.1
26.15
22.46
35.98
31,2
36.4
22.9
26.2
36.1
11.92
9.78
14.71
54.8
60.6
74.6
60.1
60.1
31.8
36.6
54,1
-2.0
0.971
60B
Avg
256 SVD
15B
Agg (0.1)
19.29
16.49
25.51
35.2
42.5
55.6
41.5
25.6
29.4
25.12
21.53
34.53
23.1
32.0
63.2
49.7
23.0
36.1
14.00
62.1
34.9
38.5
1.27B
Arg
SVD
Age (0.1)
15.79
24.34
36.7
44.9
67.2
55.3
43.8
21.19
20.70
33.20
24.4
32.4
63.9
50.8
37.9
21.9
27.4
37.0
13.07
10.84
16.49
47.7
54.4
71.7
56.1
55.9
20.4
35.2
50.1
-3.2
0.74B
60B
Avg
SVD
15B Agg
16.68
14.08
21.86
35.4
42.4
68.2
53.8
44.6
26.3
29.4
21.43
18.26
29.12
24.4
34.1
64.3
50.5
40.7
22.3
12.71
10:54
15.92
50.4
55.9
73.1
57.5
56.5
30.1
34.8
51.2
-2.9
0.828
Arg
SVD
15B Age (0.3)
16.90
14.32
22.18
37.6
43.5
67.6
45.0
25.3
21.23
18.13
28.58
25.3
34.0
61.6
51.7
40.7
23.0
26.1
38.0
12.26
10:15
15.23
53.5
58.5
73.5
58.8
30.6
0.978
60B
Avg
256 SVD
158
Agg (0.3)
x
16.56
14.09
21.08
42.6
45.1
68.2
57.7
45.7
25.9
28.8
44.8
20.78
17.72
28.29
27.9
34,3
66.3
52.2
39.7
23.6
26.8
38.7
11.80
9.68
14.45
54.1
61.2
74.0
59.0
59.9
32.9
38.0
54.1
-3.6
1,278
/ 60B
Avg.
512 SVD
15B
(03
x
16.02
13.53
20356
43.5
47.5
68.3
56.2
47.1
27.0
30.4
45.7
20.20
17.21
27.50
26.9
35.2
65.6
52.9
41.9
23.2
26.6
39.2
39

---

Published as a conference paper at ICLR 2025
Table 15: Evaluation results of TinyLlama and Pythia models after early-exit training. Delta (A) represents the accuracy change in the original last layer following early-exit post-training. In case of Pythia, these changes should be compared in reference to the performance drops observed in 75B and 60B uptraining for the full-size model.
Uptrain
Looping
LaRA
Eart-Eait Train
Perplexity
J
Few-shot Accurncy
Models
N-emb|
PT
Ausi
KD
Block
Inii
Rank
Init
CE
KD
SlimP
RedP
PG11
LD
HS
PQ
WG
ARC-
ARC-c
OB
Avg
0.978
12.26
9.37
11.96
43.3
42.2
66.8
53.4
44.7
23.3
29.2
43.3
0.480
60B
Step
10.51
11:60
44.2
43.1
68.
52.4
44.7
25.3
32.2
36.3
0.538
6BB
SVD
10.14
11:19
52.5
0.58B
6UB
SVD
10.07
11:07
50.5
45.0
0.68R
60B
53.2
0,560
10.76
47.4
52.8
26.3
0.488
GUB
Step
Ags (0.3)
10:55
12.28
11.68
13.85
38.2
29.4
52.3
41.8
0.538
60B
64 SVD
15B
Agp
21.25
10.34
9.08
11.50
41.4
44.5
69.5
48.9
53.4
46.9
15.66
32.0
33.7
45.1
+0.T
0.588
608
Aeg
128 SVD
158
(0.)
21.30
10.25
18.56
8.97
24.75
11.36
45.2
18.5
78.9
45.5
68.8
58.4
45.0
$4.0
46.5
34.1
21.88
25.0
27.4
31.6
45.2
$5.0
+0.2
TinyLaana
0.68R
60B
Avg
SVD
(0.)0.1)
20.95
10:13
18.16
24.22
20.1
28.8
48.9
33.8
22.5
25.0
33.9
0,568
60B
SVD
Agg
20.38
10.02
8.74
11.04
46.6
46.5
54.5
26.3
32.2
46.1
(0.(0.3)
10.61
54.1
26.0
0.58R
60B
178 SVD
15B
(g.(0.3)
17.10
10.50
15.03
9.22
19.99
11.71
23.5
44.2
30.1
44.2
69.2
60.8
51.3
53.0
46.0
36.5
23.8
25.5
26.4
31.2
36.0
44.8
-0.2
0,688
6UB
Aug
256 SVD
15B Ags (0.3)
10.34
17.06
14,02
9.07
11.51
1982
44.0
24.2
45,0
30.4
68.4
59.9
53.0
51.7
43.8
36.2
26.0
23.9
31.2
27.2
44.8
36.2
0.88
Avg
(0.3
1021
16:76
19:43
51.9
36.5
0.IR
60R
12:83
13.57
54.8
48/6
0.81R
75B
12:86
9.86
13.74
50.3
55.3
78.8
49.3
0,40B
CUB
Step
14.59
11.13
15.79
47.8
43.8
69.3
52.0
48.1
25.4
30.4
45.2
0.44B
60B
Avg
64
SVD
14.24
10.39
15.52
50.0
44.5
68.9
54.1
48.0
26.5
31.2
46-2
0.48
60B
Aug
128
SVD
14.10
10.70
15.27
50.1
45.5
69.0
$2.6
48.3
25.8
32.0
46.2
0.558
GUB
A4g
256
SVD
1391
10.61
14.91
50.5
45.6
68.7
51.2
48.4
25.7
32.8
46.1
0.70R
60B
Ang
512
SVD
13.59
10:38
14.45
53.0
47.0
69.6
$3.4
45.9
76.9
31.2
47.0
0,400
/
60B
2
Step
15B Ags (0.1)
15.61
14.72
11:38
14.31
20.96
16.31
47.0
38.4
28.1
44.2
69.3
67.0
53.7
53.4
48.6
43.5
24.7
24.4
29.0
30.4
45.
42.0
+0.2
0.44R
60B
Avg
Ago (0.1) 1)
24:43
14:59
11:22
27.89
16.12
49.1
26.7
43.9
31.5
53.8
90.8
26.
22.9
31.2
27.6
46.1
37.1
-0.1
(g.(0.1)
14.35
27.96
15.05
Pythis
0.558
6UB
Avg
256 SVD
15B
Agg (0.3)
24.18
14.14
17.57
10.06
27.48
15.54
28.1
50.8
45.5
32.3
61.9
68.2
$4.1
53.9
48.8
38.1
22.9
25.3
28.6
32.8
46.5
38.0
+004
0,708
60B
2
Aug
512 SVD
158 Ags (0.3)
23.50
13.81
17.49
10.72
26.72
15.11
52.4
29.5
47.0
32.8
69.3
63.2
52.7
52.5
50.1
38.8
26.9
22.5
27.8
32.0
36.2
47.2
+0.
0.44B
608
2
Avg
64
SVD
15B
(0.(0.3)
2062
14:87
11.59
15:60
23.57
16.61
47.0
32.6
49.1
33.6
68.3
63.4
51.2
53.0
47.4
40.7
23.3
35.3
28.0
$1.0
45.1
39.0
-0.9
0.488
60B
Aeg
Ags (0.3)
2060
23.65
16.36
53.2
39.5
$1.1
41.5
13.6
45.9
39.0
0.550
60B
256
SVD
(0.3)
2061
14.44
11.20
15.4
52.3
25.4
46.0
0,709
60B
2
Avg
512
SVD
15B Agg (0.3)
20.20
14.08
15.25
10.54
22.98
15:46
5t.1
34.6
34,1
46.4
68.7
63.5
53.0
52.2
41.5
50.0
26.9
23.6
27.6
31.6
39.7
46.7
-0,5
40

---

Published as a conference paper at ICLR 2025
N HYPOTHETICAL GENERATION SPEEDUP
Measuring the average generation time per token First, we measured the generation time with various model configurations using dummy weights and inputs. We measured the elapsed time for each components, such embedding matrices, Transformer blocks, and the classifier head. We measured decoding speed using FlashDecoding (Dao et al., 2022), a technique that has recently calculated the time per token by dividing the total
become standard in serving LLMs. Especially, time by the decoding length. Default prefix and decoding lengths are set to 512 and 2048, but we also used shorter context lengths, like 64 and 256 to simulate scenarios where parameter memory sizes become limiting. Using& single A A100 40GiB GPU, we measured generation times by increasing batch sizes until out-of-memory error occurred or memory usuage reached the predefined limit.
In Table 16, generation time was measured up to the maximum hatch size that single A 100 GPU could accommodate before encountering out-of-memont errors, with prefix and decoding lengths set to 512 and 2048, respectively. Mcanwhile, Table 18 presents generation times measured in a more deployment scenario, where the prefix and decoding lengths were reduced to 64 and 256, and the memory limit was set to 16GB. As anticipated, under severe memory constraints, the reduced parameter memory footprint of Recursive Transformers enabled substantially larger batch sizes. This observation indicates that Recursive Transformers, even without continuous batching techniques, can achieve higher throughput than vanilla Transformers due to their inherent memory efficiency.
When comparing the speed of the three models, Gemma 2B was the fastest, followed by TinyLlama 1.IB and then Pythia IB. This order is the exact inverse of their non-embedding parameter sizes. This speed difference is attributed to the SrroppedQueey and Multi-Query Query attention mechanisms (Ainslie et al., 2023). The main decoding bottleneck in Transformers is memory access the key-value cache. Hence, Gemma that effectively reduces the key-value cache size through the MQA mechanism, achieves fastest speeds. Despite using GQA, TinyLLlma 1.1B has similar speed to Gemma 2B due to its shallow and deep architecture (22 layers compared to Gemma's 's 18 layers). This deeper architecture likely offsets the speed gains from the attention mechanism.
Comparison of hypothetical generation throughput We conducted early-exiting simulations using language modeling datasets (SlimPajama, RedPajama, and PG19), assuming our models generated the tokens. For each dataset's 's test set, we employed an oracle-exiting algorithm to determine the earliest possible exit point for each token. We used 20K samples obtain their exit trajectories. By combining this trajectory data with previously measured per-token processing time (considering only Transformer block computations), we estimated the hypothetical throughput across various settings and datasets. The results are detailed in Tables 17 and 19.
Our analysis reveals that Recursive Transformers achieve a 2-3x throughput gain over vanilla Transformers. Relaxed models also demonstrate significant speedup despite unoptimized LoRA computations. Currently, we merge multiple LoRAs into a single, larger LoRA to enable parallel computation of samples across different looping iterations. However, this incurs extra overhead due to redundant computations. Therefore, we observed reduced throughput gains in memory-constrained scenarios (shorter context lengths and lower memory limits). This degradation stems from the increased proportion of LoRA computation time relative to overall processing time. Because attention computation has quadratic complexity with respec lengths, it becomes less expensive shorter context lengths, while the complexity of LoRA computation remains constant. This highlights the impact of unoptimized LoRA computations, leading to substantial throughput reduction. However, these findings suggest that relaxed models will yield even greater performance and throughput improvements in scenarios with longer contexts where attention computation dominates.
41

---

Published as a conference paper at ICLR 2025
Approximation errors in our hypothetical throughput Since our throughput estimations based on theoretical estimation, they may introduce certain approximation errors as follows:
As our models are not fine-tuned for any downstream task, we simulated the exit trajectories of language modeling datasets by assuming they were generated by our models. While this approach expected closely approximate actual generation, empirical validation is necessary to confirm its accuracy.
Throughput gains should be measured using realistic (confidence-based) early-exiting algorithms, rather than relying on the oracle-exiting algorithm. While early-exiting algorithms can introduce performance degradation due to inherent errors in confidence estimation, they also incur addi- tional computational costs for estimating prediction confidence, necessitating further efficiency improvements.
Our analysis solely focused on speed improvements within Transformer blocks. However, upon early exiting, the exited tokens require separate processing through the embedding layer or the classifier head for subsequent sequence generation. This necessitates non-exited tokens to wait for others, potentially reducing efficiency as the embedding layer computation may not fully utilize the maximum batch size.
Early-exiting architectures require computing key-value caches in remaining layers for already exited tokens prevent performance degradation (Bae al., 2023). While this adds negligible overhead in memory-bound scenarios, it inevitably increases overhead in compute- boutd scenarios where the maximum batch size is fully utilized. Our throughput estimation, however, excludes the computation time for these key-value caches in later loops (though we did account for their memory size). Incorporating these computations into more realistic analysis of carly-exiting generation is direction for future work.
42

---

Published as a conference paper at ICLR 2025
Table 16: Measurements of generation time across three models using a single A 40GB GPU. We measured time per token for both batch size of and the maximum batch size achievable by each model. The prefix length was set to 512 tokens, and the decoded output length to 2048 tokens. We then averaged the total elapsed time by the output length of 2048. Dummy input and dummy tensors were used for measurement. Both Gemma, employing multi-query attention, and TinyLlama utilizing grouped-quee -ueery attention, demonstrated fast generation speeds and large maximum batch sizes relative to their model sizes. TinyLlama's s deep and narrow architecture allowed for a significantly large maximum batch size, although generation speed was slower due to the increased number of layers.
Architecture
Recursive
Time (ms) per
Models
N2
dmodet
Nusod
Nxv
Voçab
N-emb
Block
Rank
Batch
Total
Emb
Transfomer
Head
18
2048
256K
1,98
22,994
0.087
21.344
0.803
43
0.657
0.002
0.616
0.023
13.918
0.088
11,059
0.827
0.336
0.002
0.265
0.023
18
2048
256K
L,788
64
15.858
0.080
13.096
0.825
41
0.398
0,002
0.323
0.024
2048
256K
1,158
128
15,708
0.080
12,969
0.822
0.398
0.002
0.324
0.024
18
2048
256K
1.308
256
15:456
0.083
12.721
0.818
39
0,450
0.002
0.372
0.025
Gemma 2B
15,489
0.078
12,775
0.817
1.60B
18
2048
256K
0.668
43
10.546
0,263
0.081
0.002
0,182
7,394
0.827
0.023
1
11,871
0.080
8,724
0.827
2048
256K
0,748
64
2048
256K
0.828
128
43
0.294
0.002
0.221
0.023
2048
256K
0.97B
256
1
12.018
0.081
8.848
0.823
0.024
12,087
0.822
2048
256K
1.278
512
39
0.325
0.002
0.237
0.025
22,016
0.082
22
2048
32K
0.97B
329
0.819
0.000
0.815
0.001
22
2048
32
32K
0,48B
1
12.657
0.077
10.370
0.209
15.243
0.211
TinyLaama 1.18
2048
0,538
211
0.454
0.000
0.421
0.002
22
2048
32
32K
0.588
128
1
15.456
0.082
13.118
0.213
209
0.454
0.000
0.002
15,223
0.081
12,908
0.208
0.68B
209
0.457
0.000
0.423
0.002
22
2048
32
32K
0.86B
2
512
15.383
0.080
13.002
0.211
209
0.461
0.000
0.428
0.002
16
2048
50K
0.818
53
1
13,280
1.227
0.002
0.080
12,286
1.206
0.005
0.235
8,423
0.081
6.378
0.262
0.856
0.001
0.606
0.005
Pythia IB
16
2048
50K
0.44B
2
64
63
10.554
0.875
0.082
0.001
0.626
8.519
0.005
0.260
0.48B
128
10.167
0.076
8,196
0.256
0.892
0.001
0.642
0.005
16
2048
S0K
0.558
2
256
59
10.410
0.913
0.001
0.079
0.662
8,402
0.258
0.008
0,708
512
12.609
0.091
10.311
0.267
0.006
43

---

Published as a conference paper at ICLR 2025
Table 17: Hypothetical generation speedup of Recursive Transformers across three models. We utilized the measurements of tokens per second calculated in Table 16. We only considered the time spent within Transformer blocks, simulating generation on the SlimPajama, RedPajama, and PG19 test sets. We used a vanilla transformer model, both with and without continuous sequence-wise batching, as our baselines. Our Recursive models further enhance throughput by applying continuous depth-wise batching, leveraging looping and early-exiting techniques. The throughput improvements over the vanilla Transformer and sequence-wil wise batching are denoted as Ay and Aseg. respectively. To aid in understanding the speedup, we also provide the performance of intermediate layers and the maximum batch size.
Uptrain
Looping
LoRA
Farty-Rait Train
Batching
Feu-shot Accuracy
Thronggoot 1
Models
Noemb]
PT
Nu.s
Black
Rank
Init
No.
CF
KD
Type
Exil
Last
Mil
Mid2
Butch
SlimP
RedP
PG19
Av
Auvs
1,998
73B
57.3
43
655
1228
1357
x1100
x0.71
1.99
73B
CSB
57.3
43
1622
1604
1357
x1.41
x1.00
0.908
V
60B
Step
15B
Agg I)
CDB
54.0
48.8
43
31.99
3050
2421
X2.06
x1.88
1,07
60B
Aug
64
SVD
15B
(02(0.3)
CDB
54.0
40.8
41
2357
2255
1858
x2.00
x1.41
1.150
600
Avg
128
SVD
150
(g2(0..)
CDB
54.6
40.2
41
2355
2250
1846
x1.99
x1.41
1,508
60B
256
SVD
158
(gx(0..)
CDB
55.2
40.5
39)
2047
1976
1740
x1.78
x1.26
1,608
64B
512
SVD
15B
(0.3)
CDB
56.2
41.7
39
1806
1754
1598
x1.59
x1.13
1,070
V
600
64
SVD
150
(0.3
CDB
S3.1
43.3
2454
2357
1929
x2.08
x1.47
1,158
60B
128
SVD
(0.3)
CDB
83.6
43.4
2415
2546
1926
x2.07
x1.47
1,308
6BB
Agz
54.6
2156
x1.31
60R
(05
0.64R
Aps
43.5
3130
1041
Agg
x2.06
0,620
Ags
x2.02
0.788
64B
Ags (0-1)
54.1
42.2
2261
2219
X2.00
1,27R
<IR
Avg
SVD
Agp (0 1)
CDB
7
43.5
37.0
2181
2122
x1.91
x1.35
8,748
6BB
Avy
SVD
Ags 31
CDB
42.9
37.7
2437
2372
X2.14
x1.51
0,829
60B
AVE
128
SVD
15B
Aex
CDB
51.2
43.2
38.0
43
2576
2521
2084
x2.00
x1.48
0,978
60
ANE
256
SVD
15B
(02
CDB
53.0
44.8
38.7
2339
2500
2009
x/2.07
x1.46
1,278
608
Avg
512
SVD
158
Agg (0.3)
CDB
54.1
45.7
39.2
39
2251
2191
1975
x1.98
x1.40
0,970
/
43.3
329
1208
1220
1194
X:1.00
x0.99
0.970
CSD
43.3
329
1227
1225
1194
x1.01
x1.00
0.458
V
6BB
Step
158
Age (1.3)
CDB
44.8
41.8
233
2038
2023
1933
X1.66
x1.64
0,538
<4B
64
SVD
158
Agx (1.3)
CDB
45.1 1
33.7
211
1733
1719
1617
x1.40
x1.39
0,53R
60R
128
SVD
15R
(0 11
CDB
45.2
33.9
700
1739
17177
1609
x1.40
x1.39
Liuma
0,640
60B
Agz
CDB
33.9
1728
1714
x1.08
0,860
600
Agx (0.1)
CDB
34.2
1716
1702
x1.38
x1.37
0.53R
Agg
0.54R
Agg
CDB
0,680
AYE
0.868
606
0818
/
738
822
X1.00
0.8IR
/ 73R
CSR
879
822
x1.07
x1.00
0,400
V
60B
Step
Ago ()
CDB
42.0
1330
1333
1281
x1.71
x1.80
0,419
60B
Avg
64
SVD
15B
Agx
CDB
46.1
37.1
63
1206
1203
1140
x1154
x1.43
Pylii
0.558
0.488
6088
60R
Ang
Ang
256
128
SVD
SVD
158
158
Ags (0:1)
Ags (0.3)
CDS
CDB
46.5
46.2
37.8
38.0
59
59
1156
1138
1139
1180
1071
1108
x1.45
X1.49
x1.85
x1.89
0.700
60B
Avg
512
SVD
150
Agg (0.1)
CDB
47.2
38.2
53
1051
1077
1021
x1.36
x1.27
0.640
/
600
Avg
64
SVD
150
(03
/
CDB
65.1
39.0
63
1254
1252
1190
x1.60
x1.49
0.455
128
SVD
158
Aex (0.3)
CDB
45.9
39.0
59
1200
1226
1155
x1.65
x1.45
0.539
256
SVD
158
(0.3)
CDB
46.0
39.4
59
1150
1180
1112
x1.50
x1.40
0.70R
40R
512
SVD
158
Agz (0-3)
CDB
46.7
39.7
53
1058
1114
1058
x1.41
x1.32
44

---

Published as a conference paper at ICLR 2025
Table 18: Generation time measurements of Gemma models single A100 GPU with 16GB memory constraint. We measured time per token for both a batch size of and the maximum batch size achievable by each model. The prefix length set 64 tokens, and the decoded output length to 256 tokens. We then averaged the total elapsed time by the output length of 256. Dummy input and dummy tensors were used for measurement.
Model Architecture
Recursive
Time per
Models
N1
desedal
Nhcad
Ngv
Vocab
N-emb
Block
Rank
Batch
Total
Emib
Transformer
Head
22.577
2048
256K
1.98B
111
0.207
0.001
0.188
0.010
2048
256K
0.99B
123
1
13.576
0.001
0.079
10.819
0.009
0.815
20488
256K
1.07B
64
117
0.140
0.001
0.112
0.009
20488
256K
1.158
128
115
1
15.631
0.141
0.001
0.082
12,899
0.113
0.010
0.816
15.317
12.639
0.811
2048
1.308
111
0.143
0.001
0.115
0.010
Gemma 28
18
2048
256K
1 1.608
512
103
1
15.379
0.158
0.080
0.001
12,692
0.127
0.011
0.807
0.66B
10.528
0.817
131
0.087
0.001
0.058
0.010
18
2048
256K
0.74B
64
123
11,957
0.105
0.081
0.001
0,075
8,855
0.009
0.815
0.82B
11,898
0.080
8,787
0.816
0.103
0.001
0.009
18
2048
256K
0.97B
256
117
11.734
0.106
0.079
0.001
0.076
8.654
0.009
0.813
2048
1.278
512
1
11.986
0.080
8.856
0.809
0.010
0,97B
23.898
0.080
22.909
0.189
1049
0.129
0.001
22
2048
32K
0.48B
1121
14.129
0.070
0.080
0.000
11,846
0.064
0.001
0.202
TinyLLaam I.IB
2048
32K
0.53B
64
14,897
0.080
12.627
0.202
2048
32K
0.58B
128
1089
0.074
0.000
0.069
0.001
2048
32K
0.68B
256
1
14.962
0.081
12.659
0.201
1065
2048
0.86B
512
1017
0.080
0.000
0.075
0.001
13.341
0.239
2048
50K
0.81B
229
0.176
0.000
0.171
0.002
2048
50K
0.40B
241
1
0.121
8.356
0.000
0.079
0.086
6.303
0.261
0.002
0.262
Pythia IB
20488
0.44B
64
233
0.133
0.000
0.097
0.002
16
2048
50K
0.488
128
221
1
10.426
0.137
0.082
0.000
0.101
8.378
0.002
0.259
0.55B
10.509
0.256
205
0.151
0.000
0.115
0.002
16
2048
50K
0.708
512
165
1
11,254
0.177
0.080
0,001
0.139
9.241
0.002
0.257
45

---

Published as a conference paper at ICLR 2025
Table 19: Hypothetical generation speedup of Recursive Transformers across three models. We utilized the measurements of tokens per second calculated in Table 18. We only considered the time spent within Transformer blocks, simulating generation on the SlimPajama, RedPajama, and PG19 test sets. We used a vanilla transformer model, both with and without continuous sequence-wise batching, as our baselines. Our Recursive models further enhance throughput by applying continuous depth-wise batching, leveraging looping and early-exiting techniques. The throughput improvements over the vanilla Transformer and sequence-wil wise batching are denoted as Ay and Aseg. respectively. To aid in understanding the speedup, we also provide the performance of intermediate layers and the maximum batch size.
Uptrain
Looping
LoRA
Farty-Rait Train
Batching
Feu-shot Accuracy
Thronggoott 1
Models
Noemb]
PT
Nu.s
Block
Rank
Init
No.
CF
KD
Type
Exil
Last
Mil
Mid2
Butch
SlimP
RedP
PG19
Av
Auvs
1,998
73B
57.3
111
1740
3159
4796
x1.00
x0.63
1,998
73B
CSB
57.3
111
5287
5060
4796
x1.68
x1.00
0.908
V
60B
Step
15B
Agg I)
CDB
54.0
48.8
43
31.59
3050
2421
x2.50
x1.59
1,07
60B
Aug
64
SVD
15B
(02(0.3)
CDB
54.0
40.8
41
2357
2255
1858
x1.87
x1.19
1.150
600
Avg
128
SVD
150
(gz(0..)
CDB
54.6
40.2
41
2355
2250
1846
x1.87
x1.19
1,508
60B
256
SVD
158
(gx(0..)
CDB
55.2
40.5
39)
2047
1976
1740
x1.86
x1.18
1,608
64B
512
SVD
15B
(0.3)
CDB
56.2
41.7
39
1806
1754
1598
x1.75
x1.10
1,070
V
600
64
SVD
150
(0.3
CDB
S3.1
43.3
2454
2357
1929
<1.95
x1.24
1,158
60B
128
SVD
(0.3)
CDB
83.6
43.4
2445
2546
1926
x1.24
1,308
6BB
Agz
54.6
2156
60R
(05
0.64R
Aps
43.5
3130
1041
Agg
0,620
Ags
0.788
64B
Ags (0-1)
54.1
42.2
2261
2219
x1.86
x1.18
1,27R
<IR
Avg
SVD
Agg (0 1)
CDB
7
43.5
37.0
2181
2122
x1.02
x1.08
8,748
6BB
Avy
SVD
Ags 31
CDB
42.9
37.7
2437
2372
x1.94
x1.25
0,829
60B
AVE
128
SVD
15B
Aex
CDB
51.2
43.2
38.0
43
2576
2521
2084
x1.97
x1.25
0,978
60
ANE
256
SVD
15D
(02
CDB
53.0
44.8
38.7
2359
2500
2009
x1192
x1.22
1,278
608
Avg
512
SVD
158
Agg (0.3)
CDB
54.1
45.7
39.2
39
2251
2191
1975
x1.67
x1.06
0,970
/
43.3
1049
6856
7581
5090
X:1.00
x0.96
0.970
CSD
43.3
1049
7709
7.5
4090
x1105
x1.00
0.458
V
6BB
Step
158
Age (1.3)
CDB
44.8
41.8
233
2018
2023
1933
x1.70
x1.62
0,538
<4B
64
SVD
158
Agx (1.3)
CDB
45.1
33.7
211
1733
1719
1617
x1.38
x1.82
0,53R
60R
128
SVD
15R
(0 11
CDB
45.2
33.9
209
1733
17177
1609
x1 11.36
x1.30
Liuma
0,640
60B
Agz
CDB
33.9
1728
1714
0,860
600
Agg (0.1)
CDB
34.2
1716
1702
x1.28
x1.23
0.53R
Agg
0.54R
Agg
CDB
0,680
AYE
0.868
606
0818
/
738
5346
x1.00
0.8IR
/ 73R
CSR
5813
5734
x1.13
x1.00
0,400
V
60B
Step
Ago ()
CDB
42.0
1330
1333
1281
x1.77
x1.57
0,419
60B
Avg
64
SVD
15B
Agx
CDB
46.1
37.1
63
1206
1203
1140
x1.44
x1.28
0.488
60R
Ang
128
SVD
158
Ags (0:1)
CDB
46.2
37.8
59
1156
1180
1108
x1.32
x1.17
0.558
6088
Ang
256
SVD
158
Ags (0.3)
CDS
46.5
38.0
59
1138
1139
1071
x1.22
x1.08
0.700
60B
Avg
512
SVD
150
Agg (0.1)
CDB
47.2
38.2
53
1051
1077
1021
x0.98
x0.87
0.640
/
600
Avg
64
SVD
150
(03
/
CDB
65.1
39.0
63
1254
1252
1190
x1.50
x1.33
0.455
128
SVD
158
Aex (0.3)
CDB
45.9
39.0
59
1200
1226
1155
x1.37
x1.22
0.539
256
SVD
158
(0.3)
CDB
46.0
39.4
59
1150
1180
1112
x1.27
x1.12
0.70R
40R
512
SVD
158
Agz (0-3)
CDB
46.7
39.7
53
1058
1114
1058
x1.02
X0.90
46


Reviews:
[{"review_id": "OplF5AkT2M", "reviewer": "Reviewer_WB6T", "summary": "The paper introduces a new approach to reducing the deployment costs of large language models by implementing recursive parameter sharing. The proposed Recursive Transformer model reuses a single block of unique layers across multiple iterations, effectively minimizing model size while maintaining strong performance. To enhance flexibility, Relaxed Recursive Transformers incorporate low-rank adaptation (LoRA) modules at each layer, allowing for slight layer-specific modifications that improve performance with minimal overhead. Additionally, the authors propose a Continuous Depth-wise Batching technique, which leverages the recursive architecture to increase inference throughput by enabling multiple depths to be processed simultaneously, achieving up to a 2-3× improvement in efficiency over traditional models.", "questions": "1. In experiment design (Line 268), authors mention “As a result, we set the target performance for Gemma and Pythia models as the performance achieved by uptraining a full-size pretrained model with an equivalent number of tokens.”. This design is unreasonable, especially since your training (SlimPajama) and testing datasets (llm harness evaluation) are different. You need to provide the original LLM’s performance to help readers understand the real performance of the proposed recursive transformer instead of fine-tuning a worse model to mislead readers.\n2. To test the effectiveness of your model compression method, it is highly recommended to achieve a larger compression ratio. E.g. compressing 7B model to 1B. Current experiments only compress 2B models to 1B, which lacks practical meaning as they are all models with similar sizes.\nReusing the transformer computation block has been explored before [1] without properly cited within this paper.\n\n[1] Liu, Zechun, Changsheng Zhao, Forrest Iandola, Chen Lai, Yuandong Tian, Igor Fedorov, Yunyang Xiong et al. \"Mobilellm: Optimizing sub-billion parameter language models for on-device use cases.\" arXiv preprint arXiv:2402.14905 (2024).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Propose a novel model structure that can effectively compress pre-trained large language models for deployment.\n2. LoRA adaptation makes the model structure more flexible to adapt to different depths.\n3. Depth-wise Batching translates the advantages of the proposed recursive transformer block to real deployment acceleration.", "weaknesses": "1. Experiment design is not so convincing, which can be further improved.\n2. Authors do not push the recursive computation block’s advantage to the limit, which can be interesting to explore\n3. Some highly-related papers are missing in the reference", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uB9Dj6Xz7m", "reviewer": "Reviewer_9eTH", "summary": "The paper presents an innovative approach to parameter efficiency in Large Language Models through Recursive Transformers, by iteratively reusing a single block. The concept is extended to Relaxed Recursive Transformers by incorporating layer-specific LoRA modules, introducing flexibility to the parameter-sharing scheme. The authors demonstrate the conversion of pre-trained LLMs into more compact recursive versions while maintaining performance and achieving improved throughput through depth-wise batching.", "questions": "- The choice of LoRA for relaxing parameter sharing appears somewhat arbitrary. And prefix-tuning results show limiting performance gain. Could you elaborate more on the relaxation technique selections and other alternatives (*e.g.* Adapter, IA3)?\n- For Section 3.8, how would the throughput gains hold up when using practical confidence-based early-exiting rather than the idealized oracle approach that assumes perfect prediction matching?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The core concept of recursive layer sharing combined with LoRA-based relaxation is straightforward yet powerful.\n- The paper provides extensive empirical validation across multiple model architectures, recursive layer initialization strategies, LoRA initialization strategies, and evaluation metrics, including thorough ablation studies and practical deployment.\n- The authors not only show the benchmarking results of Relaxed Recursive Transformers but also explore hypothetical training speedup with continuous depth-wise batching strategy and early-exiting.", "weaknesses": "**Major**\n- It compares to different-sized dense baselines. However, the evaluation lacks comparison with sparse designs, particularly layer-skipping mechanisms, and layerwise parameter-sharing approaches [1-2].\n- The paper does not fully investigate how recursive architectures affect model training dynamics and optimization landscapes. The relationship between shared parameters and gradient flow during the uptraining phase. Also, whether shared parameter amounts affect model learning capacity could be explored.\n- The study would benefit from experiments on optimal recursion depths across varying model scales and task domains.\n\n**Minor**\n- The experiments are primarily conducted on relatively small models with 1-2B parameters. The effectiveness and scalability of the approach on larger models remain unproven. This raises questions about whether the benefits would hold at scales where current LLMs demonstrate more sophisticated capabilities.\n- Typos: \"achieved a 10%p performance\" on page6.\n\n[1] Mixture-of-Depths: Dynamically allocating compute in transformer-based language models\n\n[2] Zamba: A Compact 7B SSM Hybrid Model", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NShMLwgXlA", "reviewer": "Reviewer_uD6H", "summary": "This paper introduces a novel framework for parameter sharing in LLMs, namely relaxed recursive transformers. The method revist the way of layer tying and further convert LLMs into recursive transformers by sharing parameters across layers. Moreover, the performance is further improved by relaxing the layer tying constraint. The method is demonstrated to be well performed on various LLMs, it also get significant throughput gains by paring with early exiting, namely continuous depth-wise batching.", "questions": "1. I wonder the performance of relaxed recursive transformers on larger scales.\n2. How does the authors compare the proposed parameter sharing methods with other model deployment optimization methods, such as quantization, sparcification，pruning or distillation?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well motivated on a hot topic in the field of LLMs.\n2. The paper is well organized and the insight of the method is clearly stated.\n3. The experimental results are convincing. I appreciate the way the authors present the takeaways for relaxed recursive transformers.", "weaknesses": "1, It lacks evaluation on larger scales, which makes the scalability of the method remain questionalbe. For instance, I expect results on LLMs with 7B, 70B or larger scales.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6z8LWGC78k", "reviewer": "Reviewer_LZ2e", "summary": "This paper revisits the concept of \"layer tying\" as a form of parameter sharing in Transformers for developing smaller models. These models utilize a single block of unique layers repeated in a loop derived from pre-trained models to reduce model size while maintaining high performance. The paper presents Relaxed Recursive Transformers incorporating depth-wise LoRA modules to enhance these recursive models. These modules add flexibility to the layer-tying process but maintain the model's compactness. Furthermore, the paper proposes a new inference technique, Continuous Depth-wise Batching, which leverages the recursive nature of these models combined with early exiting strategies to increase throughput by 2-3 times.", "questions": "Please refer to the weaknesses section.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The introduction of Recursive Transformers with a single repeated layer block offers a novel approach to reducing model size and complexity without substantial loss in performance.\n\n2. The incorporation of Relaxed Recursive Transformers with LoRA modules successfully improves flexibility and performance.\n\n3. The proposed Continuous Depth-wise Batching method, which takes advantage of the recursive model structure and early exiting, increases inference throughput.", "weaknesses": "1. The approach of uptraining a full-size pre-trained model is questionable, especially given that the goal is to achieve comparable performance to the original model. In practical applications, deployment considerations often prioritize model size rather than training data size or consumption. Using a smaller, lower-quality dataset for these experiments may unfairly favor smaller (1B) models, potentially skewing the results and making the claimed improvements less persuasive. \n\n2. Furthermore, this uptraining step could be prohibitive for settings where on-device retraining is required. The reliance on extensive uptraining, which consumes 60 billion tokens, presents challenges for practical deployment, particularly when computational resources are constrained.\n\n3. The current methodology appears to train reduced models from scratch, using the same number of tokens, while recursive transformers are initialized with pre-trained layers. This discrepancy in training procedures introduces a potential bias in favor of the recursive transformers.\n\n4. The paper lacks a clear rationale for determining the number of initial layers to be recursively repeated and the frequency of such recursions. Furthermore, an automated approach for selecting the optimal recursion size would greatly increase the novelty and technical contribution of the proposed method.\n\n5. All experiments are conducted with models around the 1B parameter mark, yet hardware advancements make models larger than 1B feasible for many applications. Testing the proposed approach on larger models is essential to validate its scalability and effectiveness beyond the 1B parameter range, ensuring the method’s relevance for state-of-the-art large models. \n\n6. The selection of layers for initializing looped layers, as indicated in Figure 5, appears to depend heavily on the specific model used. A clear explanation of how to automate this layer selection process would substantially improve usability and consistency across different models. Without this, the method’s applicability may be limited by model-specific manual tuning. \n\n7. The impact of using different base models (other than a 2B model) for recursive transformer initialization has not been evaluated. Testing a range of base model sizes would provide insight into how the choice of initial model affects recursive initialization outcomes, offering a more thorough assessment of the method’s robustness.\n\n8. The experiments focus exclusively on few-shot tasks, but in real-world applications, 1B models are often fine-tuned for specific tasks. Testing the proposed method on downstream QA or classification tasks with full training on task-specific data would provide a more realistic evaluation of its performance and applicability across different use cases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)