Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA

## Abstract

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.

## 1 Introduction

Efficient deployment of large language models (LLMs) demands a balance between performance and resources (Raposo et al., ; Leviathan et al., ; Rivière et al., ; Wan et al., ; Zhou et al., ).
While larger models with more parameters consistently demonstrate superior performance (Rosenfeld et al., ; Rae et al., ; Hoffmann et al., ), their substantial memory and computational demands are expensive (Pope et al., ). Parameter sharing approaches (e.g. Dehghani et al., ; Xia et al., ; Lan et al., ; Takase and Kiyono, ), 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., ; Giannou et al., ), its application to modern LLMs has yielded limited reported success.

In this work, we revisit parameter sharing for LLMs, and propose novel methodologies to *convert* existing, unshared models into smaller, and more efficient, Recursive Transformers. These models use a single block of unique layers that are recursively reused across multiple loops, yet still achieve impressive performance relative to their reduced size.
To mitigate the potential performance degradation associated with parameter sharing, we first initialize the shared block of layers based on the original model’s pre-trained parameters, and then finetune the resulting recursive model for a limited number of “uptraining” steps. Importantly, we show that our initialization strategies allow us to achieve strong performance with minimal training time.
This is aligned with observations that model compression techniques such as layer skipping (Zhang et al., ; Zeng et al., ; Fan et al., ; Elhoushi et al., ), pruning (Frankle and Carbin, ; Ramanujan et al., ) or nesting (Devvrit et al., ) can preserve surprisingly high performance—further motivating our approach of compressing models to more compact yet performant architectures (here, repeated layers with low-rank adapters).

As depicted in Figure 1, we further propose the Relaxed Recursive Transformer,
an 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., ).
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, ) on the residual matrices between 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 Batching, that their recursive nature enables. Prior work introduced continuous sequence-wise batching (Yu et al., ; Kwon et al., ), which leverages the fact that the computation performed to compute a 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 a 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 computing different iterations of the looped layer blocks per individual responses within the same batch.

Figure: Figure 1: 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 is 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.
Refer to caption: x1.png

Our key contributions are as follows:

- •
We introduce a framework for initializing and training Relaxed Recursive Transformers and demonstrate strong performance compared to non-recursive models of comparable size.
For example, when we uptrained a recursive Gemma 1B model converted from a pretrained Gemma 2B (Team et al., ), 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). Furthermore, we show that by incorporating knowledge distillation (Hinton et al., ; Kim and Rush, ), 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 a key use case for continuous depth-wise batching with early-exiting (Bae et al., ; Schuster et al., ; Elbayad et al., ; Graves, ), which opportunistically makes predictions for samples with high confidence at earlier stages. From our simulation, Early Exits reveal a substantial throughput improvement of up to 2-3$\times$ compared to a vanilla Transformer with the same architecture.
Notably, the recursive Gemma model, which outperforms the vanilla Pythia model, can theoretically achieve a nearly 4$\times$ increase in throughput (see §3.8 for details).

## 2 Effective Model Compression with Recursive Patterns

In this section, we present the main details of our method for converting a vanilla Transformer model into a parameter-shared model that outperforms models of equivalent size. We first provide a short overview of the Transformer architecture ($\S$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 ($\S$2.2).
In $\S$2.3, we relax the parameter-sharing constraint in the model design, and add a limited set 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 ($\S$2.4).

### 2.1 Basic Transformer Architecture

Large language models (Rivière et al., ; Reid et al., ; OpenAI, ; Dubey et al., ) typically leverage the Transformer architecture (Vaswani et al., ). A Transformer consists of L layers, where the hidden states at each time step t are computed by running through the series of layers:

$$ $\displaystyle\mathbf{h}^{\ell}_{t}=f(\mathbf{h}^{\ell-1}_{t};\,\Phi_{\ell}),\, \,\ell\in[1,L],$ (1) $$

with $\mathbf{h}^{0}_{t}$ representing the embedding of the token $y_{t-1}$ from the previous time step, and $\Phi_{\ell}$ denoting the trainable parameters of the $\ell$-th layer.

Each layer has two core components: a multi-head attention (MHA) mechanism and a 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 A for further details.

### 2.2 Recursive Transformer: Looped Layer Tying

In this work, we revisit parameter sharing in the context of LLMs and propose the Recursive Transformer architecture.
Among various looping strategies (refer to Appendix B), we specifically adopt the CYCLE strategy (Takase and Kiyono, ) 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:

$$ $\displaystyle\textbf{h}_{t}^{\ell}=f(\textbf{h}_{t}^{\ell-1};\,\Phi^{\prime}_{ ((\ell-1)\bmod L/B)+1}),\,\,\ell\in[1,L],$ (2) $$

where the parameter-shared model is parameterized by $\Phi^{\prime}$, and B denotes the number of looping blocks (we restrict B to be a factor of L).
For example, Gemma 2B (Team et al., ) with 18 layers can 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 and Sennrich, ).

Figure: Figure 2: Left: An example of unshared, full-size model with 6 layers. Middle: Three proposed methodologies for initializing looped layers in a 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 the Average method.
Refer to caption: x2.png

#### Initialization techniques for looped layers

To mitigate the potential performance drop associated with reduced capacity in parameter-shared models, we propose several novel initialization methodologies to facilitate effective knowledge transfer from unshared, pretrained models to Recursive Transformers.
Figure 2 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., ; Zhang et al., ; Zeng et al., ; Fan et al., ) showing minimal impact on generation quality when skipping a few layers in LLMs.
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 ($\S$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 size while maintaining strong capabilities, it has two noticeable limitations: (1) the set of possible model sizes is limited to 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., ; Houlsby et al., ) for each layer, relaxing the strict parameter sharing.
While we experiment with various approaches like layer-specific prefixes (Liu et al., ) (see Appendix H), we find low-rank adaptation (LoRA) modules (Hu et al., ) to efficiently capture the subtle variations between tied layers. Specifically, we modify Eq. 2 to:

$$ $\displaystyle\textbf{h}_{t}^{\ell}=f(\textbf{h}_{t}^{\ell-1};\,\Phi^{\prime}_{ ((\ell-1)\bmod L/B)+1},\Delta\Phi^{\prime}_{\ell}),\,\,\ell\in[1,L],$ (3) $$

where $\Delta\Phi^{\prime}$ is the (small) set of parameters for the LoRA modules.

In this relaxed model, each looped layer is augmented with multiple LoRA modules. For example, a recursive model with two loop iterations has a 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 $\textbf{h}=\textbf{W}^{\prime}\textbf{x}$, our modified forward pass yields $\textbf{h}=\textbf{W}^{\prime}\textbf{x}+\Delta\textbf{W}^{\prime}\textbf{x}=
\textbf{W}^{\prime}\textbf{x}+\textbf{B}\textbf{A}\textbf{x}$, where $\textbf{A}\in\mathbb{R}^{(r\times k)}$ and $\textbf{B}\in\mathbb{R}^{(d\times r)}$ 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 $\S$2.2, we aim for the sum of the tied weights ($\Phi^{\prime}$) and LoRA weights ($\Delta\Phi^{\prime}$) to approximately recover the full-size model’s weights ($\Phi$).
We exploit truncated Singular Value Decomposition (SVD) (Hansen, ) on residual matrices between original weights and tied weights:

$$ $\displaystyle\mathbf{U}_{r}^{\ell},\mathbf{\Sigma}_{r}^{\ell},\mathbf{V}_{r}^{ \ell}=\text{Truncated SVD}(\mathbf{W}_{\ell}-\mathbf{W^{\prime}}_{((\ell-1) \bmod L/B)+1};\,\,r),\,\,\ell\in[1,L],$ (4) $$

where outputs retain the first r columns corresponding to the r largest singular values. W denotes the weight matrices of the full-size model, and $\text{{{W}}}^{\prime}$ denotes those of the Recursive Transformer.
We initialize the LoRA’s weights with principal components in Eq. 4: B as the product of $\text{{{U}}}_{r}$ and $\mathbf{\Sigma}_{r}$, and A as the transpose of the right singular vectors $\text{{{V}}}_{r}$ (see Figure 2).

###### Remark .

By initializing LoRA weights through the proposed truncated SVD methodology, the rank of the LoRA modules serves as a pivotal hyperparameter, enabling the Relaxed Recursive Transformer to seamlessly transition between the two extremes of the vanilla and Recursive Transformer architectures.

With sufficiently large ranks, our Relaxed Recursive Transformer (Eq. 3) approximates the full-size vanilla model (Eq. 1):

$$ $\displaystyle\mathbf{W}\textbf{x}\approx\mathbf{W}^{\prime}\textbf{x}+(\mathbf {U}_{r}\mathbf{\Sigma}_{r})(\mathbf{V}^{\top}_{r})\textbf{x}=\mathbf{W}^{ \prime}\textbf{x}+\mathbf{B}\mathbf{A}\textbf{x}=\mathbf{W}^{\prime}\textbf{x} +\Delta\mathbf{W}^{\prime}\textbf{x},$ (5) $$

Meanwhile, setting the rank to zero reduces the model to a Recursive Transformer, as the LoRA modules contribute no additional parameters, highlighting the flexibility of this relaxation approach.

###### Remark .

### 2.4 Continuous Depth-wise Batching and Early-Exiting

Figure: (a) Vanilla Batching
Refer to caption: x3.png

In real-world deployments, user requests arrive sequentially and asynchronously. Recent research has introduced continuous sequence-wise batching (Yu et al., ; Kwon et al., ), a 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 a new token is functionally the same and utilize the same model parameters. By continuously scheduling requests in this manner, models can operate at 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 a 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 a maximum batch size of 32, a 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., ; Miao et al., ) through faster scheduling.

Throughput improvements from depth-wise batching are further amplified when combined with early-exiting (Bae et al., ; Schuster et al., ; Elbayad et al., ). As depicted in Figure 3(c), once some samples exit after certain looping iterations, queued requests can then be immediately scheduled. While Recursive Transformers leverage the speedup from early-exiting, they also inherently address a 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$\times$ speedup on evaluated LLMs.

**Table 1: 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, ) and Grouped-Query (Ainslie et al., ) attention mechanisms, which leads to a reduced number of key-value heads. ^∗We take an early TinyLlama checkpoint to study recursive conversions on top of an under-trained model on SlimPajama. The vanilla performance with longer pretraining is reported in Table D.1.**
|  | Model Architecture | Pretraining |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | N-emb | Emb | $N_{L}$ | $d_{model}$ | $N_{head}$ | $N_{KV}$ | $d_{head}$ | Vocab | Dataset | $N_{tok}$ | $L_{ctx}$ |
| Gemma 2B | 1.98B | 0.52B | 18 | 2048 | 8 | 1 | 256 | 256K | Unreleased | 3T | 8K |
| TinyLlama 1.1B | 0.97B | 0.13B | 22 | 2048 | 32 | 4 | 64 | 32K | SlimPajama + | 73B^∗ | 2K |
| Starcoderdata | 32B |  |  |  |  |  |  |  |  |  |  |
| Pythia 1B | 0.81B | 0.21B | 16 | 2048 | 8 | 8 | 256 | 50K | Pile | 300B | 2K |

## 3 Experiments

### 3.1 Experimental Setup

We evaluate our method on three popular pretrained LLMs: Gemma 2B (Team et al., ), TinyLlama 1.1B (Zhang et al., ), and Pythia 1B (Biderman et al., ).
Table 1 summarizes each model’s architecture and pretraining recipes, and their few-shot performance is summarized in Appendix D.
Unless stated otherwise, the number of looping blocks (B) is set to 2 for all experiments. The results for Gemma with $\text{{B}}=3$ is provided in the appendix.
After converting to Recursive Transformers, we uptrained models on the SlimPajama dataset (Soboleva et al., ).
We used the Language Model Evaluation Harness framework (Gao et al., ) to evaluate accuracy on seven few-shot tasks, and averaged them for performance comparison.
Detailed experimental setup can be found in Appendix E.

**Table 2: Uptraining the pretrained models on datasets that differ significantly in quality or distribution from their pretraining datasets can lead to decreased performance. We evaluated models after uptraining on the SlimPajama dataset. We measured perplexity on test sets of the SlimPajama, RedPajama, and PG19, and few-shot accuracy on LAMBADA, HellaSwag, PIQA, WinoGrande, ARC-easy, ARC-challenge, and OpenBookQA benchmarks.**
|  |  | Uptrain | Perplexity $\downarrow$ | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | N-emb | PT | $N_{tok}$ | SlimP | RedP | PG19 | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg |
| Gemma 2B | 1.99B | ✓ | - | 11.46 | 8.18 | 13.52 | 63.1 | 71.4 | 78.1 | 65.0 | 72.3 | 41.9 | 40.2 | 61.7 |
| 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 |  |
| TinyLlama 1.1B | 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.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 |  |
| Pythia 1B | 0.81B | ✓ | - | 15.68 | 9.90 | 12.05 | 57.5 | 49.1 | 70.4 | 52.8 | 51.9 | 26.7 | 33.4 | 48.8 |
| 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 |  |

### 3.2 Non-Recursive Model Baselines

Given that we leveraged pretrained model weights for initialization and subsequently uptrained the models, it becomes crucial to define clear performance targets for our parameter-shared models.

#### Full-size model

Our ultimate goal is for the Recursive Transformer to achieve performance comparable 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 on certain benchmarks.
Table 2 summarizes 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 is unreleased but potentially well-curated (Team et al., ), all few-shot performance metrics gradually decreased after uptraining on 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, 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.
Since TinyLlama was already pretrained on SlimPajama—which 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 are pretrained from scratch on the same training recipe (number of training tokens and distillation from full-size model), but without the benefits of the pretrained weights and the looping mechanism. This comparison serves to highlight the efficacy of our initialization techniques and the recursive function itself in attaining strong performance, even with a constrained model size.

### 3.3 Main Results

Figure: (a) Gemma
Refer to caption: x6.png

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.4% vs. 58.6%), despite using fewer parameters and uptraining on 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 more streamlined models.

In the subsequent sections, we provide a comprehensive overview of extensive ablation studies conducted prior to achieving this final performance. In §3.4, we delve into the analysis of various initialization methodologies for Recursive Transformers. Insights into the relaxation model are detailed in §3.5. Finally, we explore enhanced training strategies like knowledge distillation (§3.6).

### 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 5(a), and their few-shot performance in Figure 5(b).
Our proposed methods significantly outperformed random initialization, which simply adds recursion to a reduced-size model, suggesting that leveraging pretrained weights in any manner is beneficial for performance boost. Moreover, the Stepwise methodology consistently demonstrated best performance, aligning with insights that LLMs can preserve performance even with a few layers skipped (Raposo et al., ; Zhang et al., ; Elhoushi et al., ).
Interestingly, as summarized in Table F.1, the recursive TinyLlama model, uptrained on only 15 billion tokens, yields few-shot performance comparable to the original model pretrained on 105 billion tokens.
This suggests that with sufficient training, even a recursive architecture can match the performance of a full-size pretrained model (Dehghani et al., ; Takase and Kiyono, ).

Figure: (a) Loss curves for Gemma
Refer to caption: x9.png

#### Recursive Gemma 1B outperforms both pretrained TinyLlama 1.1B and Pythia 1B

The looped Gemma 1B model, utilizing our proposed Stepwise method, outperformed reduced-size baselines with equivalent parameter counts by up to 13.5 percentage points (51.7% vs. 38.2%). Furthermore, it even outperformed the full-size TinyLlama 1.1B and Pythia 1B models (see Figure 5(c)). This is a 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 a promising starting point, as their recursive counterparts readily outperform other ordinary vanilla models of similar size.
Further details can be found in Appendix F.

### 3.5 Relaxation of Strict Parameter Sharing via LoRA Modules

Figure: (a) Loss changes in Gemma
Refer to caption: x12.png

#### Average initialization for looped layers is most compatible with Relaxed Recursive Transformer

Figures 6(a) and 6(b) illustrate the effect of relaxing parameter sharing via layer-wise LoRA modules.
Notably, 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 is likely because capturing the nuances between entirely distinct layer weights is challenging with an insufficient rank, leading to a suboptimal solution.
Further details are provided in Appendix G.

#### 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, as shown in Figure 6(c).
Our proposed SVD strategy demonstrated an impressive performance boost of up to 6.5 points, facilitating faster convergence by updating the principal low-rank matrices (aligned with findings in Meng et al. ()).
For results across other architectures, refer to Figure G.2.

#### 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 7(a), performance generally improves with increasing rank, resulting in a clear Pareto frontier between model size and performance.
However, only Stepwise initialization showed a U-shaped performance trend: a 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 G.2.

Figure: (a) Performance by LoRA rank
Refer to caption: x15.png

### 3.6 Extended Uptraining and Knowledge Distillation

We further enhanced the performance of our low-rank models by introducing two techniques: uptraining 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 a forward Kullback-Leibler divergence (Hinton et al., ; Kim and Rush, ), which exhibited the best performance gains among the examined distillation losses. Table I.1 summarizes the results of various ablation studies conducted to investigate the impact of these two techniques.

The combined effect of these techniques is presented in Figure 7(b), demonstrating an improvement of up to 4.1 percentage points in few-shot accuracy compared to the previous 15 billion token uptraining results. Notably, the relaxed Gemma model with a rank of 512 nearly matched the performance of the full-size model.
We also expect that further performance gains can be achieved with a much lighter recursive model by utilizing a superior teacher model or conducting more extensive training on high-quality data.
Figure 7(c) 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 to a recursive architecture. Additional details can be found in Appendix I.

**Table 3: A small loss coefficient to the first loop output (intermediate output) can significantly improve intermediate performance without compromising the final performance. Performance was evaluated under a static-exiting scenario (Schuster et al., ), where all tokens exit at either first or second loop. We further trained the previously uptrained Gemma models on 15 billion tokens (post-training). Delta ($\Delta$) denotes the performance changes in the final outputs after early-exit training.**
|  | Uptrain | Looping | Early-Exit Train | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| N-emb | PT | $N_{tok}$ | Block | Init | $N_{tok}$ | CE | KD | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg | $\Delta$ |
| 0.99B | ✓ | 15B | 2 | Step | - | - | - | 53.0 | 57.3 | 73.2 | 56.2 | 56.1 | 29.2 | 36.6 | 51.7 | - |
| 0.99B | ✓ | 15B | 2 | Step | 15B | Weighted | ✗ | 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 | - |  |  |  |  |  |  |  |  |
|  |  |  |  |  |  |  |  | 53.0 | 59.1 | 73.9 | 55.4 | 57.4 | 30.6 | 37.8 | 52.5 | +0.8 |
| 0.99B | ✓ | 15B | 2 | Step | 15B | Agg (0.1) | ✗ | 45.9 | 51.2 | 71.4 | 54.5 | 48.1 | 26.8 | 32.0 | 47.1 | - |
| 0.99B | ✓ | 15B | 2 | Step | 15B | Weighted | ✓ | 47.7 | 55.1 | 73.2 | 55.6 | 54.5 | 29.1 | 37.2 | 50.4 | – 1.3 |
| 48.3 | 54.9 | 72.1 | 55.9 | 54.3 | 28.4 | 35.4 | 49.9 | - |  |  |  |  |  |  |  |  |
|  |  |  |  |  |  |  |  | 52.9 | 58.9 | 73.7 | 55.7 | 57.5 | 31.1 | 38.2 | 52.6 | +0.9 |
| 0.99B | ✓ | 15B | 2 | Step | 15B | Agg (0.1) | ✓ | 46.3 | 52.1 | 71.6 | 55.3 | 49.2 | 28.5 | 32.6 | 48.0 | - |

### 3.7 Early-Exit Training Strategy for 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 to enable token prediction.
We conducted an ablation study on various strategies, as summarized in Table 3 (more detailed results are presented in Table J.1).
Directly applying the weighted CE loss ($\mathcal{L}=\sum_{i=1}^{B}\alpha_{i}\mathcal{L}_{i}\text{\; where \;}\alpha_{i
}=i/\sum_{i}i$) commonly used in prior works (Schuster et al., ; Bae et al., ) led to an overemphasis on the training of intermediate representations.
To address this, we employ an aggressive coefficient strategy that aggressively reduces the loss coefficient for intermediate outputs while maintaining a coefficient of 1 for the final output.
Our experiments demonstrated that an aggressive coefficient of 0.1, utilizing knowledge distillation from the detached final outputs (Bae et al., ), effectively preserves final performance while enhancing intermediate performance. Notably, the first loop output yielded only a difference of 4.6 percentage points in accuracy compared to the final output.
This underscores the potential to maximize the benefits of early-exiting in parameter-shared 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 J.
The aggressive coefficient strategy, combined with self-distillation, consistently achieved the best performance for intermediate outputs while maintaining strong performance for the final loop output across all models.
However, as the optimal strategy derived from the non-relaxed models was directly applied to the relaxed models, a more tailored training approach might further enhance the performance of intermediate loop outputs in Relaxed Recursive Transformers.

### 3.8 Hypothetical Generation Speedup via Continuous Depth-wise Batching

Figure: 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. $\Delta_{V}$ measures throughput relative to the vanilla Gemma model, while $\Delta_{Seq}$ measures throughput relative to the vanilla Gemma model with continuous sequence-wise batching (CSB).
Refer to caption: x18.png

#### 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., ; Bae et al., ).
This assumes that tokens exit at the earliest looping block where their prediction aligns with the final loop’s prediction.
We simulated the generation of language modeling datasets as if they were generated by our models, to obtain 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, we conducted simulations to estimate theoretical throughput.
Detailed explanations and limitations are discussed in Appendix K.

#### Continuous depth-wise batching paired with early-exiting can substantially boost throughput

Figure 8 illustrates the throughput of our proposed models and the vanilla Transformer across three architectures. 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 et al., ; Kwon et al., ).
In particular, Recursive models demonstrate up to 2.66$\times$ speedup in generation compared to vanilla counterparts.
Additionally, the recursive Gemma model significantly outperforms the vanilla pretrained Pythia model, with nearly 4$\times$ improvement in throughput.
Relaxed recursive models show a clear trade-off between achievable few-shot performance and throughput, modulated by the degree of relaxation through the LoRA ranks. This characteristic enables flexible model selection tailored to specific deployment scenarios.
Comprehensive results are presented in Tables K.2 and K.4.

## 4 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, ; Graves, ), CNNs (Eigen et al., ; Savarese and Maire, ; Guo et al., ), and the popular Transformer architecture. The Universal Transformer (Dehghani et al., ), 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., ), language modeling (Bai et al., ), and machine translation, through stacking a single layer (Dabre and Fujita, ), tying encoder (Milbauer et al., ) and decoder components (Xia et al., ), or partially tying layers (Takase and Kiyono, ). These methods often claim to achieve comparable performance with more compact models and increased computational speed, while also setting the ground for effective adaptive compute solutions (Dehghani et al., ; Graves, ; Schuster et al., ).

Concurrently, there has been growing interest in exploiting recurrent architectures for algorithmic or logical reasoning tasks (Saunshi et al., ).
Prior research (Schwarzschild et al., ; McLeish and Tran-Thanh, ) 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., ), to learn iterative algorithms for data-fitting problems (Yang et al., ), to achieve length generalization in algorithmic tasks (Fan et al., ), and promising theoretical potential for few-shot learning (Gatmiry et al., ).

However, previous work has predominantly focused on relatively small Transformer models, trained from scratch without leveraging pretrained model weights.
Our work distinguishes itself by investigating 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 in weight tying constraints through layer-specific modules.

In this paper, we also discuss how Recursive Transformers can be well suited for early-exiting techniques to 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 a synchronization issue with early-exiting, where the model must forward all layers if even a single token in a batch requires full processing (exited tokens must wait for them).
Several approaches attempt to 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 et al., ; Elbayad et al., ), SkipDecode (Del Corro et al., ), and parallel decoding (which can be combined with Speculative Decoding) (Bae et al., ; Elhoushi et al., ; Liu et al., ; Chen et al., ; Tang et al., ).
Nevertheless, the heterogeneous parameters across varying model depths still hinder the efficient progression of exited tokens to subsequent sequences.
In contrast, our Recursive Transformers enable parallel computation for tokens at different depths and sequences (in a continuous depth-wise batching paradigm)—also allow for parallel computation of missing KV caches with minimal overhead during the memory-bounded decoding phase.

## 5 Discussion and Future Work

In this work, we introduced Recursive Transformers, in which we compress LLMs via parameter sharing across recursively looped blocks of layers. Additionally, we presented a 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 a continuous depth-wise batching paradigm tailored for efficient serving systems of Recursive Transformers. We theoretically demonstrated that an oracle-exiting strategy can yield substantial throughput gains, reaching up to 2-3$\times$ speedup.
This work motivates further research on recursive patterns in modern LLMs such as:

#### Beyond hypothetical generation 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., ; Bae et al., ). Moreover, practical deployment needs to address decoding bottlenecks like key-value cache computation for exited tokens in remaining loops.
We leave the potential solutions, such as efficient confidence estimation or key-value cache sharing (Sun et al., ; Brandon et al., ), to 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., ; Wang et al., ), hindering parallel computation.
We concatenated LoRA weights into a single weight to improve efficiency over sequential computation, yet it introduces redundancy.
To mitigate this, we can explore optimized CUDA kernels for LoRA serving (Sheng et al., ; Chen et al., ) and parallelization across accelerators, inspired by distributed training for Mixture of Experts (Fedus et al., ; Gale et al., ).

#### 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-wise batching will yield considerable gains in serving efficiency.

## Acknowledgements

We thank Jacob Eisenstein for valuable feedback on an earlier version of 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, we thank the Google Cloud Platform for awarding Google Cloud credits for this project.

## Appendix A Components in Transformer Architecture

The Transformer block consists of two core components: a multi-head attention (MHA) mechanism and a 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:

$$ $\displaystyle\text{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V})=\text{softmax} \left(\frac{\mathbf{Q}\mathbf{K}^{T}}{\sqrt{d_{k}}}\right)\mathbf{V},$ $$

where $\mathbf{Q}$, $\mathbf{K}$, and $\mathbf{V}$ are linear projections of the input, parameterized by learned weight matrices $\mathbf{W}_{\ell}^{Q}$, $\mathbf{W}_{\ell}^{K}$, and $\mathbf{W}_{\ell}^{V}$, respectively. The outputs from each head of the multi-head attention are concatenated and then projected back to the original hidden size using a learned weight matrix $\mathbf{W}_{\ell}^{out}$.

While the FFN structure typically consists of two linear transformations, in the Gemma model, it deviates from this standard architecture as follows:

$$ $\displaystyle\text{FFN}(\mathbf{x})=\mathbf{W}_{\ell}^{down}(\text{GELU}( \mathbf{x}\mathbf{W}_{\ell}^{gate})*\mathbf{x}\mathbf{W}_{\ell}^{up})$ $$

with three learned linear weight matrices and a GeGLU activation (Shazeer, ).

## Appendix B Parameter Sharing Strategy

Takase and Kiyono () discuss three strategies for partial layer tying in Transformer models, as depicted in Figure B.1. The SEQUENE strategy is the simplest, assigning the same parameters to consecutive layers. The CYCLE strategy repeatedly stacks a 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.
We specifically utilized the CYCLE strategy due to its superior compatibility with early-exiting, which can amplify the throughput of Recursive Transformers through continuous depth-wise batching.

Figure: (a) SEQUENCE
Refer to caption: x19.png

## Appendix C Illustrative Examples of SVD Initialization in Relaxed Recursive Transformer

We propose an SVD initialization approach for LoRA modules within a Relaxed Recursive Transformer, effectively steering the summation of base and LoRA weights towards the pretrained weights of their corresponding depth. Figure C.1 illustrates an overview of how the LoRA module is initialized under three different initialization techniques (Stepwise, Average, and Lower) for looped layers.
One crucial point is 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 A and zero for B. For example, with the Stepwise method, the first loop’s LoRA module receives standard zero initialization, while the second loop’s LoRA is initialized using our proposed initialization.

Figure: Figure C.1: We visualize LoRA modules to show which residual matrices they target for initialization under three different looping initialization methods, assuming a full-size model with six layers and two looping blocks. For ease of understanding, A matrices are colored according to the full-size model weights at the corresponding depth, while B matrices are colored based on the looped layer weights. White B matrices indicate cases where the full-size model and looped model weights are identical, resulting in standard zero initialization.
Refer to caption: x22.png

## Appendix D Overview of Three Pretrained LLMs

We utilized three pretrained models—Gemma 2B (Team et al., ), TinyLlama 1.1B (Zhang et al., ), and Pythia 1B (Biderman et al., )—and converted them into Recursive Transformers.
Their corresponding few-shot performance results are presented in Table D.1.

**Table D.1: Few-shot performance of three pretrained models. Few-shot accuracy is measured on the LAMBADA, HellaSwag, PIQA, WinoGrande, ARC-easy, ARC-challenge, and OpenBookQA 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 under-trained model.**
|  |  |  |  | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | N-emb | Dataset | $N_{token}$ | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg |
| Gemma 2B | 1.99B | Unreleased | 3T | 63.13 | 71.38 | 78.13 | 65.04 | 72.26 | 41.89 | 40.20 | 61.72 |
| TinyLlama 1.1B | 0.97B | SlimPajama + | 105B | 43.26 | 42.23 | 66.81 | 53.35 | 44.74 | 23.21 | 29.20 | 43.26 |
| Starcoderdata | 503B | 48.92 | 49.56 | 69.42 | 55.80 | 48.32 | 26.54 | 31.40 | 47.14 |  |  |
| 1T | 53.00 | 52.52 | 69.91 | 55.96 | 52.36 | 27.82 | 33.40 | 49.28 |  |  |  |
| 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 1B | 0.81B | Pile | 300B | 57.52 | 49.10 | 70.40 | 52.80 | 51.89 | 26.71 | 33.40 | 48.83 |

This diversity offers several benefits. First, with three versions of recursive models, we can compare 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 to closely match the performance of 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 2 presents the evaluation results obtained after uptraining each of the pretrained models. While TinyLlama readily improves its performance due to uptraining on the same dataset, Gemma and Pythia show a decline in few-shot performance with SlimPajama uptraining, which can be attributed to the differences in data distribution and the lower quality of the uptraining dataset.

## Appendix E 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., ). SlimPajama is an open-source dataset designed for training large language models, which is created by cleaning and deduplicating the RedPajama dataset (Computer, ). The source data primarily consists of web-crawled data, along with data from Github, books, Arxiv, Wikipedia, and StackExchange.
We employed the HuggingFace training framework (Wolf et al., ) and enhanced memory efficiency through the Zero Redundancy Optimizer (ZeRO) (Rajbhandari et al., ) from the DeepSpeed library (Rasley et al., ), along with mixed precision training.
The context length was set to 2048, and the batch size was approximately 2 million tokens. We used the AdamW optimizer (Loshchilov and Hutter, ) with a learning rate of 2e-4, utilizing a cosine annealing learning rate scheduler (Loshchilov and Hutter, ). 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 underwent additional training on a 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 (SlimPajama dataset), 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.

#### Evaluation setting

We evaluated perplexity on test sets from three language modeling datasets: SlimPajama, RedPajama, and PG19 (Rae et al., ). Additionally, we used the Language Model Evaluation Harness framework (Gao et al., ) to evaluate accuracy on seven few-shot tasks: LAMBADA (Paperno et al., ), HellaSwag (Zellers et al., ), PIQA (Bisk et al., ), WinoGrande (Sakaguchi et al., ), ARC-easy and ARC-challenge (Clark et al., ), and OpenBookQA (Mihaylov et al., ). We adhered to the standard number of shots specified by the evaluation framework for each dataset. For few-shot datasets, excluding LAMBADA and WinoGrande, we normalized accuracy by the byte length of the target string. All evaluation performance measurements were conducted using a 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
(final throughput comparisons were based solely on the time spent within Transformer blocks.)
We tested two settings of prefix and decoding lengths (512 / 2048 and 64 / 256), calculating the per-token time by dividing the total elapsed time by the decoding length.
Using a single A100 40GB GPU, we recorded 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 prediction.
Since our models are not finetuned on any specific downstream tasks, we did simulation with three language modeling datasets (SlimPajama, RedPajama, and PG19) as if they were generated by our models.
For simplicity, we assumed a queue of 20K samples, rather than considering their arrival in static or dynamic time intervals.
We then recorded the exit loop of each token in these samples using the oracle-exiting algorithm.
With these two measurement (per-token generation time and exit trajectories), we present the hypothetical throughput of Recursive Transformers under various simulation scenarios.

## Appendix F Expanded Results of Initialization Methods for Looped Layers

#### Ablation study of Stepwise method

We initially hypothesized that the Stepwise method’s performance could be significantly influenced by the specific rule used for layer selection from the pretrained model. To investigate this, we conducted a controlled experiment (illustrated in Figure 1(a)), where layers were selected at 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 significant differences in training loss or few-shot performance, other models like Gemma exhibited 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 LLMs can yield performance benefits (Kim et al., ).

#### Ablation study of Average method

The Average initialization method exhibited notably poor performance, particularly when applied to the Gemma model. We hypothesized that this could be attributed to instability in the model’s learned distribution, potentially arising from averaging of normalization layer weights. To investigate this further, we experimented with three different methods for initializing normalization weights, as outlined in Figure 1(b): averaging weights (Norm-avg), selecting weights from a single layer (Norm-choice), and zero initialization (Norm-zero).
The performance trend observed among these methods varied across different model architectures. However, zero initialization of normalization layers resulted in a huge performance drop in certain architectures like TinyLlama and Pythia. Conversely, we observed no 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.

Figure: (a) Stepwise ablations
Refer to caption: x23.png

#### Overall comparison of training perplexity

Figure F.2 presents a comparative analysis of training loss across three model architectures and varying looping blocks, incorporating our proposed initialization methodologies. To set an upper bound on performance, we utilized a full-size model further uptrained on SlimPajama, accounting for the distribution shift between uptraining and pretraining data. Additionally, we trained a Recursive Transformer with a 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.

Figure: (a) Gemma (2 blocks)
Refer to caption: x26.png

#### Overall comparison of few-shot performance

Few-shot performance exhibited a consistent trend with training perplexity. Table F.1 provides a comparative summary of the proposed looping initialization methods against the full-size model, the reduced-size model, and Recursive Transformers utilizing random initialization. Moreover, Figure F.3 visually illustrates the performance differences across different few-shot datasets. Notably, the Stepwise method consistently demonstrated the best performance, showing a performance improvement of up to 14.1%p compared to random initialization.

**Table F.1: 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 ($\Delta$) show improvements over random initialization.**
|  |  | Uptrain | Looping | Perplexity $\downarrow$ | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | N-emb | PT | $N_{tok}$ | Block | Init | SlimP | RedP | PG19 | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg | $\Delta$ |
| 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 | - |
| 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.66B | ✗ | 15B | - | - | 24.44 | 21.69 | 36.03 | 27.2 | 30.6 | 63.8 | 50.5 | 40.6 | 22.0 | 27.0 | 37.4 | - |  |
| 0.99B | ✓ | 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 | ​​​+14.1 |  |
| 0.99B | ✓ | 15B | 2 | 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 | 2 | 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 | 2 | Rand | 22.66 | 20.06 | 32.86 | 27.4 | 31.6 | 63.4 | 50.5 | 39.7 | 21.9 | 28.8 | 37.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 | +9.9 |  |
| 0.66B | ✓ | 15B | 3 | 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 | 3 | 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 | 3 | Rand | 22.67 | 20.09 | 32.77 | 28.1 | 31.4 | 63.8 | 51.1 | 41.0 | 23.0 | 26.6 | 37.9 | - |  |
|  | 0.97B | ✓ | 15B | - | - | 12.26 | 9.37 | 11.94 | 43.3 | 42.2 | 66.8 | 53.4 | 44.7 | 23.2 | 29.2 | 43.3 | - |
|  | 0.48B | ✗ | 15B | - | - | 16.61 | 15.66 | 20.27 | 22.3 | 30.0 | 60.9 | 50.6 | 37.0 | 23.0 | 28.0 | 36.0 | - |
| TinyLlama | 0.48B | ✓ | 15B | 2 | Step | 11.61 | 9.89 | 13.00 | 39.6 | 39.8 | 66.5 | 52.9 | 44.3 | 24.9 | 30.6 | 42.7 | +6.2 |
|  | 0.48B | ✓ | 15B | 2 | Avg | 11.86 | 10.29 | 13.42 | 38.6 | 39.4 | 66.1 | 52.8 | 42.7 | 25.4 | 30.6 | 42.2 | +5.7 |
|  | 0.48B | ✓ | 15B | 2 | Lower | 14.67 | 12.67 | 16.68 | 31.9 | 32.3 | 62.6 | 52.0 | 39.1 | 22.1 | 27.8 | 38.3 | +1.8 |
|  | 0.48B | ✗ | 15B | 2 | Rand | 16.14 | 15.11 | 19.55 | 24.7 | 30.7 | 61.2 | 50.6 | 36.4 | 22.6 | 29.2 | 36.5 | - |
|  | 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.40B | ✗ | 15B | - | - | 25.69 | 20.00 | 32.08 | 24.3 | 30.0 | 61.9 | 50.7 | 38.3 | 22.3 | 26.0 | 36.2 | - |
| Pythia | 0.40B | ✓ | 15B | 2 | 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 | ✓ | 15B | 2 | Avg | 16.76 | 12.76 | 18.63 | 43.6 | 39.1 | 68.2 | 51.9 | 45.4 | 25.1 | 29.8 | 43.3 | +7.1 |
|  | 0.40B | ✓ | 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 | +7.4 |
|  | 0.40B | ✗ | 15B | 2 | Rand | 24.45 | 18.93 | 29.63 | 25.2 | 30.2 | 62.1 | 51.1 | 39.2 | 22.4 | 23.6 | 36.2 | - |

Figure: (a) Recursive Gemma with 2 blocks
Refer to caption: x31.png

## Appendix G Expanded Results of Relaxed Recursive Transformers

#### Training perplexity changes with LoRA modules

Figure G.1 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 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 learn the difference between the original pretrained weights and the averaged looped weights using low-rank matrices.

Figure: (a) Gemma
Refer to caption: x35.png

#### 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 G.2, 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.
This appears to be attributed to LoRA ranks being insufficient to adequately approximate the low-rank deltas across layers, resulting in initialization at a sub-optimal point.

Figure: (a) Gemma
Refer to caption: x39.png

#### Ablation study on 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 G.3, 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, a larger rank value results in enhanced performance. The Stepwise method, consistent with previous experimental findings, exhibited a U-shaped trend: with extremely low or high ranks, a clear performance increase results. However, with mid-range values, the approximation becomes less accurate, leading to a further performance decrease.

Figure: (a) Gemma
Refer to caption: x42.png

We further experimented with assigning different ranks to LoRA modules associated with each linear weights. Given the computational overhead inherent to LoRA modules, allocating different ranks to each module can offer a better balance between performance and computational efficiency. The experimental results in Table G.1 reveal a strong correlation between performance and overall model sizes.
Due to the substantial hidden dimension of the linear weights 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 drop.
It is intriguing that even minimal relaxation of key-value weights (achieved through small ranks) yielded comparable performance, despite the inherent strong sharing of key-value caches between attention heads in the Multi-Query attention structure (Ainslie et al., ).

**Table G.1: Evaluation results of relaxed recursive Gemma models with varying LoRA ranks for different 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 $\downarrow$ | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| N-emb | PT | $N_{tok}$ | Block | Init | Q | KV | Out | FFN | Init | SlimP | RedP | PG19 | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg |
| 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.99B | ✗ | 15B | - | - | - | - | - | - | - | 22.63 | 20.03 | 32.60 | 28.9 | 31.6 | 63.1 | 52.3 | 41.2 | 22.5 | 27.8 | 38.2 |
| 1.30B | ✓ | 15B | 2 | Avg | 256 | 256 | 256 | 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.28B | ✓ | 15B | 2 | 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 | 2 | 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 | 2 | 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 | 2 | 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 |
| 1.15B | ✓ | 15B | 2 | Avg | 128 | 128 | 128 | 128 | SVD | 12.52 | 10.07 | 15.51 | 56.1 | 58.2 | 72.3 | 55.8 | 57.1 | 30.7 | 37.2 | 52.5 |
| 1.14B | ✓ | 15B | 2 | 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 |
| 1.14B | ✓ | 15B | 2 | Avg | 128 | 64 | 128 | 128 | SVD | 12.72 | 10.18 | 15.76 | 55.5 | 57.7 | 72.7 | 57.0 | 56.9 | 30.1 | 38.2 | 52.6 |
| 1.08B | ✓ | 15B | 2 | 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 |
| 1.13B | ✓ | 15B | 2 | Avg | 64 | 64 | 64 | 128 | SVD | 12.77 | 10.29 | 15.95 | 55.2 | 57.4 | 73.0 | 56.7 | 56.5 | 30.5 | 37.2 | 52.3 |

#### Overall performance comparison of Relaxed Recursive Transformers

A comprehensive performance comparison is presented in Table G.2. 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, a clear trend of improved performance was observed with increasing rank.

**Table G.2: Overall evaluation results of Relaxed Recursive Transformers. Delta ($\Delta$) represent the accuracy differences between relaxed and non-relaxed models using the same looping initialization.**
|  |  | Uptrain | Looping | LoRA | Perplexity $\downarrow$ | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | N-emb | PT | $N_{tok}$ | Block | Init | Rank | Init | SlimP | RedP | PG19 | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg | $\Delta$ |
|  | 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.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.99B | ✓ | 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 | - |
|  | 1.07B | ✓ | 15B | 2 | Step | 64 | SVD | 12.76 | 10.21 | 15.99 | 52.1 | 57.2 | 73.0 | 57.8 | 56.9 | 28.9 | 36.8 | 51.8 | +0.1 |
|  | 1.15B | ✓ | 15B | 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 | 2 | 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.9 |
|  | 1.60B | ✓ | 15B | 2 | Step | 512 | SVD | 13.13 | 10.66 | 16.63 | 53.0 | 54.3 | 72.1 | 54.9 | 54.8 | 28.8 | 35.4 | 50.5 | – 1.2 |
|  | 1.60B | ✓ | 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.99B | ✓ | 15B | 2 | Avg | - | - | 15.15 | 12.57 | 19.86 | 43.6 | 47.4 | 70.4 | 52.6 | 50.5 | 27.8 | 34.4 | 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.15B | ✓ | 15B | 2 | 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.30B | ✓ | 15B | 2 | 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.60B | ✓ | 15B | 2 | 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 |
|  | 1.60B | ✓ | 15B | 2 | Avg | 512 | Zero | 13.78 | 11.31 | 17.71 | 49.8 | 52.4 | 71.7 | 53.3 | 51.2 | 29.4 | 35.0 | 49.0 | +2.3 |
|  | 0.99B | ✓ | 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 | 2 | Lower | 64 | SVD | 14.21 | 11.77 | 18.40 | 47.5 | 50.5 | 70.9 | 54.2 | 54.1 | 29.2 | 36.0 | 48.9 | +1.8 |
|  | 1.15B | ✓ | 15B | 2 | 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 | 2 | 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 |
|  | 1.60B | ✓ | 15B | 2 | Lower | 512 | SVD | 12.54 | 10.22 | 15.90 | 57.1 | 58.2 | 73.7 | 58.6 | 57.6 | 31.5 | 35.6 | 53.2 | +6.1 |
|  | 1.60B | ✓ | 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.48B | ✗ | 15B | - | - | - | - | 16.61 | 15.66 | 20.27 | 22.3 | 30.0 | 60.9 | 50.6 | 37.0 | 23.0 | 28.0 | 36.0 | - |
|  | 0.48B | ✓ | 15B | 2 | Step | - | - | 11.61 | 9.89 | 13.00 | 39.6 | 39.8 | 66.5 | 52.9 | 44.3 | 24.9 | 30.6 | 42.7 | - |
|  | 0.53B | ✓ | 15B | 2 | Step | 64 | 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.58B | ✓ | 15B | 2 | 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.3 |
|  | 0.68B | ✓ | 15B | 2 | 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.86B | ✓ | 15B | 2 | Step | 512 | SVD | 11.33 | 9.79 | 12.61 | 42.2 | 40.9 | 67.7 | 51.1 | 45.0 | 25.3 | 30.2 | 43.2 | +0.5 |
|  | 0.86B | ✓ | 15B | 2 | Step | 512 | Zero | 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.48B | ✓ | 15B | 2 | Avg | - | - | 11.86 | 10.29 | 13.42 | 38.6 | 39.4 | 66.1 | 52.8 | 42.7 | 25.4 | 30.6 | 42.2 | - |
|  | 0.53B | ✓ | 15B | 2 | 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 | 2 | 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.68B | ✓ | 15B | 2 | Avg | 256 | SVD | 10.71 | 9.18 | 11.82 | 44.1 | 43.2 | 68.1 | 53.5 | 44.4 | 25.7 | 30.4 | 44.2 | +2.0 |
|  | 0.86B | ✓ | 15B | 2 | Avg | 512 | SVD | 10.46 | 8.92 | 11.50 | 46.0 | 44.1 | 68.2 | 53.0 | 45.8 | 25.1 | 31.2 | 44.8 | +2.6 |
|  | 0.86B | ✓ | 15B | 2 | Avg | 512 | Zero | 11.28 | 9.75 | 12.69 | 41.5 | 41.0 | 66.8 | 53.2 | 44.8 | 25.5 | 31.2 | 43.4 | +1.2 |
|  | 0.48B | ✓ | 15B | 2 | Lower | - | - | 14.67 | 12.67 | 16.68 | 31.9 | 32.3 | 62.6 | 52.0 | 39.1 | 22.1 | 27.8 | 38.3 | - |
|  | 0.53B | ✓ | 15B | 2 | Lower | 64 | SVD | 13.68 | 11.77 | 15.48 | 35.5 | 34.0 | 63.8 | 51.0 | 40.0 | 24.6 | 28.0 | 39.5 | +1.2 |
|  | 0.58B | ✓ | 15B | 2 | 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 | 2 | 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 | 2 | Lower | 512 | SVD | 11.31 | 9.61 | 12.49 | 43.2 | 40.5 | 66.0 | 50.8 | 43.9 | 24.8 | 30.0 | 42.8 | +4.5 |
|  | 0.86B | ✓ | 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.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.40B | ✗ | 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 | 2 | Step | - | - | 16.38 | 12.37 | 17.74 | 43.4 | 40.5 | 67.4 | 50.8 | 46.3 | 25.7 | 30.0 | 43.5 | - |
|  | 0.44B | ✓ | 15B | 2 | Step | 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.48B | ✓ | 15B | 2 | 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 | 2 | 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.70B | ✓ | 15B | 2 | Step | 512 | SVD | 15.68 | 11.88 | 17.25 | 45.4 | 41.3 | 68.5 | 52.6 | 46.7 | 25.4 | 31.2 | 44.4 | +0.9 |
|  | 0.70B | ✓ | 15B | 2 | 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.40B | ✓ | 15B | 2 | Avg | - | - | 16.76 | 12.76 | 18.63 | 43.6 | 39.1 | 68.2 | 51.9 | 45.4 | 25.1 | 29.8 | 43.3 | - |
|  | 0.44B | ✓ | 15B | 2 | 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 | 2 | 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 | 2 | Avg | 256 | SVD | 15.22 | 11.54 | 16.47 | 48.5 | 43.3 | 67.2 | 51.4 | 46.7 | 25.5 | 32.0 | 44.9 | +1.6 |
|  | 0.70B | ✓ | 15B | 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 | +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.40B | ✓ | 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.44B | ✓ | 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.48B | ✓ | 15B | 2 | Lower | 128 | SVD | 16.63 | 12.49 | 18.17 | 45.2 | 39.2 | 66.8 | 51.0 | 45.6 | 24.9 | 29.6 | 43.2 | – 0.4 |
|  | 0.55B | ✓ | 15B | 2 | Lower | 256 | SVD | 15.93 | 11.99 | 17.30 | 47.6 | 41.4 | 68.3 | 53.2 | 46.0 | 25.8 | 31.0 | 44.8 | +1.2 |
|  | 0.70B | ✓ | 15B | 2 | Lower | 512 | SVD | 15.11 | 11.34 | 16.07 | 50.2 | 43.5 | 67.8 | 51.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 |

## Appendix H Relaxation of Parameter Sharing with Prefix Tuning

As discussed in §2.3, we employed LoRA modules to relax parameter sharing. However, efficiently computing multiple LoRA modules requires specialized kernels and sequential computations of the base layers and LoRA modules, incurring additional computational costs.
Consequently, we explored replacing layer-specific prompts (Liu et al., ) as an alternative. In this approach, prompts specific to each layer are integrated as prefix tokens, being utilized to generate layer-wise key and value states for self-attention computation. This approach is significantly more amenable to parallel computation, leading to reduced computational overhead.

Table H.1 summarizes performance of the prefix tuning method. Due to the reliance on small learnable prompts, performance gains were limited. Additionally, without a mechanism to leverage the original pretrained weights, the performance of prefix tuning was significantly lower (52.1% vs. 47.6% with the Average method in 1.07B model size). While offering parallel computation advantages, further research is needed to enhance its effectiveness.

**Table H.1: Evaluation results of relaxation through prefix tuning methods. Prefix length denotes the sequence length of trainable vectors used to generate key-value states in each self-attention layer. Non-embedding parameter sizes include the sizes of these trainable prefixes. Delta ($\Delta$) represent the accuracy differences to non-relaxed models using the same looping initialization.**
|  |  | Uptrain | Looping | Prefix | Perplexity $\downarrow$ | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | N-emb | PT | $N_{tok}$ | Block | Init | Len | Size | SlimP | RedP | PG19 | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg | $\Delta$ |
|  | 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.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.99B | ✓ | 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 | - |
|  | 1.00B | ✓ | 15B | 2 | Step | 256 | 9.4M | 12.62 | 10.06 | 15.80 | 53.5 | 58.3 | 73.9 | 57.6 | 57.5 | 29.3 | 35.6 | 52.2 | +0.5 |
|  | 1.01B | ✓ | 15B | 2 | Step | 512 | 18.9M | 12.67 | 10.10 | 15.85 | 54.1 | 57.8 | 73.8 | 58.4 | 57.2 | 28.7 | 35.8 | 52.3 | +0.6 |
|  | 1.03B | ✓ | 15B | 2 | Step | 1024 | 37.7M | 12.89 | 10.34 | 16.22 | 53.5 | 57.1 | 72.4 | 57.2 | 56.9 | 28.6 | 36.8 | 51.8 | +0.1 |
|  | 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.99B | ✓ | 15B | 2 | Avg | - | - | 15.15 | 12.57 | 19.86 | 43.6 | 47.4 | 70.4 | 52.6 | 50.5 | 27.8 | 34.4 | 46.7 | - |
|  | 1.00B | ✓ | 15B | 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.03B | ✓ | 15B | 2 | Avg | 1024 | 37.7M | 14.60 | 12.02 | 18.89 | 46.9 | 49.7 | 71.1 | 52.3 | 51.0 | 28.6 | 34.2 | 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 | 47.6 | +0.9 |
|  | 0.99B | ✓ | 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.00B | ✓ | 15B | 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.03B | ✓ | 15B | 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 | 9.37 | 11.94 | 43.3 | 42.2 | 66.8 | 53.4 | 44.7 | 23.2 | 29.2 | 43.3 | - |
|  | 0.48B | ✗ | 15B | - | - | - | - | 16.61 | 15.66 | 20.27 | 22.3 | 30.0 | 60.9 | 50.6 | 37.0 | 23.0 | 28.0 | 36.0 | - |
|  | 0.48B | ✓ | 15B | 2 | Step | - | - | 11.61 | 9.89 | 13.00 | 39.6 | 39.8 | 66.5 | 52.9 | 44.3 | 24.9 | 30.6 | 42.7 | - |
|  | 0.49B | ✓ | 15B | 2 | Step | 256 | 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 | ✓ | 15B | 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.53B | ✓ | 15B | 2 | Step | 1024 | 46.1M | 11.60 | 9.88 | 13.00 | 39.7 | 39.9 | 66.7 | 53.0 | 44.3 | 25.1 | 30.6 | 42.8 | +0.1 |
|  | 0.57B | ✓ | 15B | 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 |
| TinyLlama | 0.48B | ✓ | 15B | 2 | Avg | - | - | 11.86 | 10.29 | 13.42 | 38.6 | 39.4 | 66.1 | 52.8 | 42.7 | 25.4 | 30.6 | 42.2 | - |
|  | 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 | Avg | 512 | 23.1M | 11.86 | 10.28 | 13.41 | 38.1 | 39.3 | 66.3 | 52.2 | 42.6 | 25.6 | 30.8 | 42.1 | – 0.1 |
|  | 0.53B | ✓ | 15B | 2 | Avg | 1024 | 46.1M | 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 | 11.86 | 10.28 | 13.42 | 38.5 | 39.5 | 65.9 | 52.7 | 42.4 | 25.7 | 31.0 | 42.2 | +0.0 |
|  | 0.48B | ✓ | 15B | 2 | Lower | - | - | 14.67 | 12.67 | 16.68 | 31.9 | 32.3 | 62.6 | 52.0 | 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.3 | 27.8 | 38.2 | – 0.1 |
|  | 0.50B | ✓ | 15B | 2 | Lower | 512 | 23.1M | 14.67 | 12.67 | 16.69 | 31.9 | 32.3 | 62.8 | 51.7 | 38.9 | 22.2 | 27.8 | 38.2 | – 0.1 |
|  | 0.53B | ✓ | 15B | 2 | Lower | 1024 | 46.1M | 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.57B | ✓ | 15B | 2 | Lower | 2048 | 92.3M | 14.67 | 12.67 | 16.67 | 34.1 | 32.5 | 62.8 | 52.4 | 38.5 | 23.0 | 27.6 | 38.7 | +0.4 |

## Appendix I Expanded Results of Extended Uptraining and Knowledge 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., ; Kim and Rush, ). Specifically, we expanded the token number from 15 billion to 60 billion. Furthermore, we designated the teacher model as the full-size model for each architecture, uptrained on 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 and Rush ()), reverse KL (RKL; Gu et al. ()), Jensen–Shannon divergence (JSD; Agarwal et al. ()), and total variation distance (TVD; Wen et al. ()).
Table I.1 summarizes the controlled experimental results for each method. We finally selected forward KL as the distillation loss function due to its simplicity and superior performance. Especially, we observed a performance improvement of 1.7%p attributed to the extended uptraining and up to 1.7%p from the KD loss. This suggests that combining both techniques could yield even greater gains.

**Table I.1: 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 as the final configuration due to its simplicity and superior performance.**
|  | Uptrain | Looping | LoRA | Perplexity $\downarrow$ | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| N-emb | PT | $N_{tok}$ | KD | Func | Block | Init | Rank | Init | SlimP | RedP | PG19 | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg | $\Delta$ |
| 0.99B | ✓ | 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 | ✓ | 60B | ✗ | - | 2 | 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 | ✓ | 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 | ✓ | 15B | ✓ | FKL | 2 | Step | - | - | 12.36 | 9.85 | 15.45 | 56.8 | 58.6 | 74.8 | 58.6 | 59.1 | 29.2 | 36.6 | 53.4 | +1.7 |
| 0.99B | ✓ | 15B | ✓ | RKL | 2 | Step | - | - | 12.56 | 10.09 | 15.80 | 55.6 | 58.3 | 74.3 | 58.6 | 58.3 | 30.4 | 37.4 | 53.3 | +1.6 |
| 0.99B | ✓ | 15B | ✓ | JSD | 2 | Step | - | - | 12.60 | 10.06 | 15.77 | 56.1 | 58.4 | 73.4 | 57.0 | 58.4 | 29.8 | 37.2 | 52.9 | +1.2 |
| 0.99B | ✓ | 15B | ✓ | TVD | 2 | 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 | ✗ | - | 2 | 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.30B | ✓ | 15B | ✓ | FKL | 2 | 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 |
| 1.30B | ✓ | 15B | ✓ | RKL | 2 | Avg | 256 | SVD | 11.95 | 9.62 | 14.79 | 60.0 | 61.6 | 74.5 | 60.0 | 58.1 | 32.9 | 37.8 | 55.0 | +0.8 |
| 1.30B | ✓ | 15B | ✓ | JSD | 2 | 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 |
| 1.30B | ✓ | 15B | ✓ | TVD | 2 | Avg | 256 | SVD | 12.05 | 9.62 | 14.78 | 59.3 | 61.5 | 73.9 | 60.5 | 59.0 | 33.0 | 38.2 | 55.1 | +0.9 |

Figure: (a) Gemma
Refer to caption: x46.png

#### Overall performance after longer training with distillation loss

Figure I.1 and Table I.2 summarize the performance gains achieved by incorporating these advanced training techniques.
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.

**Table I.2: 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 ($\Delta$) represents the accuracy differences between the longer uptrained models with KD and their 15 billion uptrained counterparts.**
|  |  | Uptrain | Looping | LoRA | Perplexity $\downarrow$ | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | N-emb | PT | $N_{tok}$ | KD | Block | Init | Rank | Init | SlimP | RedP | PG19 | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg | $\Delta$ |
|  | 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.99B | ✗ | 60B | ✓ | - | - | - | - | 15.33 | 13.04 | 20.37 | 42.3 | 43.0 | 68.8 | 53.4 | 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.66B | ✗ | 60B | ✓ | - | - | - | - | 16.79 | 14.39 | 22.85 | 37.5 | 38.4 | 68.7 | 50.4 | 46.5 | 24.6 | 31.6 | 42.5 | - |
|  | 0.66B | ✗ | 15B | ✗ | - | - | - | - | 24.44 | 21.69 | 36.03 | 27.2 | 30.6 | 63.8 | 50.5 | 40.6 | 22.0 | 27.0 | 37.4 | - |
|  | 0.99B | ✓ | 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.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 | ✗ | 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 | - |
|  | 1.15B | ✓ | 15B | ✗ | 2 | 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 | - |
|  | 1.30B | ✓ | 15B | ✗ | 2 | 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 | - |
| Gemma | 1.60B | ✓ | 15B | ✗ | 2 | 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 | - |
|  | 0.99B | ✓ | 60B | ✓ | 2 | Step | - | - | 11.44 | 9.14 | 13.98 | 56.5 | 62.1 | 75.2 | 59.4 | 59.8 | 32.5 | 38.6 | 54.9 | +3.2 |
|  | 1.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 | 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 | +4.1 |
|  | 1.30B | ✓ | 60B | ✓ | 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 | +2.8 |
|  | 1.60B | ✓ | 60B | ✓ | 2 | 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 |
|  | 0.66B | ✓ | 60B | ✓ | 3 | Step | - | - | 12.27 | 9.90 | 15.24 | 55.6 | 58.1 | 73.1 | 60.2 | 58.8 | 30.2 | 37.2 | 53.3 | +5.5 |
|  | 0.74B | ✓ | 60B | ✓ | 3 | Avg | 64 | SVD | 12.13 | 9.80 | 14.95 | 55.5 | 58.3 | 73.5 | 60.1 | 58.0 | 31.1 | 36.8 | 53.3 | - |
|  | 0.82B | ✓ | 60B | ✓ | 3 | Avg | 128 | SVD | 11.83 | 9.53 | 14.51 | 56.7 | 60.2 | 74.2 | 59.8 | 59.1 | 33.0 | 35.4 | 54.1 | - |
|  | 0.97B | ✓ | 60B | ✓ | 3 | Avg | 256 | SVD | 11.43 | 9.17 | 13.87 | 59.3 | 62.6 | 74.7 | 61.2 | 61.6 | 32.9 | 40.2 | 56.1 | - |
|  | 1.27B | ✓ | 60B | ✓ | 3 | 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.97B | ✓ | - | - | - | - | - | - | 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.48B | ✗ | 15B | ✗ | - | - | - | - | 16.61 | 15.66 | 20.27 | 22.3 | 30.0 | 60.9 | 50.6 | 37.0 | 23.0 | 28.0 | 36.0 | - |
|  | 0.48B | ✓ | 15B | ✗ | 2 | Step | - | - | 11.61 | 9.89 | 13.00 | 39.6 | 39.8 | 66.5 | 52.9 | 44.3 | 24.9 | 30.6 | 42.7 | - |
|  | 0.53B | ✓ | 15B | ✗ | 2 | 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 | - |
|  | 0.58B | ✓ | 15B | ✗ | 2 | 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 | - |
| TinyLlama | 0.68B | ✓ | 15B | ✗ | 2 | Avg | 256 | SVD | 10.71 | 9.18 | 11.82 | 44.1 | 43.2 | 68.1 | 53.5 | 44.4 | 25.7 | 30.4 | 44.2 | - |
|  | 0.86B | ✓ | 15B | ✗ | 2 | Avg | 512 | SVD | 10.46 | 8.92 | 11.50 | 46.0 | 44.1 | 68.2 | 53.0 | 45.8 | 25.1 | 31.2 | 44.8 | - |
|  | 0.48B | ✓ | 60B | ✓ | 2 | Step | - | - | 10.51 | 9.01 | 11.60 | 44.2 | 43.1 | 68.2 | 52.4 | 44.7 | 25.3 | 32.2 | 44.3 | +1.6 |
|  | 0.53B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 10.14 | 8.77 | 11.19 | 44.3 | 44.9 | 69.5 | 52.5 | 46.5 | 26.1 | 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 |
|  | 0.68B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 9.96 | 8.56 | 10.93 | 46.2 | 45.7 | 69.0 | 53.2 | 47.9 | 25.9 | 31.6 | 45.6 | +1.4 |
|  | 0.86B | ✓ | 60B | ✓ | 2 | 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 | 51.9 | 27.7 | 31.6 | 48.6 | - |
|  | 0.40B | ✗ | 60B | ✓ | - | - | - | - | 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 | ✗ | - | - | - | - | 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 | ✗ | 2 | Step | - | - | 16.38 | 12.37 | 17.74 | 43.4 | 40.5 | 67.4 | 50.8 | 46.3 | 25.7 | 30.0 | 43.5 | - |
|  | 0.44B | ✓ | 15B | ✗ | 2 | 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.48B | ✓ | 15B | ✗ | 2 | 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 | - |
| Pythia | 0.55B | ✓ | 15B | ✗ | 2 | Avg | 256 | SVD | 15.22 | 11.54 | 16.47 | 48.5 | 43.3 | 67.2 | 51.4 | 46.7 | 25.5 | 32.0 | 44.9 | - |
|  | 0.70B | ✓ | 15B | ✗ | 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.40B | ✓ | 60B | ✓ | 2 | Step | - | - | 14.59 | 11.13 | 15.79 | 47.8 | 43.8 | 69.3 | 52.0 | 48.1 | 25.4 | 30.4 | 45.2 | +1.7 |
|  | 0.44B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 14.24 | 10.89 | 15.52 | 50.0 | 44.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.55B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 13.91 | 10.61 | 14.91 | 50.5 | 45.6 | 68.7 | 51.2 | 48.4 | 25.7 | 32.8 | 46.1 | +1.2 |
|  | 0.70B | ✓ | 60B | ✓ | 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 |

## Appendix J Expanded Results of Early-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 on various strategies, demonstrating Recursive Transformers can be transformed into early-exiting models without compromising final loop output’s performance.
Table J.1 presents a comprehensive summary of experimental results 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 (co-training) significantly degraded the final output performance, even with an aggressive loss coefficient strategy.
- •
While freezing all parameters, we attempted to train intermediate loop outputs by attaching trainable LoRA modules to frozen classifier head. However, we found that this proved ineffective.
- •
The aggressive coefficient strategy successfully maintained final loop output performance while enhancing intermediate output performance. Moreover, incorporating knowledge distillation from detached final outputs further enhanced intermediate performance.
- •
No significant performance differences were observed when using the overlapped uptraining dataset versus new SlimPajama tokens for post-training.

Finally, we opted to perform post-training with new tokens sourced from the same SlimPajama dataset. Moreover, we incorporated a distillation loss from the final loop output, while aggressively reducing the loss coefficients of intermediate outputs.

#### Early-exit training results on final models

We applied the aggressive coefficient strategy with distillation loss to our final models (uptrained on 60 billion tokens with knowledge distillation). Tables J.2 and J.3 summarize the performance of intermediate loops and the final loop across 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.
Consistent with previous findings, the aggressive coefficient strategy yielded the best performance across both intermediate and final outputs.

However, we find that intermediate loop outputs in LoRA-relaxed models underperformed their non-relaxed counterparts (recursive models). This could potentially reduce throughput gain, as early loop performance directly influences the number of tokens eligible for early-exit.
In perfectly tied looping blocks, intermediate outputs seem to be distilled from the last, as all gradients are backpropagated to the same parameters.
Conversely, since LoRA modules allow each layer to specialize based on its depth, intermediate representations appear to be optimized to facilitate performance of the final output, not for their own sake.
Hence, relaxation introduces a trade-off between final performance and early-exiting benefits.
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 is left for future work.

**Table J.1: Ablation studies on early-exit training for recursive Gemma models. We evaluated performance in a static-exiting scenario (Schuster et al., ; Bae et al., ), where all tokens exit at either first or second iteration loops (9th or 18th depths). We explored post-training (after uptraining) and co-training (during uptraining) approaches. Moreover, we explored freezing uptrained weights and adding LoRA with the rank of 128 to the classifier head. Different coefficient values were tested for the aggressive CE loss function. Early-exit training utilized 15 billion tokens, either overlapping with uptraining data or entirely new. Delta ($\Delta$) indicates the performance changes of the final loop outputs. We highlight the final configuration: post-training with aggressive CE and KD loss on 15 billion new tokens.**
|  | Uptrain | Looping | Early-Exit Train | Perplexity $\downarrow$ | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| N-emb | PT | $N_{tok}$ | Block | Init | Train | Freeze | $N_{tok}$ | CE | KD | Data | SlimP | RedP | PG19 | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg | $\Delta$ |
| 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.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.99B | ✓ | 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 | ✓ | 15B | 2 | Step | Post- | ✗ | 15B | Weighted | ✗ | Ovlp | 12.97 | 10.51 | 16.55 | 48.9 | 55.5 | 72.7 | 55.3 | 54.9 | 30.1 | 36.0 | 50.5 | – 1.2 |
| 13.11 | 10.59 | 16.71 | 49.5 | 54.8 | 72.0 | 53.4 | 54.1 | 29.1 | 35.6 | 49.8 | - |  |  |  |  |  |  |  |  |  |  |  |
| 0.99B | ✓ | 15B | 2 | Step | Post- | ✗ | 15B | Agg (0.3) | ✗ | Ovlp | 12.60 | 10.21 | 15.75 | 51.8 | 58.2 | 73.7 | 56.8 | 57.0 | 29.9 | 37.8 | 52.2 | +0.5 |
| 13.63 | 11.02 | 17.55 | 47.5 | 53.0 | 71.2 | 54.9 | 50.2 | 28.2 | 34.8 | 48.5 | - |  |  |  |  |  |  |  |  |  |  |  |
|  |  |  |  |  |  |  |  |  |  |  | 12.37 | 9.94 | 15.37 | 53.0 | 59.1 | 73.9 | 55.4 | 57.4 | 30.6 | 37.8 | 52.5 | +0.8 |
| 0.99B | ✓ | 15B | 2 | Step | Post- | ✗ | 15B | Agg (0.1) | ✗ | Ovlp | 14.55 | 11.87 | 19.00 | 45.9 | 51.2 | 71.4 | 54.5 | 48.1 | 26.8 | 32.0 | 47.1 | - |
| 0.99B | ✓ | 15B | 2 | Step | Post- | ✗ | 15B | Agg (0.05) | ✗ | Ovlp | 12.33 | 9.90 | 15.31 | 52.8 | 59.2 | 73.6 | 57.5 | 57.7 | 30.5 | 37.2 | 52.6 | +0.9 |
| 15.70 | 12.93 | 20.69 | 43.1 | 49.8 | 69.8 | 55.2 | 46.0 | 26.9 | 31.2 | 46.0 | - |  |  |  |  |  |  |  |  |  |  |  |
|  |  |  |  |  |  |  |  |  |  |  | 12.28 | 9.80 | 15.23 | 52.9 | 59.5 | 73.3 | 56.5 | 57.2 | 30.1 | 37.2 | 52.4 | +0.7 |
| 0.99B | ✓ | 15B | 2 | Step | Post- | ✗ | 15B | Agg (0.01) | ✗ | Ovlp | 22.76 | 20.37 | 30.39 | 32.2 | 45.2 | 67.5 | 53.9 | 40.3 | 26.3 | 29.2 | 42.1 | - |
| 0.99B | ✓ | 15B | 2 | Step | Post- | ✗ | 15B | Weighted | ✓ | Ovlp | 13.04 | 10.57 | 16.66 | 47.7 | 55.1 | 73.2 | 55.6 | 54.5 | 29.1 | 37.2 | 50.4 | – 1.3 |
| 13.04 | 10.54 | 16.66 | 48.3 | 54.9 | 72.1 | 55.9 | 54.3 | 28.4 | 35.4 | 49.9 | - |  |  |  |  |  |  |  |  |  |  |  |
|  |  |  |  |  |  |  |  |  |  |  | 12.40 | 9.97 | 15.42 | 52.9 | 58.9 | 73.7 | 55.7 | 57.5 | 31.1 | 38.2 | 52.6 | +0.9 |
| 0.99B | ✓ | 15B | 2 | Step | Post- | ✗ | 15B | Agg (0.1) | ✓ | Ovlp | 14.11 | 11.47 | 18.32 | 46.3 | 52.1 | 71.6 | 55.3 | 49.2 | 28.5 | 32.6 | 48.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.6 | 51.7 | +0.0 |
| 43.74 | 41.63 | 56.78 | 5.3 | 37.9 | 61.4 | 52.6 | 35.3 | 24.0 | 29.0 | 35.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.6 | 51.7 | +0.0 |
| 43.09 | 39.97 | 55.37 | 5.6 | 37.7 | 62.5 | 52.7 | 34.5 | 24.7 | 29.2 | 35.3 | - |  |  |  |  |  |  |  |  |  |  |  |
| 0.99B | ✓ | 15B | 2 | Step | Co- | ✗ | 15B | Agg (0.1) | ✗ | Ovlp | 13.24 | 10.67 | 16.98 | 50.1 | 54.2 | 72.2 | 53.7 | 54.7 | 28.9 | 37.4 | 50.2 | – 1.5 |
| 13.59 | 10.89 | 17.42 | 50.6 | 52.7 | 71.2 | 54.4 | 53.0 | 27.5 | 35.0 | 49.2 | - |  |  |  |  |  |  |  |  |  |  |  |
|  |  |  |  |  |  |  |  |  |  |  | 12.34 | 9.92 | 15.31 | 52.3 | 59.0 | 73.8 | 57.6 | 55.5 | 30.4 | 37.2 | 52.3 | +0.6 |
| 0.99B | ✓ | 15B | 2 | Step | Post- | ✗ | 15B | Agg (0.1) | ✗ | New | 14.49 | 11.86 | 18.89 | 43.9 | 51.3 | 71.0 | 54.9 | 48.1 | 27.5 | 31.4 | 46.9 | - |

**Table J.2: Evaluation results of Gemma models after early-exit training. Delta ($\Delta$) represent the accuracy changes in original last loop outputs 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. The relaxed model with three blocks shows a more significant performance drop because KD loss (from final loop output) could not be utilized due to out-of-memory issues.**
|  | Uptrain | Looping | LoRA | Early-exit Train | Perplexity $\downarrow$ | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| N-emb | PT | $N_{tok}$ | KD | Block | Init | Rank | Init | $N_{tok}$ | CE | KD | SlimP | RedP | PG19 | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg | $\Delta$ |
| 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 | - |
| 1.99B | ✓ | 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.99B | ✓ | 60B | ✓ | 2 | Step | - | - | - | - | - | 11.44 | 9.14 | 13.98 | 56.5 | 62.1 | 75.2 | 59.4 | 59.8 | 32.5 | 38.6 | 54.9 | - |
| 1.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 | - |
| 1.15B | ✓ | 60B | ✓ | 2 | 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 | ✓ | 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 | - |
| 1.60B | ✓ | 60B | ✓ | 2 | 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 | - |
|  |  |  |  |  |  |  |  |  |  |  | 11.71 | 9.56 | 14.46 | 54.0 | 61.7 | 75.1 | 58.9 | 58.6 | 31.9 | 37.6 | 54.0 | – 0.9 |
| 0.99B | ✓ | 60B | ✓ | 2 | Step | - | - | 15B | Agg (0.1) | ✓ | 13.68 | 11.39 | 17.60 | 45.0 | 54.1 | 71.9 | 58.5 | 49.8 | 28.8 | 33.4 | 48.8 | - |
|  |  |  |  |  |  |  |  |  |  |  | 11.79 | 9.70 | 14.52 | 53.7 | 60.8 | 73.6 | 61.1 | 58.7 | 32.9 | 37.2 | 54.0 | – 1.8 |
| 1.07B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.1) | ✓ | 19.45 | 16.46 | 26.10 | 30.7 | 37.9 | 66.5 | 55.3 | 42.2 | 25.3 | 27.6 | 40.8 | - |
|  |  |  |  |  |  |  |  |  |  |  | 11.66 | 9.59 | 14.32 | 53.3 | 62.1 | 74.9 | 60.0 | 59.9 | 33.4 | 38.8 | 54.6 | – 2.0 |
| 1.15B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.1) | ✓ | 19.65 | 16.77 | 26.44 | 29.7 | 37.7 | 66.8 | 52.6 | 41.4 | 25.3 | 28.0 | 40.2 | - |
|  |  |  |  |  |  |  |  |  |  |  | 11.47 | 9.39 | 14.03 | 54.9 | 63.0 | 74.5 | 61.7 | 60.5 | 33.1 | 38.4 | 55.2 | – 1.8 |
| 1.30B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.1) | ✓ | 19.67 | 16.82 | 26.40 | 29.7 | 38.3 | 66.4 | 53.1 | 43.8 | 24.7 | 27.6 | 40.5 | - |
|  |  |  |  |  |  |  |  |  |  |  | 11.20 | 9.14 | 13.58 | 57.2 | 64.1 | 75.2 | 61.7 | 62.1 | 34.6 | 38.2 | 56.2 | – 2.2 |
| 1.60B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.1) | ✓ | 19.29 | 16.47 | 25.73 | 32.0 | 39.6 | 67.6 | 53.3 | 43.2 | 25.8 | 30.2 | 41.7 | - |
|  |  |  |  |  |  |  |  |  |  |  | 12.11 | 9.98 | 14.97 | 52.6 | 59.8 | 74.4 | 59.4 | 57.6 | 31.1 | 37.0 | 53.1 | – 2.7 |
| 1.07B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.3) | ✓ | 16.09 | 13.54 | 21.19 | 35.4 | 42.8 | 69.8 | 52.8 | 45.8 | 25.8 | 31.0 | 43.3 | - |
|  |  |  |  |  |  |  |  |  |  |  | 11.96 | 9.87 | 14.76 | 52.3 | 60.5 | 74.2 | 59.1 | 58.9 | 33.0 | 37.2 | 53.6 | – 3.0 |
| 1.15B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.3) | ✓ | 16.28 | 13.77 | 21.45 | 35.2 | 42.1 | 69.8 | 53.5 | 46.5 | 25.8 | 31.2 | 43.4 | - |
|  |  |  |  |  |  |  |  |  |  |  | 11.73 | 9.63 | 14.43 | 54.3 | 61.4 | 75.0 | 60.7 | 58.8 | 33.1 | 38.6 | 54.6 | – 2.4 |
| 1.30B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.3) | ✓ | 16.41 | 13.89 | 21.68 | 35.6 | 42.3 | 69.0 | 52.7 | 46.8 | 26.4 | 29.8 | 43.2 | - |
|  |  |  |  |  |  |  |  |  |  |  | 11.47 | 9.36 | 13.93 | 56.2 | 62.7 | 75.4 | 60.9 | 60.4 | 34.0 | 37.0 | 55.2 | – 3.2 |
| 1.60B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.3) | ✓ | 16.24 | 13.72 | 21.42 | 37.8 | 43.6 | 69.0 | 54.4 | 45.5 | 26.4 | 31.2 | 44.0 | - |
| 0.66B | ✓ | 60B | ✓ | 3 | Step | - | - | - | - | - | 12.27 | 9.90 | 15.24 | 55.6 | 58.1 | 73.1 | 60.2 | 58.8 | 30.2 | 37.2 | 53.3 | - |
| 0.74B | ✓ | 60B | ✓ | 3 | Avg | 64 | SVD | - | - | - | 12.13 | 9.80 | 14.95 | 55.5 | 58.3 | 73.5 | 60.1 | 58.0 | 31.1 | 36.8 | 53.3 | - |
| 0.82B | ✓ | 60B | ✓ | 3 | Avg | 128 | SVD | - | - | - | 11.83 | 9.53 | 14.51 | 56.7 | 60.2 | 74.2 | 59.8 | 59.1 | 33.0 | 35.4 | 54.1 | - |
| 0.97B | ✓ | 60B | ✓ | 3 | Avg | 256 | SVD | - | - | - | 11.43 | 9.17 | 13.87 | 59.3 | 62.6 | 74.7 | 61.2 | 61.6 | 32.9 | 40.2 | 56.1 | - |
| 1.27B | ✓ | 60B | ✓ | 3 | 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 | - |
|  |  |  |  |  |  |  |  |  |  |  | 12.75 | 10.48 | 16.01 | 50.2 | 57.0 | 72.7 | 58.6 | 56.7 | 30.0 | 38.2 | 51.9 | – 1.4 |
|  |  |  |  |  |  |  |  |  |  |  | 13.81 | 11.47 | 17.80 | 48.4 | 53.0 | 72.4 | 55.6 | 51.6 | 27.2 | 35.2 | 49.0 | - |
| 0.66B | ✓ | 60B | ✓ | 3 | Step | - | - | 15B | Agg (0.1) | ✓ | 16.72 | 14.23 | 22.97 | 37.7 | 44.2 | 69.8 | 53.6 | 44.2 | 24.6 | 30.2 | 43.5 | - |
|  |  |  |  |  |  |  |  |  |  |  | 12.64 | 10.43 | 15.81 | 51.4 | 56.3 | 72.2 | 57.9 | 56.7 | 30.4 | 35.0 | 51.4 | – 1.9 |
|  |  |  |  |  |  |  |  |  |  |  | 19.90 | 16.88 | 26.26 | 30.4 | 39.3 | 66.3 | 54.1 | 41.2 | 24.8 | 29.2 | 40.8 | - |
| 0.74B | ✓ | 60B | ✓ | 3 | Avg | 64 | SVD | 15B | Agg (0.1) | ✗ | 26.31 | 22.49 | 36.10 | 20.9 | 31.2 | 62.6 | 50.8 | 37.2 | 22.0 | 28.0 | 36.1 | - |
|  |  |  |  |  |  |  |  |  |  |  | 12.37 | 10.21 | 15.38 | 52.0 | 58.0 | 72.0 | 56.5 | 58.4 | 30.0 | 35.2 | 51.7 | – 2.4 |
|  |  |  |  |  |  |  |  |  |  |  | 20.07 | 17.09 | 26.47 | 30.9 | 40.5 | 66.3 | 55.4 | 40.8 | 24.4 | 29.6 | 41.1 | - |
| 0.82B | ✓ | 60B | ✓ | 3 | Avg | 128 | SVD | 15B | Agg (0.1) | ✗ | 26.15 | 22.46 | 35.98 | 21.3 | 31.2 | 62.7 | 51.8 | 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 |
|  |  |  |  |  |  |  |  |  |  |  | 19.29 | 16.49 | 25.51 | 35.2 | 42.5 | 65.8 | 55.6 | 41.5 | 25.6 | 29.4 | 42.2 | - |
| 0.97B | ✓ | 60B | ✓ | 3 | Avg | 256 | SVD | 15B | Agg (0.1) | ✗ | 25.12 | 21.53 | 34.53 | 23.1 | 32.0 | 63.2 | 49.7 | 36.1 | 23.0 | 25.2 | 36.1 | - |
|  |  |  |  |  |  |  |  |  |  |  | 11.49 | 9.38 | 14.00 | 56.1 | 62.7 | 74.4 | 60.5 | 62.1 | 34.9 | 38.8 | 55.7 | – 3.0 |
|  |  |  |  |  |  |  |  |  |  |  | 18.52 | 15.79 | 24.34 | 36.7 | 44.9 | 67.2 | 55.3 | 43.8 | 26.0 | 30.4 | 43.5 | - |
| 1.27B | ✓ | 60B | ✓ | 3 | Avg | 512 | SVD | 15B | Agg (0.1) | ✗ | 24.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 | 29.4 | 35.2 | 50.1 | – 3.2 |
|  |  |  |  |  |  |  |  |  |  |  | 16.68 | 14.08 | 21.86 | 35.4 | 42.4 | 68.2 | 53.8 | 44.6 | 26.3 | 29.4 | 42.9 | - |
| 0.74B | ✓ | 60B | ✓ | 3 | Avg | 64 | SVD | 15B | Agg (0.3) | ✗ | 21.43 | 18.26 | 29.12 | 24.4 | 34.1 | 64.3 | 50.5 | 40.7 | 22.3 | 27.8 | 37.7 | - |
|  |  |  |  |  |  |  |  |  |  |  | 12.71 | 10.54 | 15.92 | 50.4 | 55.9 | 73.1 | 57.5 | 56.8 | 30.1 | 34.8 | 51.2 | – 2.9 |
|  |  |  |  |  |  |  |  |  |  |  | 16.90 | 14.32 | 22.18 | 37.6 | 43.5 | 67.6 | 54.5 | 45.0 | 25.3 | 29.0 | 43.2 | - |
| 0.82B | ✓ | 60B | ✓ | 3 | Avg | 128 | SVD | 15B | Agg (0.3) | ✗ | 21.23 | 18.13 | 28.88 | 25.3 | 34.0 | 64.6 | 51.7 | 40.7 | 23.0 | 26.4 | 38.0 | - |
|  |  |  |  |  |  |  |  |  |  |  | 12.26 | 10.15 | 15.23 | 53.5 | 58.5 | 73.5 | 58.8 | 58.3 | 30.6 | 37.6 | 53.0 | – 3.1 |
|  |  |  |  |  |  |  |  |  |  |  | 16.56 | 14.09 | 21.68 | 42.6 | 45.1 | 68.2 | 57.7 | 45.7 | 25.9 | 28.8 | 44.8 | - |
| 0.97B | ✓ | 60B | ✓ | 3 | Avg | 256 | SVD | 15B | Agg (0.3) | ✗ | 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 |
|  |  |  |  |  |  |  |  |  |  |  | 16.02 | 13.53 | 20.86 | 43.5 | 47.5 | 68.3 | 56.2 | 47.1 | 27.0 | 30.4 | 45.7 | - |
| 1.27B | ✓ | 60B | ✓ | 3 | Avg | 512 | SVD | 15B | Agg (0.3) | ✗ | 20.20 | 17.21 | 27.50 | 28.9 | 35.2 | 65.6 | 52.9 | 41.9 | 23.2 | 26.6 | 39.2 | - |

**Table J.3: Evaluation results of TinyLlama and Pythia models after early-exit training. Delta ($\Delta$) represents the accuracy change in the original last loop outputs after early-exit post-training.**
|  |  | Uptrain | Looping | LoRA | Early-exit Train | Perplexity $\downarrow$ | Few-shot Accuracy $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | N-emb | PT | $N_{tok}$ | KD | Block | Init | Rank | Init | $N_{tok}$ | CE | KD | SlimP | RedP | PG19 | LD | HS | PQ | WG | ARC-e | ARC-c | OB | Avg | $\Delta$ |
|  | 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.48B | ✓ | 60B | ✓ | 2 | Step | - | - | - | - | - | 10.51 | 9.01 | 11.60 | 44.2 | 43.1 | 68.2 | 52.4 | 44.7 | 25.3 | 32.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.1 | 31.6 | 45.0 | - |
|  | 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 | - |
|  | 0.68B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | - | - | - | 9.96 | 8.56 | 10.93 | 46.2 | 45.7 | 69.0 | 53.2 | 47.9 | 25.9 | 31.6 | 45.6 | - |
|  | 0.86B | ✓ | 60B | ✓ | 2 | 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 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 10.55 | 9.16 | 11.68 | 45.0 | 43.7 | 68.9 | 53.4 | 44.8 | 25.3 | 32.2 | 44.8 | + 0.5 |
|  | 0.48B | ✓ | 60B | ✓ | 2 | Step | - | - | 15B | Agg (0.1) | ✓ | 12.28 | 10.62 | 13.83 | 38.2 | 39.4 | 65.8 | 52.3 | 41.5 | 24.7 | 30.6 | 41.8 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 10.34 | 9.08 | 11.50 | 43.4 | 44.8 | 69.5 | 53.4 | 46.9 | 25.6 | 32.0 | 45.1 | + 0.1 |
|  | 0.53B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.1) | ✓ | 21.23 | 18.63 | 24.85 | 16.8 | 29.0 | 57.6 | 48.9 | 33.2 | 23.1 | 27.0 | 33.7 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 10.25 | 8.97 | 11.36 | 45.2 | 45.5 | 68.8 | 54.0 | 46.5 | 25.0 | 31.6 | 45.2 | + 0.2 |
|  | 0.58B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.1) | ✓ | 21.30 | 18.56 | 24.75 | 18.5 | 28.9 | 58.4 | 48.0 | 34.1 | 21.8 | 27.4 | 33.9 | - |
| TinyLlama |  |  |  |  |  |  |  |  |  |  |  | 10.13 | 8.84 | 11.23 | 45.2 | 45.9 | 69.6 | 53.6 | 46.9 | 25.9 | 32.0 | 45.6 | + 0.0 |
|  | 0.68B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.1) | ✓ | 20.95 | 18.16 | 24.22 | 20.1 | 28.8 | 57.8 | 48.9 | 33.8 | 22.5 | 25.0 | 33.9 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 10.02 | 8.74 | 11.04 | 46.6 | 46.5 | 68.6 | 54.5 | 47.9 | 26.3 | 32.2 | 46.1 | + 0.2 |
|  | 0.86B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.1) | ✓ | 20.38 | 17.70 | 23.57 | 19.9 | 28.8 | 58.2 | 49.0 | 34.7 | 22.8 | 25.8 | 34.2 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 10.61 | 9.36 | 11.87 | 42.1 | 43.7 | 68.6 | 54.1 | 46.1 | 26.0 | 31.2 | 44.6 | – 0.4 |
|  | 0.53B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.3) | ✓ | 16.83 | 14.88 | 19.77 | 22.0 | 30.3 | 60.7 | 50.7 | 36.9 | 24.1 | 27.8 | 36.1 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 10.50 | 9.22 | 11.71 | 44.2 | 44.2 | 69.2 | 53.0 | 46.0 | 25.5 | 31.2 | 44.8 | – 0.2 |
|  | 0.58B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.3) | ✓ | 17.10 | 15.03 | 19.99 | 23.5 | 30.1 | 60.8 | 51.3 | 36.5 | 23.8 | 26.4 | 36.0 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 10.34 | 9.07 | 11.51 | 44.0 | 45.0 | 68.4 | 53.0 | 45.8 | 26.0 | 31.2 | 44.8 | – 0.8 |
|  | 0.68B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.3) | ✓ | 17.06 | 14.92 | 19.82 | 24.2 | 30.4 | 59.9 | 51.7 | 36.2 | 23.9 | 27.2 | 36.2 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 10.21 | 8.94 | 11.28 | 45.1 | 45.8 | 69.3 | 54.5 | 46.7 | 25.9 | 33.4 | 45.8 | – 0.1 |
|  | 0.86B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.3) | ✓ | 16.76 | 14.68 | 19.43 | 24.4 | 30.0 | 61.1 | 51.9 | 37.1 | 22.9 | 28.2 | 36.5 | - |
|  | 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 | - |
|  | 0.81B | ✓ | 75B | ✗ | - | - | - | - | - | - | - | 12.86 | 9.86 | 13.74 | 54.8 | 50.3 | 70.5 | 55.3 | 52.2 | 28.8 | 33.0 | 49.3 | + 0.7 |
|  | 0.40B | ✓ | 60B | ✓ | 2 | 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 | ✓ | 2 | Avg | 64 | SVD | - | - | - | 14.24 | 10.89 | 15.52 | 50.0 | 44.5 | 68.9 | 54.1 | 48.0 | 26.5 | 31.2 | 46.2 | - |
|  | 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 | - |
|  | 0.55B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | - | - | - | 13.91 | 10.61 | 14.91 | 50.5 | 45.6 | 68.7 | 51.2 | 48.4 | 25.7 | 32.8 | 46.1 | - |
|  | 0.70B | ✓ | 60B | ✓ | 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 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 14.72 | 11.38 | 16.31 | 47.0 | 44.2 | 69.2 | 53.4 | 48.6 | 24.7 | 30.4 | 45.4 | + 0.2 |
|  | 0.40B | ✓ | 60B | ✓ | 2 | Step | - | - | 15B | Agg (0.1) | ✓ | 18.61 | 14.11 | 20.96 | 38.4 | 38.1 | 67.0 | 53.7 | 43.3 | 24.4 | 29.0 | 42.0 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 14.49 | 11.22 | 16.12 | 49.1 | 43.9 | 69.8 | 53.8 | 48.6 | 26.1 | 31.2 | 46.1 | – 0.1 |
|  | 0.44B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.1) | ✓ | 24.43 | 18.19 | 27.89 | 26.7 | 31.6 | 61.6 | 50.8 | 38.2 | 22.9 | 27.6 | 37.1 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 14.35 | 11.17 | 15.93 | 50.1 | 44.7 | 69.0 | 52.1 | 49.9 | 25.3 | 32.6 | 46.2 | + 0.0 |
|  | 0.48B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.1) | ✓ | 24.33 | 18.09 | 27.96 | 28.2 | 32.3 | 61.1 | 53.0 | 38.8 | 23.7 | 27.4 | 37.8 | - |
| Pythia |  |  |  |  |  |  |  |  |  |  |  | 14.14 | 10.96 | 15.54 | 50.8 | 45.5 | 68.2 | 53.9 | 48.8 | 25.3 | 32.8 | 46.5 | + 0.4 |
|  | 0.55B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.1) | ✓ | 24.18 | 17.87 | 27.48 | 28.1 | 32.3 | 61.9 | 54.1 | 38.1 | 22.9 | 28.6 | 38.0 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 13.81 | 10.72 | 15.11 | 52.4 | 47.0 | 69.3 | 52.7 | 50.1 | 26.9 | 32.0 | 47.2 | + 0.2 |
|  | 0.70B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.1) | ✓ | 23.50 | 17.49 | 26.72 | 29.5 | 32.8 | 63.2 | 52.3 | 38.8 | 22.8 | 27.8 | 38.2 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 14.87 | 11.53 | 16.61 | 47.0 | 43.1 | 68.7 | 53.0 | 47.4 | 25.7 | 31.0 | 45.1 | – 0.9 |
|  | 0.44B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.3) | ✓ | 20.62 | 15.60 | 23.57 | 32.6 | 33.6 | 63.4 | 51.2 | 40.7 | 23.3 | 28.0 | 39.0 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 14.69 | 11.46 | 16.36 | 48.9 | 43.8 | 68.4 | 53.0 | 49.1 | 26.2 | 31.6 | 45.9 | – 0.3 |
|  | 0.48B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.3) | ✓ | 20.60 | 15.56 | 23.63 | 33.2 | 33.6 | 62.7 | 51.1 | 41.3 | 23.6 | 27.8 | 39.0 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 14.44 | 11.20 | 15.94 | 50.0 | 44.7 | 69.2 | 52.3 | 48.1 | 25.4 | 32.2 | 46.0 | – 0.1 |
|  | 0.55B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.3) | ✓ | 20.61 | 15.48 | 23.45 | 33.3 | 34.2 | 63.4 | 52.2 | 40.8 | 23.0 | 28.8 | 39.4 | - |
|  |  |  |  |  |  |  |  |  |  |  |  | 14.08 | 10.94 | 15.44 | 51.1 | 46.4 | 68.7 | 52.2 | 50.0 | 26.9 | 31.6 | 46.7 | – 0.3 |
|  | 0.70B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.3) | ✓ | 20.20 | 15.25 | 22.98 | 34.6 | 34.1 | 63.5 | 53.0 | 41.5 | 23.6 | 27.6 | 39.7 | - |

## Appendix K Expanded Results of Hypothetical Generation Speedup

#### Measuring the average per-token generation time

First, we measured the generation time with various model configurations using dummy weights and inputs. We measured the elapsed time for each components, such as embedding matrices, Transformer blocks, and the classifier head.
Especially, we calculated the time per token by dividing the total 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.
We measured decoding speed using FlashDecoding (Dao et al., ), a technique that has recently become standard in serving LLMs.
Using a single A100 40GB GPU, we measured generation times by increasing batch sizes until an out-of-memory error occurred or memory usuage reached the predefined limit.

In Table K.1, generation time was measured up to the maximum batch size that a single A100 GPU could accommodate before encountering out-of-memory errors, with prefix and decoding lengths set to 512 and 2048, respectively. Meanwhile, Table K.3 presents generation times measured in a more memory-constrained 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 memory efficiency.

When comparing the speed of the three models, Gemma 2B was the fastest, followed by TinyLlama 1.1B and then Pythia 1B. This order is the exact inverse of their non-embedding parameter sizes. This speed difference is attributed to the Grouped-Query (Shazeer, ) and Multi-Query attention (Ainslie et al., ). The main decoding bottleneck in Transformers is memory access of heavy key-value caches. Hence, Gemma that effectively reduces the key-value cache size through the MQA mechanism, achieves fastest speeds. Despite using GQA, TinyLlama 1.1B has a similar speed to Gemma 2B due to its shallow and deep architecture (22 layers compared to Gemma’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 (test sets of SlimPajama, RedPajama, and PG19), assuming our models generated those tokens. We used 20K samples to obtain their exit trajectories, and we employed an oracle-exiting algorithm to determine the earliest possible exit point for each token.
Combining this trajectory data with previously measured per-token generation times across various batch sizes (considering only Transformer block computations), we estimated the hypothetical throughput for various settings and datasets.
The results are detailed in Tables K.2 and K.4.

Our analysis reveals that Recursive Transformers achieve a 2-3$\times$ throughput gain over vanilla model counterparts. Relaxed models also demonstrated significant speedup despite their unoptimized LoRA computations. Currently, we merge multiple LoRAs into a single, large LoRA module to enable parallel computation of samples at different looping iterations.
However, this introduces extra overhead due to redundant computations, resulting in 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. Since attention computation has quadratic complexity with respect to sequence length, it becomes less expensive at shorter lengths, while the complexity of LoRA computation remains constant. This highlights the need for highly optimized LoRA computations to achieve substantial throughput gains in all scenarios. Nevertheless, these findings also suggest that relaxed models will yield even greater performance and throughput improvements with longer contexts where attention computation dominates.

#### Approximation errors in our hypothetical throughput

Since our throughput estimations are 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 is expected to 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 additional 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 to prevent performance degradation (Bae et al., ). While this adds negligible overhead in memory-bound scenarios, it inevitably increases overhead in compute-bound 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 a more realistic analysis of early-exiting generation is a direction for future work.

**Table K.1: Measurements of generation time across three models using a single A100 40GB GPU. We measured time per token for both a batch size of 1 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-query attention, demonstrated fast generation speeds and large maximum batch sizes relative to their model sizes. TinyLlama’s deep and narrow architecture allowed for a significantly large maximum batch size, although its generation speed was slower due to the increased number of layers.**
|  | Model Architecture |  | Recursive |  | Time (ms) per token |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | $N_{L}$ | $d_{model}$ | $N_{head}$ | $N_{KV}$ | Vocab | N-emb | Block | Rank | Batch | Total | Emb | Transformer | Head |
|  |  |  |  |  |  |  |  |  | 1 | 22.994 | 0.087 | 21.344 | 0.803 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.98B | - | - | 43 | 0.657 | 0.002 | 0.616 | 0.023 |
|  |  |  |  |  |  |  |  |  | 1 | 13.918 | 0.088 | 11.059 | 0.827 |
|  | 18 | 2048 | 8 | 1 | 256K | 0.99B | 2 | - | 43 | 0.336 | 0.002 | 0.265 | 0.023 |
|  |  |  |  |  |  |  |  |  | 1 | 15.858 | 0.080 | 13.096 | 0.825 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.07B | 2 | 64 | 41 | 0.398 | 0.002 | 0.323 | 0.024 |
|  |  |  |  |  |  |  |  |  | 1 | 15.708 | 0.080 | 12.969 | 0.822 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.15B | 2 | 128 | 41 | 0.398 | 0.002 | 0.324 | 0.024 |
|  |  |  |  |  |  |  |  |  | 1 | 15.456 | 0.083 | 12.721 | 0.818 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.30B | 2 | 256 | 39 | 0.450 | 0.002 | 0.372 | 0.025 |
| Gemma |  |  |  |  |  |  |  |  | 1 | 15.489 | 0.078 | 12.775 | 0.817 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.60B | 2 | 512 | 39 | 0.499 | 0.002 | 0.422 | 0.025 |
|  |  |  |  |  |  |  |  |  | 1 | 10.546 | 0.081 | 7.394 | 0.827 |
|  | 18 | 2048 | 8 | 1 | 256K | 0.66B | 3 | - | 43 | 0.263 | 0.002 | 0.182 | 0.023 |
|  |  |  |  |  |  |  |  |  | 1 | 11.871 | 0.080 | 8.724 | 0.827 |
|  | 18 | 2048 | 8 | 1 | 256K | 0.74B | 3 | 64 | 43 | 0.306 | 0.002 | 0.182 | 0.023 |
|  |  |  |  |  |  |  |  |  | 1 | 11.768 | 0.080 | 8.649 | 0.825 |
|  | 18 | 2048 | 8 | 1 | 256K | 0.82B | 3 | 128 | 43 | 0.294 | 0.002 | 0.221 | 0.023 |
|  |  |  |  |  |  |  |  |  | 1 | 12.018 | 0.081 | 8.848 | 0.823 |
|  | 18 | 2048 | 8 | 1 | 256K | 0.97B | 3 | 256 | 41 | 0.311 | 0.002 | 0.226 | 0.024 |
|  |  |  |  |  |  |  |  |  | 1 | 12.087 | 0.082 | 8.932 | 0.822 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.27B | 3 | 512 | 39 | 0.325 | 0.002 | 0.237 | 0.025 |
|  |  |  |  |  |  |  |  |  | 1 | 22.016 | 0.082 | 21.010 | 0.188 |
|  | 22 | 2048 | 32 | 4 | 32K | 0.97B | - | - | 329 | 0.819 | 0.000 | 0.815 | 0.001 |
|  |  |  |  |  |  |  |  |  | 1 | 12.657 | 0.077 | 10.370 | 0.209 |
|  | 22 | 2048 | 32 | 4 | 32K | 0.48B | 2 | - | 233 | 0.446 | 0.000 | 0.413 | 0.001 |
|  |  |  |  |  |  |  |  |  | 1 | 15.243 | 0.079 | 12.908 | 0.211 |
| TinyLlama | 22 | 2048 | 32 | 4 | 32K | 0.53B | 2 | 64 | 211 | 0.454 | 0.000 | 0.421 | 0.002 |
|  |  |  |  |  |  |  |  |  | 1 | 15.456 | 0.082 | 13.118 | 0.213 |
|  | 22 | 2048 | 32 | 4 | 32K | 0.58B | 2 | 128 | 209 | 0.454 | 0.000 | 0.421 | 0.002 |
|  |  |  |  |  |  |  |  |  | 1 | 15.223 | 0.081 | 12.908 | 0.208 |
|  | 22 | 2048 | 32 | 4 | 32K | 0.68B | 2 | 256 | 209 | 0.457 | 0.000 | 0.423 | 0.002 |
|  |  |  |  |  |  |  |  |  | 1 | 15.383 | 0.080 | 13.062 | 0.211 |
|  | 22 | 2048 | 32 | 4 | 32K | 0.86B | 2 | 512 | 209 | 0.461 | 0.000 | 0.428 | 0.002 |
|  |  |  |  |  |  |  |  |  | 1 | 13.280 | 0.080 | 12.286 | 0.235 |
|  | 16 | 2048 | 8 | 8 | 50K | 0.81B | - | - | 53 | 1.227 | 0.002 | 1.206 | 0.005 |
|  |  |  |  |  |  |  |  |  | 1 | 8.423 | 0.081 | 6.378 | 0.262 |
|  | 16 | 2048 | 8 | 8 | 50K | 0.40B | 2 | - | 61 | 0.856 | 0.001 | 0.606 | 0.005 |
|  |  |  |  |  |  |  |  |  | 1 | 10.554 | 0.082 | 8.519 | 0.260 |
| Pythia | 16 | 2048 | 8 | 8 | 50K | 0.44B | 2 | 64 | 63 | 0.875 | 0.001 | 0.626 | 0.005 |
|  |  |  |  |  |  |  |  |  | 1 | 10.167 | 0.076 | 8.196 | 0.256 |
|  | 16 | 2048 | 8 | 8 | 50K | 0.48B | 2 | 128 | 59 | 0.892 | 0.001 | 0.642 | 0.005 |
|  |  |  |  |  |  |  |  |  | 1 | 10.410 | 0.079 | 8.402 | 0.258 |
|  | 16 | 2048 | 8 | 8 | 50K | 0.55B | 2 | 256 | 59 | 0.913 | 0.001 | 0.662 | 0.005 |
|  |  |  |  |  |  |  |  |  | 1 | 12.609 | 0.091 | 10.311 | 0.267 |
|  | 16 | 2048 | 8 | 8 | 50K | 0.70B | 2 | 512 | 53 | 0.956 | 0.002 | 0.702 | 0.006 |

**Table K.2: Hypothetical generation speedup of Recursive Transformers across three models. We utilized the measurements of per-token generation time calculated in Table K.1, which were calculated using an A100 40GB GPU, a prefix length of 512, and a decoding length of 2048. We only considered the time spent within Transformer blocks, simulating generation on the SlimPajama, RedPajama, and PG19 test sets. We used vanilla Transformer models, both with and without continuous sequence-wise batching (CSB), as our baselines. Our recursive models further enhance throughput by applying continuous depth-wise batching (CDB), leveraging looping and early-exiting techniques. The throughput improvements over the vanilla Transformer without and with sequence-wise batching are denoted as $\Delta_{V}$ and $\Delta_{Seq}$, respectively.**
|  |  | Uptrain | Looping | LoRA | Early-Exit Train | Batching | Few-shot Accuracy |  | Throughput $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | N-emb | PT | $N_{tok}$ | KD | Block | Init | Rank | Init | $N_{tok}$ | CE | KD | Type | Exit | Last | Mid 1 | Mid 2 | Batch | SlimP | RedP | PG19 | $\Delta_{V}$ | $\Delta_{Seq}$ |
|  | 1.99B | ✓ | 75B | ✗ | - | - | - | - | - | - | - | - | ✗ | 57.3 | - | - | 43 | 655 | 1228 | 1357 | $\mathbf{\times 1.00}$ | $\mathbf{\times 0.71}$ |
|  | 1.99B | ✓ | 75B | ✗ | - | - | - | - | - | - | - | CSB | ✗ | 57.3 | - | - | 43 | 1622 | 1604 | 1357 | $\mathbf{\times 1.41}$ | $\mathbf{\times 1.00}$ |
|  | 0.99B | ✓ | 60B | ✓ | 2 | Step | - | - | 15B | Agg (0.1) | ✓ | CDB | ✓ | 54.0 | 48.8 | - | 43 | 3159 | 3050 | 2421 | $\mathbf{\times{2.66}}$ | $\mathbf{\times{1.88}}$ |
|  | 1.07B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 54.0 | 40.8 | - | 41 | 2357 | 2255 | 1858 | $\mathbf{\times{2.00}}$ | $\mathbf{\times{1.41}}$ |
|  | 1.15B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 54.6 | 40.2 | - | 41 | 2355 | 2250 | 1844 | $\mathbf{\times{1.99}}$ | $\mathbf{\times{1.41}}$ |
|  | 1.30B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 55.2 | 40.5 | - | 39 | 2047 | 1976 | 1740 | $\mathbf{\times{1.78}}$ | $\mathbf{\times{1.26}}$ |
|  | 1.60B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 56.2 | 41.7 | - | 39 | 1806 | 1754 | 1598 | $\mathbf{\times{1.59}}$ | $\mathbf{\times{1.13}}$ |
|  | 1.07B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 53.1 | 43.3 | - | 41 | 2454 | 2357 | 1929 | $\mathbf{\times{2.08}}$ | $\mathbf{\times{1.47}}$ |
|  | 1.15B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 53.6 | 43.4 | - | 41 | 2445 | 2346 | 1926 | $\mathbf{\times{2.07}}$ | $\mathbf{\times{1.47}}$ |
|  | 1.30B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 54.6 | 43.2 | - | 39 | 2123 | 2056 | 1804 | $\mathbf{\times{1.85}}$ | $\mathbf{\times{1.31}}$ |
| Gemma | 1.60B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 55.2 | 44.0 | - | 39 | 1870 | 1819 | 1655 | $\mathbf{\times{1.65}}$ | $\mathbf{\times{1.17}}$ |
|  | 0.66B | ✓ | 60B | ✓ | 3 | Step | - | - | 15B | Agg (0.1) | ✓ | CDB | ✓ | 51.9 | 49.0 | 43.5 | 43 | 3120 | 3041 | 2729 | $\mathbf{\times{2.74}}$ | $\mathbf{\times{1.94}}$ |
|  | 0.74B | ✓ | 60B | ✓ | 3 | Avg | 64 | SVD | 15B | Agg (0.1) | ✗ | CDB | ✓ | 51.4 | 40.8 | 36.1 | 43 | 2334 | 2274 | 2059 | $\mathbf{\times{2.06}}$ | $\mathbf{\times{1.45}}$ |
|  | 0.82B | ✓ | 60B | ✓ | 3 | Avg | 128 | SVD | 15B | Agg (0.1) | ✗ | CDB | ✓ | 51.7 | 41.1 | 36.1 | 43 | 2290 | 2230 | 2007 | $\mathbf{\times{2.02}}$ | $\mathbf{\times{1.42}}$ |
|  | 0.97B | ✓ | 60B | ✓ | 3 | Avg | 256 | SVD | 15B | Agg (0.1) | ✗ | CDB | ✓ | 54.1 | 42.2 | 36.1 | 41 | 2281 | 2219 | 1984 | $\mathbf{\times{2.00}}$ | $\mathbf{\times{1.41}}$ |
|  | 1.27B | ✓ | 60B | ✓ | 3 | Avg | 512 | SVD | 15B | Agg (0.1) | ✗ | CDB | ✓ | 55.7 | 43.5 | 37.0 | 39 | 2181 | 2122 | 1900 | $\mathbf{\times{1.91}}$ | $\mathbf{\times{1.35}}$ |
|  | 0.74B | ✓ | 60B | ✓ | 3 | Avg | 64 | SVD | 15B | Agg (0.3) | ✗ | CDB | ✓ | 50.1 | 42.9 | 37.7 | 43 | 2427 | 2372 | 2143 | $\mathbf{\times{2.14}}$ | $\mathbf{\times{1.51}}$ |
|  | 0.82B | ✓ | 60B | ✓ | 3 | Avg | 128 | SVD | 15B | Agg (0.3) | ✗ | CDB | ✓ | 51.2 | 43.2 | 38.0 | 43 | 2376 | 2321 | 2084 | $\mathbf{\times{2.09}}$ | $\mathbf{\times{1.48}}$ |
|  | 0.97B | ✓ | 60B | ✓ | 3 | Avg | 256 | SVD | 15B | Agg (0.3) | ✗ | CDB | ✓ | 53.0 | 44.8 | 38.7 | 41 | 2359 | 2300 | 2039 | $\mathbf{\times{2.07}}$ | $\mathbf{\times{1.46}}$ |
|  | 1.27B | ✓ | 60B | ✓ | 3 | Avg | 512 | SVD | 15B | Agg (0.3) | ✗ | CDB | ✓ | 54.1 | 45.7 | 39.2 | 39 | 2251 | 2191 | 1975 | $\mathbf{\times{1.98}}$ | $\mathbf{\times{1.40}}$ |
|  | 0.97B | ✓ | - | - | - | - | - | - | - | - | - | - | ✗ | 43.3 | - | - | 329 | 1205 | 1220 | 1194 | $\mathbf{\times 1.00}$ | $\mathbf{\times 0.99}$ |
|  | 0.97B | ✓ | - | - | - | - | - | - | - | - | - | CSB | ✗ | 43.3 | - | - | 329 | 1227 | 1225 | 1194 | $\mathbf{\times 1.01}$ | $\mathbf{\times 1.00}$ |
|  | 0.48B | ✓ | 60B | ✓ | 2 | Step | - | - | 15B | Agg (0.1) | ✓ | CDB | ✓ | 44.8 | 41.8 | - | 233 | 2038 | 2023 | 1933 | $\mathbf{\times{1.66}}$ | $\mathbf{\times{1.64}}$ |
|  | 0.53B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 45.1 | 33.7 | - | 211 | 1733 | 1719 | 1617 | $\mathbf{\times{1.40}}$ | $\mathbf{\times{1.39}}$ |
|  | 0.58B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 45.2 | 33.9 | - | 209 | 1733 | 1717 | 1609 | $\mathbf{\times{1.40}}$ | $\mathbf{\times{1.39}}$ |
| TinyLlama | 0.68B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 45.6 | 33.9 | - | 209 | 1728 | 1714 | 1606 | $\mathbf{\times{1.39}}$ | $\mathbf{\times{1.38}}$ |
|  | 0.86B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 46.1 | 34.2 | - | 209 | 1716 | 1702 | 1581 | $\mathbf{\times{1.38}}$ | $\mathbf{\times{1.37}}$ |
|  | 0.53B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 44.6 | 36.1 | - | 211 | 1810 | 1796 | 1688 | $\mathbf{\times{1.46}}$ | $\mathbf{\times{1.45}}$ |
|  | 0.58B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 44.8 | 36.0 | - | 209 | 1802 | 1787 | 1668 | $\mathbf{\times{1.45}}$ | $\mathbf{\times{1.44}}$ |
|  | 0.68B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 44.8 | 36.2 | - | 209 | 1793 | 1779 | 1668 | $\mathbf{\times{1.45}}$ | $\mathbf{\times{1.44}}$ |
|  | 0.86B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 45.8 | 36.5 | - | 209 | 1778 | 1763 | 1637 | $\mathbf{\times{1.43}}$ | $\mathbf{\times{1.42}}$ |
|  | 0.81B | ✓ | 75B | ✗ | - | - | - | - | - | - | - | - | ✗ | 49.3 | - | - | 53 | 702 | 785 | 822 | $\mathbf{\times 1.00}$ | $\mathbf{\times 0.93}$ |
|  | 0.81B | ✓ | 75B | ✗ | - | - | - | - | - | - | - | CSB | ✗ | 49.3 | - | - | 53 | 829 | 827 | 822 | $\mathbf{\times 1.07}$ | $\mathbf{\times 1.00}$ |
|  | 0.40B | ✓ | 60B | ✓ | 2 | Step | - | - | 15B | Agg (0.1) | ✓ | CDB | ✓ | 45.4 | 42.0 | - | 61 | 1339 | 1333 | 1281 | $\mathbf{\times{1.71}}$ | $\mathbf{\times{1.60}}$ |
|  | 0.44B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 46.1 | 37.1 | - | 63 | 1205 | 1203 | 1140 | $\mathbf{\times{1.54}}$ | $\mathbf{\times{1.43}}$ |
|  | 0.48B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 46.2 | 37.8 | - | 59 | 1156 | 1180 | 1108 | $\mathbf{\times{1.49}}$ | $\mathbf{\times{1.39}}$ |
| Pythia | 0.55B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 46.5 | 38.0 | - | 59 | 1138 | 1139 | 1071 | $\mathbf{\times{1.45}}$ | $\mathbf{\times{1.35}}$ |
|  | 0.70B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 47.2 | 38.2 | - | 53 | 1051 | 1077 | 1021 | $\mathbf{\times{1.36}}$ | $\mathbf{\times{1.27}}$ |
|  | 0.44B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 45.1 | 39.0 | - | 63 | 1254 | 1252 | 1190 | $\mathbf{\times{1.60}}$ | $\mathbf{\times{1.49}}$ |
|  | 0.48B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 45.9 | 39.0 | - | 59 | 1200 | 1226 | 1153 | $\mathbf{\times{1.55}}$ | $\mathbf{\times{1.45}}$ |
|  | 0.55B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 46.0 | 39.4 | - | 59 | 1180 | 1180 | 1112 | $\mathbf{\times{1.50}}$ | $\mathbf{\times{1.40}}$ |
|  | 0.70B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 46.7 | 39.7 | - | 53 | 1088 | 1114 | 1058 | $\mathbf{\times{1.41}}$ | $\mathbf{\times{1.32}}$ |

**Table K.3: Generation time measurements of Gemma models on a single A100 40GB GPU with 16GB memory constraint. We measured generation time per token for both a batch size of 1 and the maximum batch size achievable by each model. The prefix length was set to 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 (ms) per token |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | $N_{L}$ | $d_{model}$ | $N_{head}$ | $N_{KV}$ | Vocab | N-emb | Block | Rank | Batch | Total | Emb | Transformer | Head |
|  |  |  |  |  |  |  |  |  | 1 | 22.577 | 0.084 | 20.937 | 0.801 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.98B | - | - | 111 | 0.207 | 0.001 | 0.188 | 0.010 |
|  |  |  |  |  |  |  |  |  | 1 | 13.576 | 0.079 | 10.819 | 0.815 |
|  | 18 | 2048 | 8 | 1 | 256K | 0.99B | 2 | - | 123 | 0.118 | 0.001 | 0.091 | 0.009 |
|  |  |  |  |  |  |  |  |  | 1 | 15.372 | 0.080 | 12.675 | 0.813 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.07B | 2 | 64 | 117 | 0.140 | 0.001 | 0.112 | 0.009 |
|  |  |  |  |  |  |  |  |  | 1 | 15.631 | 0.082 | 12.899 | 0.816 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.15B | 2 | 128 | 115 | 0.141 | 0.001 | 0.113 | 0.010 |
|  |  |  |  |  |  |  |  |  | 1 | 15.317 | 0.079 | 12.639 | 0.811 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.30B | 2 | 256 | 111 | 0.143 | 0.001 | 0.115 | 0.010 |
| Gemma |  |  |  |  |  |  |  |  | 1 | 15.379 | 0.080 | 12.692 | 0.807 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.60B | 2 | 512 | 103 | 0.158 | 0.001 | 0.127 | 0.011 |
|  |  |  |  |  |  |  |  |  | 1 | 10.528 | 0.080 | 7.411 | 0.817 |
|  | 18 | 2048 | 8 | 1 | 256K | 0.66B | 3 | - | 131 | 0.087 | 0.001 | 0.058 | 0.010 |
|  |  |  |  |  |  |  |  |  | 1 | 11.957 | 0.081 | 8.855 | 0.815 |
|  | 18 | 2048 | 8 | 1 | 256K | 0.74B | 3 | 64 | 123 | 0.105 | 0.001 | 0.075 | 0.009 |
|  |  |  |  |  |  |  |  |  | 1 | 11.898 | 0.080 | 8.787 | 0.816 |
|  | 18 | 2048 | 8 | 1 | 256K | 0.82B | 3 | 128 | 121 | 0.103 | 0.001 | 0.074 | 0.009 |
|  |  |  |  |  |  |  |  |  | 1 | 11.734 | 0.079 | 8.654 | 0.813 |
|  | 18 | 2048 | 8 | 1 | 256K | 0.97B | 3 | 256 | 117 | 0.106 | 0.001 | 0.076 | 0.009 |
|  |  |  |  |  |  |  |  |  | 1 | 11.986 | 0.080 | 8.856 | 0.809 |
|  | 18 | 2048 | 8 | 1 | 256K | 1.27B | 3 | 512 | 107 | 0.125 | 0.001 | 0.090 | 0.010 |
|  |  |  |  |  |  |  |  |  | 1 | 23.898 | 0.080 | 22.909 | 0.189 |
|  | 22 | 2048 | 32 | 4 | 32K | 0.97B | - | - | 1049 | 0.131 | 0.000 | 0.129 | 0.001 |
|  |  |  |  |  |  |  |  |  | 1 | 14.129 | 0.080 | 11.846 | 0.202 |
|  | 22 | 2048 | 32 | 4 | 32K | 0.48B | 2 | - | 1121 | 0.070 | 0.000 | 0.064 | 0.001 |
|  |  |  |  |  |  |  |  |  | 1 | 14.897 | 0.080 | 12.627 | 0.202 |
| TinyLlama | 22 | 2048 | 32 | 4 | 32K | 0.53B | 2 | 64 | 1105 | 0.073 | 0.000 | 0.068 | 0.001 |
|  |  |  |  |  |  |  |  |  | 1 | 15.090 | 0.081 | 12.778 | 0.205 |
|  | 22 | 2048 | 32 | 4 | 32K | 0.58B | 2 | 128 | 1089 | 0.074 | 0.000 | 0.069 | 0.001 |
|  |  |  |  |  |  |  |  |  | 1 | 14.962 | 0.081 | 12.659 | 0.201 |
|  | 22 | 2048 | 32 | 4 | 32K | 0.68B | 2 | 256 | 1065 | 0.076 | 0.000 | 0.071 | 0.001 |
|  |  |  |  |  |  |  |  |  | 1 | 15.284 | 0.083 | 12.950 | 0.206 |
|  | 22 | 2048 | 32 | 4 | 32K | 0.86B | 2 | 512 | 1017 | 0.080 | 0.000 | 0.075 | 0.001 |
|  |  |  |  |  |  |  |  |  | 1 | 13.341 | 0.081 | 12.326 | 0.239 |
|  | 16 | 2048 | 8 | 8 | 50K | 0.81B | - | - | 229 | 0.176 | 0.000 | 0.171 | 0.002 |
|  |  |  |  |  |  |  |  |  | 1 | 8.336 | 0.079 | 6.303 | 0.261 |
|  | 16 | 2048 | 8 | 8 | 50K | 0.40B | 2 | - | 241 | 0.121 | 0.000 | 0.086 | 0.002 |
|  |  |  |  |  |  |  |  |  | 1 | 10.408 | 0.081 | 8.353 | 0.262 |
| Pythia | 16 | 2048 | 8 | 8 | 50K | 0.44B | 2 | 64 | 233 | 0.133 | 0.000 | 0.097 | 0.002 |
|  |  |  |  |  |  |  |  |  | 1 | 10.426 | 0.082 | 8.378 | 0.259 |
|  | 16 | 2048 | 8 | 8 | 50K | 0.48B | 2 | 128 | 221 | 0.137 | 0.000 | 0.101 | 0.002 |
|  |  |  |  |  |  |  |  |  | 1 | 10.509 | 0.080 | 8.471 | 0.256 |
|  | 16 | 2048 | 8 | 8 | 50K | 0.55B | 2 | 256 | 205 | 0.151 | 0.000 | 0.115 | 0.002 |
|  |  |  |  |  |  |  |  |  | 1 | 11.254 | 0.080 | 9.241 | 0.257 |
|  | 16 | 2048 | 8 | 8 | 50K | 0.70B | 2 | 512 | 165 | 0.177 | 0.001 | 0.139 | 0.002 |

**Table K.4: Hypothetical generation speedup of Recursive Transformers across three models. We utilized the measurements of per-token generation time calculated in Table K.3, which were calculated using an A100 GPU with 16GB memory constraint, a prefix length of 64, and a decoding length of 256. We only considered the time spent within Transformer blocks, simulating generation on the SlimPajama, RedPajama, and PG19 test sets. We used vanilla transformer models, both with and without continuous sequence-wise batching (CSB), as our baselines. Our recursive models further enhance throughput by applying continuous depth-wise batching (CDB), leveraging looping and early-exiting techniques. The throughput improvements over the vanilla Transformer without and with sequence-wise batching are denoted as $\Delta_{V}$ and $\Delta_{Seq}$, respectively.**
|  |  | Uptrain | Looping | LoRA | Early-Exit Train | Batching | Few-shot Accuracy |  | Throughput $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | N-emb | PT | $N_{tok}$ | KD | Block | Init | Rank | Init | $N_{tok}$ | CE | KD | Type | Exit | Last | Mid 1 | Mid 2 | Batch | SlimP | RedP | PG19 | $\Delta_{V}$ | $\Delta_{Seq}$ |
|  | 1.99B | ✓ | 75B | ✗ | - | - | - | - | - | - | - | - | ✗ | 57.3 | - | - | 111 | 1740 | 3059 | 4796 | $\mathbf{\times 1.00}$ | $\mathbf{\times 0.63}$ |
|  | 1.99B | ✓ | 75B | ✗ | - | - | - | - | - | - | - | CSB | ✗ | 57.3 | - | - | 111 | 5287 | 5060 | 4796 | $\mathbf{\times 1.58}$ | $\mathbf{\times 1.00}$ |
|  | 0.99B | ✓ | 60B | ✓ | 2 | Step | - | - | 15B | Agg (0.1) | ✓ | CDB | ✓ | 54.0 | 48.8 | - | 43 | 3159 | 3050 | 2421 | $\mathbf{\times{2.50}}$ | $\mathbf{\times{1.59}}$ |
|  | 1.07B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 54.0 | 40.8 | - | 41 | 2357 | 2255 | 1858 | $\mathbf{\times{1.87}}$ | $\mathbf{\times{1.19}}$ |
|  | 1.15B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 54.6 | 40.2 | - | 41 | 2355 | 2250 | 1844 | $\mathbf{\times{1.87}}$ | $\mathbf{\times{1.19}}$ |
|  | 1.30B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 55.2 | 40.5 | - | 39 | 2047 | 1976 | 1740 | $\mathbf{\times{1.86}}$ | $\mathbf{\times{1.18}}$ |
|  | 1.60B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 56.2 | 41.7 | - | 39 | 1806 | 1754 | 1598 | $\mathbf{\times{1.73}}$ | $\mathbf{\times{1.10}}$ |
|  | 1.07B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 53.1 | 43.3 | - | 41 | 2454 | 2357 | 1929 | $\mathbf{\times{1.95}}$ | $\mathbf{\times{1.24}}$ |
|  | 1.15B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 53.6 | 43.4 | - | 41 | 2445 | 2346 | 1926 | $\mathbf{\times{1.95}}$ | $\mathbf{\times{1.24}}$ |
|  | 1.30B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 54.6 | 43.2 | - | 39 | 2123 | 2056 | 1804 | $\mathbf{\times{1.93}}$ | $\mathbf{\times{1.22}}$ |
| Gemma | 1.60B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 55.2 | 44.0 | - | 39 | 1870 | 1819 | 1655 | $\mathbf{\times{1.79}}$ | $\mathbf{\times{1.14}}$ |
|  | 0.66B | ✓ | 60B | ✓ | 3 | Step | - | - | 15B | Agg (0.1) | ✓ | CDB | ✓ | 51.9 | 49.0 | 43.5 | 43 | 3120 | 3041 | 2729 | $\mathbf{\times{2.62}}$ | $\mathbf{\times{1.66}}$ |
|  | 0.74B | ✓ | 60B | ✓ | 3 | Avg | 64 | SVD | 15B | Agg (0.1) | ✗ | CDB | ✓ | 51.4 | 40.8 | 36.1 | 43 | 2334 | 2274 | 2059 | $\mathbf{\times{1.87}}$ | $\mathbf{\times{1.19}}$ |
|  | 0.82B | ✓ | 60B | ✓ | 3 | Avg | 128 | SVD | 15B | Agg (0.1) | ✗ | CDB | ✓ | 51.7 | 41.1 | 36.1 | 43 | 2290 | 2230 | 2007 | $\mathbf{\times{1.90}}$ | $\mathbf{\times{1.20}}$ |
|  | 0.97B | ✓ | 60B | ✓ | 3 | Avg | 256 | SVD | 15B | Agg (0.1) | ✗ | CDB | ✓ | 54.1 | 42.2 | 36.1 | 41 | 2281 | 2219 | 1984 | $\mathbf{\times{1.86}}$ | $\mathbf{\times{1.18}}$ |
|  | 1.27B | ✓ | 60B | ✓ | 3 | Avg | 512 | SVD | 15B | Agg (0.1) | ✗ | CDB | ✓ | 55.7 | 43.5 | 37.0 | 39 | 2181 | 2122 | 1900 | $\mathbf{\times{1.62}}$ | $\mathbf{\times{1.03}}$ |
|  | 0.74B | ✓ | 60B | ✓ | 3 | Avg | 64 | SVD | 15B | Agg (0.3) | ✗ | CDB | ✓ | 50.1 | 42.9 | 37.7 | 43 | 2427 | 2372 | 2143 | $\mathbf{\times{1.94}}$ | $\mathbf{\times{1.23}}$ |
|  | 0.82B | ✓ | 60B | ✓ | 3 | Avg | 128 | SVD | 15B | Agg (0.3) | ✗ | CDB | ✓ | 51.2 | 43.2 | 38.0 | 43 | 2376 | 2321 | 2084 | $\mathbf{\times{1.97}}$ | $\mathbf{\times{1.25}}$ |
|  | 0.97B | ✓ | 60B | ✓ | 3 | Avg | 256 | SVD | 15B | Agg (0.3) | ✗ | CDB | ✓ | 53.0 | 44.8 | 38.7 | 41 | 2359 | 2300 | 2039 | $\mathbf{\times{1.92}}$ | $\mathbf{\times{1.22}}$ |
|  | 1.27B | ✓ | 60B | ✓ | 3 | Avg | 512 | SVD | 15B | Agg (0.3) | ✗ | CDB | ✓ | 54.1 | 45.7 | 39.2 | 39 | 2251 | 2191 | 1975 | $\mathbf{\times{1.67}}$ | $\mathbf{\times{1.06}}$ |
|  | 0.97B | ✓ | - | - | - | - | - | - | - | - | - | - | ✗ | 43.3 | - | - | 1049 | 6856 | 7481 | 4090 | $\mathbf{\times 1.00}$ | $\mathbf{\times 0.96}$ |
|  | 0.97B | ✓ | - | - | - | - | - | - | - | - | - | CSB | ✗ | 43.3 | - | - | 1049 | 7709 | 7481 | 4090 | $\mathbf{\times 1.05}$ | $\mathbf{\times 1.00}$ |
|  | 0.48B | ✓ | 60B | ✓ | 2 | Step | - | - | 15B | Agg (0.1) | ✓ | CDB | ✓ | 44.8 | 41.8 | - | 233 | 2038 | 2023 | 1933 | $\mathbf{\times{1.70}}$ | $\mathbf{\times{1.62}}$ |
|  | 0.53B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 45.1 | 33.7 | - | 211 | 1733 | 1719 | 1617 | $\mathbf{\times{1.38}}$ | $\mathbf{\times{1.32}}$ |
|  | 0.58B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 45.2 | 33.9 | - | 209 | 1733 | 1717 | 1609 | $\mathbf{\times{1.36}}$ | $\mathbf{\times{1.30}}$ |
| TinyLlama | 0.68B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 45.6 | 33.9 | - | 209 | 1728 | 1714 | 1606 | $\mathbf{\times{1.34}}$ | $\mathbf{\times{1.28}}$ |
|  | 0.86B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 46.1 | 34.2 | - | 209 | 1716 | 1702 | 1581 | $\mathbf{\times{1.28}}$ | $\mathbf{\times{1.23}}$ |
|  | 0.53B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 44.6 | 36.1 | - | 211 | 1810 | 1796 | 1688 | $\mathbf{\times{1.45}}$ | $\mathbf{\times{1.38}}$ |
|  | 0.58B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 44.8 | 36.0 | - | 209 | 1802 | 1787 | 1668 | $\mathbf{\times{1.41}}$ | $\mathbf{\times{1.35}}$ |
|  | 0.68B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 44.8 | 36.2 | - | 209 | 1793 | 1779 | 1668 | $\mathbf{\times{1.39}}$ | $\mathbf{\times{1.33}}$ |
|  | 0.86B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 45.8 | 36.5 | - | 209 | 1778 | 1763 | 1637 | $\mathbf{\times{1.33}}$ | $\mathbf{\times{1.27}}$ |
|  | 0.81B | ✓ | 75B | ✗ | - | - | - | - | - | - | - | - | ✗ | 49.3 | - | - | 229 | 4273 | 5346 | 5149 | $\mathbf{\times 1.00}$ | $\mathbf{\times 0.89}$ |
|  | 0.81B | ✓ | 75B | ✗ | - | - | - | - | - | - | - | CSB | ✗ | 49.3 | - | - | 229 | 5813 | 5724 | 5149 | $\mathbf{\times 1.13}$ | $\mathbf{\times 1.00}$ |
|  | 0.40B | ✓ | 60B | ✓ | 2 | Step | - | - | 15B | Agg (0.1) | ✓ | CDB | ✓ | 45.4 | 42.0 | - | 61 | 1339 | 1333 | 1281 | $\mathbf{\times{1.77}}$ | $\mathbf{\times{1.57}}$ |
|  | 0.44B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 46.1 | 37.1 | - | 63 | 1205 | 1203 | 1140 | $\mathbf{\times{1.44}}$ | $\mathbf{\times{1.28}}$ |
|  | 0.48B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 46.2 | 37.8 | - | 59 | 1156 | 1180 | 1108 | $\mathbf{\times{1.32}}$ | $\mathbf{\times{1.17}}$ |
| Pythia | 0.55B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 46.5 | 38.0 | - | 59 | 1138 | 1139 | 1071 | $\mathbf{\times{1.22}}$ | $\mathbf{\times{1.08}}$ |
|  | 0.70B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.1) | ✓ | CDB | ✓ | 47.2 | 38.2 | - | 53 | 1051 | 1077 | 1021 | $\mathbf{\times{0.98}}$ | $\mathbf{\times{0.87}}$ |
|  | 0.44B | ✓ | 60B | ✓ | 2 | Avg | 64 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 45.1 | 39.0 | - | 63 | 1254 | 1252 | 1190 | $\mathbf{\times{1.50}}$ | $\mathbf{\times{1.33}}$ |
|  | 0.48B | ✓ | 60B | ✓ | 2 | Avg | 128 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 45.9 | 39.0 | - | 59 | 1200 | 1226 | 1153 | $\mathbf{\times{1.37}}$ | $\mathbf{\times{1.22}}$ |
|  | 0.55B | ✓ | 60B | ✓ | 2 | Avg | 256 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 46.0 | 39.4 | - | 59 | 1180 | 1180 | 1112 | $\mathbf{\times{1.27}}$ | $\mathbf{\times{1.12}}$ |
|  | 0.70B | ✓ | 60B | ✓ | 2 | Avg | 512 | SVD | 15B | Agg (0.3) | ✓ | CDB | ✓ | 46.7 | 39.7 | - | 53 | 1088 | 1114 | 1058 | $\mathbf{\times{1.02}}$ | $\mathbf{\times{0.90}}$ |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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.
2. 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.
Reusing the transformer computation block has been explored before [1] without properly cited within this paper.

[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).
soundness: 3
strengths: 1. Propose a novel model structure that can effectively compress pre-trained large language models for deployment.
2. LoRA adaptation makes the model structure more flexible to adapt to different depths.
3. Depth-wise Batching translates the advantages of the proposed recursive transformer block to real deployment acceleration.
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.
weaknesses: 1. Experiment design is not so convincing, which can be further improved.
2. Authors do not push the recursive computation block’s advantage to the limit, which can be interesting to explore
3. Some highly-related papers are missing in the reference

---

confidence: 3
contribution: 4
initial_rating: 6
presentation: 3
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)?
- 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?
soundness: 3
strengths: - The core concept of recursive layer sharing combined with LoRA-based relaxation is straightforward yet powerful.
- 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.
- 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.
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.
weaknesses: **Major**
- 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].
- 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.
- The study would benefit from experiments on optimal recursion depths across varying model scales and task domains.

**Minor**
- 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.
- Typos: "achieved a 10%p performance" on page6.

[1] Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

[2] Zamba: A Compact 7B SSM Hybrid Model

---

confidence: 5
contribution: 3
initial_rating: 8
presentation: 4
questions: 1. I wonder the performance of relaxed recursive transformers on larger scales.
2. How does the authors compare the proposed parameter sharing methods with other model deployment optimization methods, such as quantization, sparcification，pruning or distillation?
soundness: 3
strengths: 1. The paper is well motivated on a hot topic in the field of LLMs.
2. The paper is well organized and the insight of the method is clearly stated.
3. The experimental results are convincing. I appreciate the way the authors present the takeaways for relaxed recursive transformers.
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.
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.

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 3
questions: Please refer to the weaknesses section.
soundness: 2
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.

2. The incorporation of Relaxed Recursive Transformers with LoRA modules successfully improves flexibility and performance.

3. The proposed Continuous Depth-wise Batching method, which takes advantage of the recursive model structure and early exiting, increases inference throughput.
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.
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. 

2. 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.

3. 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.

4. 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.

5. 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. 

6. 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. 

7. 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.

8. 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.

Meta-review: The paper proposes relaxed recursive transformers, a new approach to parameter sharing in LLMs by iteratively reusing a single block of layers. The incorporation of LoRA enhances model flexibility without significant performance loss. Additionally, the paper introduces continuous depth-wise batching for improved inference throughput. Reviewers praised the innovative combination of recursive architecture and LoRA, as well as the comprehensive experiments. I recommend acceptance.

Decision: Accept (Poster)