Boosting Inference Efficiency: Unleashing the Power of Parameter-Shared Pre-trained Language Models
Authors: Weize Chen
Xiaoyue Xu
Xu Han
Yankai Lin
Ruobing Xie
Zhiyuan Liu
Maosong Sun
Jie Zhou
Venue: EMNLP 2023 Findings (2023)
Keywords: pre-trained language model
parameter sharing
inference acceleration

Abstract: Parameter-shared pre-trained language models (PLMs) have emerged as a successful approach in resource-constrained environments, enabling substantial reductions in model storage and memory costs without significant performance compromise. However, it is important to note that parameter sharing does not alleviate computational burdens associated with inference, thus impeding its practicality in situations characterized by limited stringent latency requirements or computational resources. Building upon neural ordinary differential equations (ODEs), we introduce a straightforward technique to enhance the inference efficiency of parameter-shared PLMs. Additionally, we propose a simple pre-training technique that leads to fully or partially shared models capable of achieving even greater inference acceleration. The experimental results demonstrate the effectiveness of our methods on both autoregressive and autoencoding PLMs, providing novel insights into more efficient utilization of parameter-shared models in resource-constrained settings.

Boosting Inference Efficiency: Unleashing the Power of Parameter-Shared Pre-trained Language Models
Weize Chen'; Xiaoyue Xu"; Xu Han'; Yankai Lin Ruobing Xie Zhiyuan L Maosong Sun Jie Zhou
'NLP Group, DCST, IAI, BNRIST, Tsinghua University. Beijing
"Gaoling School of Artificial Intelligence, Renmin University of China, Beijing 3 Patterr Recognition Center, WeChat AI, Tencent Inc. chenwz21 tsinghua. edu. cn hanxu2822@tsinghua. edu. cn
Abstract
few billion parameters such as LLaMA (Touvron et al., 2023), Vieuna (Chiang et al., 2023), and Al- paca (Taori et al., 2023) have exhibited remarkable capabilities, their application remains constricted in numerous resource-constraingg environments.
Parameter-sharee pre- trained language mod- els (PLMs) have emerged as a successful ap- proach in resource-constrainee environments, enabling substantial reductions in model stor- age and memory costs without significant per- formance compromise. However, impor- tant to note that parameter sharing does not al- leviate computational burdens associated with inference, thus impeding its practicality in sit- uations characterized by limited stringent la- tency requirements or computational resources. Building upon neural ordinary differential equa- tions (ODEs), we introduce a straightforward technique to enhance the inference efficiency of parameter-shared PLMs. Additionally, we propose a simple pre-training technique that leads fully or partially shared models capa- ble of achieving even greater inference accel- eration. The experimental results demonstrate the effectiveness of methods both au- loregressive and autoencoding PL.Ms, provid- ing novel insights into more efficient utiliza- tion of parameter-shared models in resource- constrained settings.
In contrast to the monumental advances in PLMs, the real-world applications often still favor more established models such as BERT and GPT-2 (Rad- ford et al., 2019). These models, despite their rela- tively fewer parameters, deliver satisfactory perfor- mance across many tasks while requiring signifi- cantly less resources. This balance offers an appeal- ing trade-off between performance and cost. More- over, parameter sharing techniques have success- fully demonstrated that model size can be greatly reduced without significant performance degrada- tion, mitigating the storage burden and yielding better cost-effectivenness. This has sparked inter- est in parameter-shared PLMs (PSPLMs) like AL- BERT (Lan et al., 2020), a derivative of the BERT architecture that shares parameters across lay- ers, effectively reducing model size and memory requirements. Still, it's critical to recognize that parameter sharing alone doesn't guarantee reduced inference time since the number of layers processed during each forward pass remains unchanged. In other words, while resolves the storage issue, it does not address the computational challenge.
1 Introduction
In recent years, there has been a significant increase in the number of parameters of PLMs. This be- gan with the advent of BERT (Devlin et al., 2019), containing 340 million parameters, and has esca- lated to models like T5 (Raffel et al., 2020), GPT- 3 (Brown et al., 2020), and PALM (Chowdhery et al., 2022), with the latter reaching astounding 540 billion parameters. The trend of PLM expan- sion has, undeniably, improved performance across numerous tasks. Nonetheless, the corresponding in- crease in computation and storage requirements has raised substantial barriers for scenarios character- ized by stringent latency requirements or resource limitations. While PLMS encompassing merely a
Early exit techniques promise reduce the num- ber of layers processed during inference by halt- ing computation at early layers (Zhou et al., 2020; Wang et al., 2022; Schuster et al., 2022). While ef- fective, these methods typically require additional trained classifiers or computationally expensive dot products between the vocabulary matrix and the hidden states at each layer. This circumstance prompts the question: Can a method be proposed to reduce the inference cost without introducing extra modules or computations, and could be complementary to early exit techniques, allowing their combined use for further acceleration?
"Equal Contribution Corresponding author.

---

In this study, we show that the problem can be well addressed in PSPLMs. Specifically, we illus- trate how significant acceleration in PSPLM infer- ence can be achieved through our straightforward yet effective technique. This technique, inspired by the principles of neural ODEs, accelcrate the infer- ence without necessitating the addition of modules calculations to the model. Hence, in addition to the inherent storage efficiency, our method notably makes PSPLMs be computational efficient. We also introduce pre-training method for PSPLMs with a slightly altered forward propagation rule. Experiments reveal that our proposed pre-training method prepares the model for even greater accel- eration during inference, and we give theoretical explanation to aptly support our method.
Lakge step inference
lvr
hr
Area Correct Prediction
ho
Figure 1: An illustration of reducing the iteration count during inference by enlarging the step size.
Euler's method approximates the solution of this ODE from = to T by dividing the interval into n steps, each with step size 8i, such that Ee S; = T. The method iteratively computes the following formula from = 0 to n 1: 1:
We further extend the application of our method beyond the domain of fully shared PSPLMs. Our research demonstrates the potential of our acceler- ation strategy in the context of more complex and capable partially-shared PLMs, and even hints at its applicability to unshared models. This broader applicability shows the flexibility and robustness of our approach. Additionally, competition with other acceleration strategies. We demonstrate that our method can be combined or- thogonally with early exit techniques, thereby facil- itating further acceleration. Remarkably, this syn- ergy of methods makes the partially-shared model surpass its unshared equivalent within an equiva- lent computational budget.
y y: + 8, fe(1))
(3)
The final value yn the solution Eq. at time T. The correspondence between Eq. 1 and Eq. 3 evident. T-layer resid- ual network can be interpreted as parameterizing the vector field f that characterizes the derivative along the path from the input space to the final out- put space. The ODE perspective generalizes the concept of depth in residual networks to the con- tinuous domain, where the notion of progression from input to output is captured by the continuous time rather than discrete depth or layer index.
as an approximation of
method is not in
In essence, our work offers novel route to accelerate the inference for PSPLMs, and lay a novel foundation for unleashing the PSPLMs' po- tential, offering critical insights for deployment of PSPLMs in resource-constrrange settings.
During the inference process of a trained model, the vector field remains fixed because the param- eters are frozen. As a result, the model's infer- ence can be seen as solving an ODE within this vector field using Euler's method, where the ini- tial value corresponds the input embedding, and the solution time is T. Furthermore, the pre-norm Transformer architecture also type of residual network (see Appendix A). Therefore, the ODE perspective we have presented can be applied to the pre-norm Transformer architecture as well.
ODE Perspective on Residual Networks 2
We begin by providing brief overview of the re- lationship between residual networks and ODEs, which forms the fundamental basis of our research.
In T T-layer residual network, we denote the layer as fe.- The update formulation for the hid- den state h can be expressed as:
3 Method
3.1
Enlarging the Step Size
h h4 + fo. (he).
(1)
In the process of solving the ODE, the choice of the step size s; in Eq. 3 has significant impact on the speed and accuracy of the solution. Given the final time T, a larger step size reduces the number of iterations, resulting in faster computation but (2) decreased accuracy. This trade-off allows us to
Remarkably, this update scheme aligns with Euler's method for solving ODEs. Consider an ODE:
00) yo: y'(1)= = f(()(()),

---

Small step pre- -raatraned model
K/g
Bi gi. bly
P12
Iay herl 114. Tef.l
( L-24, 5-1, sets of parameters
(3) Linear Interpolation
Loyer Paramilioss
Area of Correct Prediction
P1D)
Layer Panamiteess
Step=1 pre-trained model
23
Figure 2: An illustration of the difference between the models pre-trained with different step size.
Figure An illustration of purtially-sharde model. The example shows how a model with number of layers 24, step size and n 3 sets of parameters determines the layer parameters for 2.
sacrifice certain degree of solution accuracy in exchange for inference speed. To expedite model inference, we propose employing larger step size during inference compared to training (Fig. 1).
h hrl < K(1 + 8*)s
(5)
Specifically, for PSPLMs, the vector field f at any given time t is parameterized by exactly the same set of parameters. Consequently, the time de- pendence in 0g (as shown in Eq. 1) can be omitted, leading the forward rule 1 h, + fo(h.) where 0 now represents the shared layer parame- ters. By applying different scaling factors B, > 1 to the original step size 8 = 1 at different lay- ers, and reducing the number of layers such that E. 3e s T, the model still mathematically solves the ODE from t = 0 to T, albeit using larger step sizes 3is Bg at each layer. The updated forward rule can be written
where 8r is the largest scaling factor used across all the layers, and K is constant determined by the model parameters. This inequality indicates that the difference between hr and hy bounded to the magnitude of the largest scaled-up step size 3 s employed during inference. Assuming that the value of K is approximately the same for models pre-trained with different step scales, then when the step sizes are scaled up by the same factors, the model pre-trained with smaller produces final approximated hidden state that is closer to the hidden state obtained using its original step size. Empirically, we will show in Section 4.3 that PSPLMs pre-trained with reasonably small step size can achieve improved performance when re- ducing the number of iterations during inference.
h 9. = h4 Br fo())) B B 1.
In practice, we perform minimal search to deter- mine set of suitable { values. Section 4.2, we will demonstrate that by simply changing the forward rule to Eq. 4, the inference of existing PSPLMs can already be accelerated while main- taining overall performance to satisfactory extent.
3.3 Generalizing to Partially-Shared PLMs
Shared-parameter models employ the same set of parameters to parameterize the derivative at differ- ent discrete time steps during training. This prop- erty offers the models the ability generalize from discrete time to continuous time during inference (Eq. 4). On the other hand, unshared models use distinct parameters to parameterize the derivative at each discrete time step, making challenging to apply continuous scaling factor 3 to the step size during inference. For instance, if we use 30 1.3 and 1, the model would need to provide the derivative at t 1.3 the next iteration, which is unattainable the unshared model can only pro- vide the derivative at using layer 3, but not at any intermediate time. However, we will demonstrate that pre- pre-training a partially-shared PLM with dependdet parame-
3.2 Pre- Te-Trriingg with Smaller Step Size
Ideally, if the approximate result hr obtained us- ing scaled-up step sizes is close to the result hr obtained with the original step size, the overall performance can be greatly preserved. Convention- ally, pre-training of PSPLMs is conducted with a step size of 1. However, from theoretical standpoint, the error analysis of Buler's method suggests that selecting smaller step size during pre-training may enable better acceleration during inference. Under certain mild conditions, we prove in Appendix the following inequality:
using layer at =

---

ters represented by a piece-wise linear function can enhance the model's 's capabilities while benefiting from the accleration method we introduce in Sec- tion 3.1. Given a language model with L layers and step size 8, and n (2 S n I L) sets of layer parameters, denoted as 0 = { 01....00-111 we uniformly position the sets of parameters within the range from 0 to (L 1)s, with each interval spanning A (I. 1)s/(n 1).
et al.. 2021). Detailed information on hyper- parameters and pre-training configurations can be found in Appendix D. It is important to mention that while we exclusively focus on parameter shar- ing among layers, our proposed method can be scamlessly incorporated alongside other parameter- reduction techniques such as embedding factoriza- tion used in ALBERT.
For downstream task evaluation, we measure the zero-shot perplexity (PPL) on Wikitext- 103 (Mer- ity et al., 2017), and zero-shot accuracy on LAM- BADA (Paperno et al., 2016) for GPT-2 models. And as for BERT models, they are fine-tuned on different tasks including MNLI, SST-2 (Wang et al., 2019), RACE (Lai et al., 2017). SQuAD and SQuAD2.0 O (Rajpurkar et al., 2016) separately. Configuration details and metrics for these down- stream tasks can be found in Appendix E. Dur- ing the inference, we experiment with different iteration counts and for each count, we perform a minimal search on the 3 for each layer within the set (1.0, 1.1,.... 3.00} using Optuma (Akiba et al., 2019), and report the best results. Unless explicitly stated otherwise, both BERT and GPT-2 models mentioned hereafter are parameter-shared.
To determine the parameters at a specific time, denoted as t, we define the function P(t) that re- turns the parameters at time f. We first identify the left and right boundary indices of the interval in which & resides, denoted as 1 and r, respectively. Subsequently, we perform linear interpolation be- tween Or and 0, to obtain the parameters at time I. P(t) can be formally written as:
A (L- -1)s n 1-1 1 t
t 1 r=
1A (0, - 8.)-
P(t) - 8
(6)
Fig. 3 shows an example of this process. The model is referred to partially-shared since it does not share all layer parameters; instead, it shares a set of parameters and uses interpolation to obtain the parameters at different time. Notably, when n L, it becomes the unshared model, and the fully- shared model is special case if we allow n = 1. By learning to use linear interpolation to derive parameters for different time steps during training. the model can naturally generalize to the continu- ous domain and provide derivatives for any time step during inference. We will demonstrate at Sec- tion 4.4 that these partially-shared PLMs exhibit notable advantages, including better or compara- ble performance to their unshared counterparts, as well as the ability to enable accelerated inference through a scaled-up step size.
4.2 Inherent Highways: Scaling Up Step Sizes
As described in Section 3.1, from the perspective of ODEs, we can naturally accelerate PSPLMs by increasing the step sizes. In other words, there may be inherent highways in PSPLMs, and we may utilize them by increasing the step size and decreasing the number of iterations.
To validate the presence of these inherent high- ways in PSPLMs, we pre-train GPT--2arge and BERT lage models under the conventional setting (i.e., 1), and evaluate their inference perfor- mance on variety of downstream tasks with differ- ent iteration counts and step sizes. The results are shown in Table 1. Additionally, we compute rela- tive changes in performances for reduced iterations as Produced -piii where represents the performance, Porig
4 Experiments and Analyses
4.1 Experimental Setups
We investigate the effectiveness of our suggested in- ference acceleration technique on both autoregres- sive and autoencoding models. Specifically, we pre- train GPT-2targe models and pre-norm BERT models both with shared parameters under diverse settings, as elucidated in the subsequent sections. All GPT-2 models are pre-trained on the OpenWeb- Web- Text dataset (Radford et al., 2019), and all BERT models are pre-trained on the Pile dataset (Gao
and we report these values in parentheses.
Our experimental results reveal that a clever re- duction in the iteration count presents an oppor- tunity for substantial computational savings while maintaining most of the model performance. When the iteration count decreases from 24 to 20, the performance impact across all datasets is virtually negligible. For BERT, variations in performance are consistently contained within a margin of +

---

BERT
GPT-2
#lters
Speed
MNLIT
SST-21
RACE1
SQuADt
SQuAD21
Speed
Wiki-103
LAMBADAT
24
1.00x
83.6
91.1
64.3
90.2
81.5
1.00x
33.0
31.1
20
1.20x
83.6 (10,08)
90.8 (.0,3%)
64.1 (0.39)
90.0 (-22%)
81.5 (eousss)
1.16x
33.5 (11.4%)
(9.6(
16
1.47x
83.3 (-0,45)
90.9 (0..%)
62.9 (.222%
89.3 (-10%)
80.3 (-155%)
1.40x
35.3
(x7.0%)
30.9 (..8%)
12
1,92x
81.1
(-3,0%
90.4 (-08%)
59.4 (-7.5%
83.0 (-8.0%)
65.0 (-202%)
1.77x
105.1
(421849)
5.3 (-83.19)
Table I: Inference performance of PSPLMs pre-trained with step size 1. Values in parentheses indicates relative change from non-reduced iteration counts. Speed reflects the acceleration of the forward pass wallclock time.
0.3% across all tasks. Simultaneously, the GPT-2 model shows only slight increase in perplexity on Wikitext- 103 from 33.0 to 33.5. Even with further reduction iteration count to 16, the mod- els continue to deliver respectable performance. For BERT, the majority of tasks report a minimal performance decrease, with the highest decrease appearing in the RACE task at -2.2%. For GPT-2 model, although the perplexity on Wikitext- 103 in- creases to 35.3 and the accuracy on LAMBADA decreases 30.9, the performance still stays within an acceptable range.
performs marginally better. However, it should still be noted that extremely small step size of 0.01 still negatively impacts the model's performance across all tasks. But overall, reasonably small step size does not impact the model's capacity.
Small step sizes enhance performance reten- tion when reducing the iteration count. Looking at the performance at 12 iterations across models with varying step sizes, as expected, we generally observed decline in comparison to the perfor- mance attained at 24 iterations. However, models with smaller step sizes exhibit remarkably better performance retention. Particularly noteworthy is GPT models pretrained with small step sizes, as they exhibit a significantly better zero-shot per- formance retention on both the LAMBADA and Wikitext- 103 tasks. Notably, as we do not inro- duce any additional computational overhead, the speedups of BERT and GPT are the same those reported in Table 1, which almost linear the reduced number of iteration, while the performance retention is significantly improved.
Overall, these results suggest that the step size per iteration can be scaled up and leads to a re- duction in the number of iterations in conventional PSPLMs without a significant compromise on per- formance. In essence, our approach enables a com- putational reduction by approximately 1/3 to 1/6. However, further reductions in iteration does result in performance degration, which can be addressed in the subsequent section.
4.3 Acceleration Boost: Mini-Step Pretraining
Reducing iteration count enhances perfor- mance on certain datasets. This finding aligns with observations made in earlier studies on early exits, suggesting that preventing models from over- thinking can enhance both accuracy and robust- ness (Zhou al., 2020; Balagansky and Gavrilov, 2022). However, unlike these previous studies, our approach demonstrates that we can effectively and easily prevent overthinking for models pretrained with smaller step sizes without auxiliary modules.
In Section 3.2, we posited that pre-training PSPLMs with small step sizes may make the mod- els more conducive to acceleration during the infer- ence. This section provides empirical validation of these theoretical insights.
4.3.1 Performance Across Downstream Tasks
For fair comparison, we maintain identical pre- training configurations and data and train 4 models with step size 1, 0.1, 0.05, and 0.01 respectively. The performance is presented in Fig. 4, where we have several noteworthy observations:
4.3.2 Analyzing the Possible Mechanism
We further explore why models pretrained with a small step size result in PSPLMs that are more effi- ciently accelerated during inference. Our analysis reveals two main advantages for models pretrained with smaller step sizes:
Small step sizes do not detrimentally affect performance within a reasonable range. We first look at the performances when the iteration count is not reduced (24 in the figures). BERT mod- els pretrained with smaller step sizes demonstrate comparable, and in some instances superior, perfor- mance on various downstream tasks in comparison to the conventionally pretrained BERT. The only exception is the MNLI task, where the latter model
Reduced absolute and relative difference when the iteration count is decreased. We de- crease the iteration count for all models to 20, 16, and 12, use the searched step scales and calculate the absolute difference, denoted as | by - hy

---

06.
633
22
62.
96.3
119 Mare
84
0.01
0.03
26
24
14
16
25
23
24
22
14
10
20
22
22
14
26
28
20
24
(a) MNLI
(b) SST-2
(c) RACE
(d) SQuAD
Sice S.c
3.05
0.2
3.0L
25.0
12.5
2xp 2.2
0.1
926
3.05
0.6
1.00
2011
(c) SQuAD2.0
(f) LAMBADA
(g) Wikitext-108
Figure 4: The inference performance of parameter-sharee models pre-traine pe-ttrined with different step size. (a-e) The accuracy of BERT MNLI, SST-2 and RACT, and Fl score SQuAD and SQuAD2.0. (f-g) The zero-shot accuracy of GPT-2 on LAMBADA, and the zero-shot perplexity of GPT-2 on Wikitext-100
and the relative difference, expressed as || ht hr]]. Here we keep the notations consistant with Eq. 5. These values represent the difference between the approximated and the original final hidden states.
over 0.8 across all layers, suggesting an almost par- allel alignment of the derivatives at different time, that is, a smoother vector field. In other words, the paths from the input embedding to the final output are more "straight" for models pre-trained with small step size, thus allowing us reduce the number of iteration and enlarge the step size during inference.
As demonstrated in Fig. 10, the final hidden state approximation from model pre-trained with smaller step size presents a closer resemblance to the original final hidden state, both terms ab- solute and relative difference. These observations suggest that when the number of iterations de- creased, models pre-trained with smaller step sizes could yield results that align more closely with those from models with unreduced iterations most tasks, thereby better preserving performance.
4.4 Expanding Horizons: Partial Sharing
In this section, we pre-train partially-shared PLMs, and apply our method described in Section 3.3 to them. This experimental validation substantiates the possibility of inference acceleration in more complex, partially-shared, and even unshared mod- els.
Enhanced smoothness in the vector field. To further our analysis, we compute the cosine sim- ilarity between the derivatives (x) produced by the model at two consecutive iterations, denoted as 1)) C 1,2,....22 23. The
For BERT lurge and GPT-2 we conduct pre- training with 12 sets of parameters and step sizes of 0.1 and 0.05, respectively. We establish baselines by pre-training unshared BERT and GPT models using the equivalent configurations, except that the total number of layers is set 12. This en- sured the same number of parameters between our partially-sharee models and the baseline models.
CosSim (f results are plotted in Fig. 5.
Figs. 5 and 8 reveal an increasing trend of co- sine similarity as the layer index increases, with smaller step size generally resulting in higher co- sine similarity in the early layers. Although step size of 0.1 also appears have lower similarities for the first few layers, there is swift increase as the layer index increases. The cosine similarities for step sizes of 0.01 and 0.05 consistently remain
The results of downstream tasks illustrated in Figs. 6 and The unshared model performance represented by the red dashed line in each figure. As anticipated, partially-share models, benefiting from the increased number of parameters, signifi-

---

63
2A
64
8
Step Sue
1a
3.4
1.5
30
15
20
Layar hoda
Iidox
(a) MNLI
(b) RACE
(c) LAMBADA
(d) Wikitext- 103
Figure 5: The cosine similarity between the derivatives given by the model at two consecutive iterations.
Size
Size
6-1
0.02
0:22
0.28
6.2
70
$5.0
$4.5
90
Stag
130
0.33
0.5
18
Inttrrttora Infernnc
#irtandaraa Infenmnce
(a) MNLI
(b) RACE
(c) LAMBADA
(d) Wikitext 103
Figure 6: inference performance partially parameter-shane models with n 12 sets parameters. The red dashed line represents the performance of the unshared 12 layer model.
cantly outperform those fully-shared counterparts in Section 4.3. Furthermore, because they have larger iteration count, they outperform unshared models with the same parameter count. The results show the feasibility of pre-training partiall--shared PLMs via linear parameter interpolation.
85.5
84.2
Iatal faanatora
Our focus, however, is on the performance post reduction in iteration count. 14 iterations, BERT pre-trained with a step size of 0.1 cither surpasses or matches the unshared 12-layer baseline across all tasks. However, when the iteration count is further decreased to 12, a performance drop is observed, making it marginally underperform the baseline. Nonetheless, it still achieves over 98% of the bascline performance in most tasks. As for GPT, the model pre-trained with a step size of 0.05 exhibits impressive performance retention on LAM- BADA across all iteration counts, consistently beat- ing the unshared baseline. Although perplexity on Wikitext- 103 rises with reduced iterations, it remains at an acceptable level.
38 30
20
24
Argruge mnrrene Miter 015
Figure 7: Performance of early-exit BERT on MNLI. The black dashed line represents partially- shared BERT pre-trained with 0.1, while the red dashed line denotes the unshared 12-layer BERT.
We have also tried the 24 setting, which is equivalent to the unshared 24-layer model, the performance retention is satisfactory when the scal- ing factors are integers. We place the results and analysis in Appendix G.
4.5 Rapid Inference: Early Exit Integration
This section aims to demonstrate the potential for combining our approach with the early exit tech- nique to further enhance model performance under reduced iteration counts. We adopt the strategy em- ployed by DeeBERT (Xin al., 2020), in which internal classifiers trained at every layer of frozen, fine-tuned BERT model to predict the fi- nal label. Similarily, we train classifiers on the partially-shared BERT that has been pre-trained
It is crucial to underline that this iteration reduc- tion is achieved without additional training post pre-training and fine-tuning. While the partially- shared model may lag behind the unshared baseline in some tasks post-reduction, it stays competitive, which is non- trivial achievement considering that we merely increased the step size during inference.

---

with 0.1. See Appendix F for more details. Initially, we conduct a step scale search for the model at 24, 20, and 16 iterations, as we have done in Section 4.4. Subsequently, classifiers are trained on these models with their respective it- eration counts and searched step scales. During inference, the entropy of the prediction distribu- tion at each layer guide the decision to halt. By manage
or training as separate stage with the PLMs held frozen (Xin et al., 2020; Liu et al., 2020). Our method distinguishes itself by reducing the number of layers during inference without the need for an additional classifier that requires training. Also, our method is complementary to early exit tech- nique, and can be jointly leveraged to accelerate the inference.
adjusting the entropy thresholds, the trade-off between performance and efficiency.
Neural ODEs. The connection between resid- ual networks and ordinary differential equations has been extensively explored in prior research (E, 2017), where different designs of residual networks can be linked diverse numerical discretizations of ODEs (Chen et al., 2018; Lu et al., 2018), Neu- ral ODEs extend the concept of residual networks to continuous-depth architectures. In our work, we build upon the ODE perspective of residual net- works and propose to accelerate the PSPLMs by increasing step size, and from the error analysis of Buler's method, we propose simple pre-training technique enable further inference acceleration.
The results are presented in Fig. 7. Evidently, when the DeeBERT technique is applied to mod- els with reduced iteration counts of 16 and 20, it succeeds in outperforming the unshared 12-layer model under the equivalent computational budget. This finding indicates that the early exit strategy and reduction of iteration count using large step scales are indeed synergistic. Their integration ef- fectively bolsters performance retention, yielding a remarkable inference acceleration.
5 Related Work
PSPLMs. The increasing size and memory us- age of PLMs have prompted research efforts fo- cused on parameter sharing in these models. Sev- eral approaches have been proposed, demonstrating the potential to maintain comparable performance while significantly reducing the model size. Uni- versal Transformers (Dehghani et al., 2019) and ALBERT (Lan et al., 2020) share parameters across all layers. Takase and Kiyono (2021) propose to share parameters for every two consecutive lay- ers or share layer parameters cyclically, while Xue al. (2022) propose to share all layer parameters except bias terms and layer normalization mod- ules. These advanced strategies enhance model capacity at the cost of increased parameter number. However, none of these methods reduce the com- putational cost during inference: the computations required for the inference of shared and unshared PLMs are still identical.
We adopt the linear interpola- tion of a piece-wise linear function as parameters for different layers to build partially-shared PLMs. This bears resemblance to hypernetworks (Ha et al., 2017), where the parameters of a neural network are generated by a specific function or another neural network. The parameterization of model parameters in hypernetwor style has found wide application in various domains, including neural architecture search (Brock et al., 2018), meta-learning (Requeima et al., 2019), and neu- ral ODEs (Chen et al., 2018).
Hyper-Networks.
Conclusion
In this study, we draw inspiration from the ODE perspective residual networks. Our research pro- poses straightforward strategies expedite the in- ference process for both fully and partially-shared PLMs. The results of our work reveal that PSPLMs can not only reduce the storage and memory costs, but also reduce the time costs. Furthermore, when our approach is coupled with the early exit tech- nique, the partially-shared PLMs demonstrate su- perior performance compared to unshared models under the same computational budget. We believe that our methodology harbors substantial poten- tial, particularly in the acceleration of inference in unshared PLMs a promising avenue for future research. We anticipate the extension of our tech-
Early Exit. The efficiency of inference in PLMs has become significant concern for deployment, leading to extensive research efforts focused on in- ference acceleration. Early exit techniques aim to terminate the inference process in early layers and bear close relevance to our work. Many early exit methods necessitate an internal classifier be ap- plied to the intermediate hidden states of the early layers, thus requiring joint training with the PLMs themselves (Zhou et al., 2020; Wang et al., 2022),

---

niques the acceleration of large language models encompassing billions of parameters, and look for- ward to further explorations in this field.
consequences. For example, accelerated PLMs may be used produce fake text more efficiently. potentially contributing to misinformation or fraud.
Acknowledgements
References
This work is supported by the National Key R&D Program of China (No.2022ZD01 16312), Na- tional Natural Science Foundation of China (No. 62236004) and Institute Guo Qiang Tsinghua University.
Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. 2019. Optuna: A next- generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD Interna- tional Conference on Knowledge Discovery & Data Mining, KDD 2019, Anchorage, AK, USA, August 4-8, 2019, pages 2623-2631, ACM.
Author Contribution
the preparation and discussion of the project, Weize Chen, Xiaoyue Xu, Yankai Lin designed the algorithm. Weize Chen and Xiaoyue Xu wrote the code and conducted the experiments. Weize Chen and Xiaoyue Xu wrote the initial draft. Xu Han, Yankai Lin, Ruobing Xie, and Zhiyuan Liu signif- icantly edited and improved the paper. Maosong Sun and Jie Zhou provided valuable advice to the research.
Nikita Balagansky and Daniil Gavrilov. 2022. PAL- BERT: teaching ALBERT to ponder. In NeurIPS.
Andrew Brock, Theodore Lim, James M. Ritchie, and Nick Weston. 2018 SMASH: one-shot model archi- tecture search through hypernetworks. In 6th Inter- national Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenRe- view.net.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Bemner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language models are few-shot learners. In Ad- vances in Neural Information Processing Systems 33: Annual Conference on Neural Information Process- ing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.
Limitations
The effect of a larger step size on model's infer- ence performance could vary across different mod- els and tasks. Although we have tried to include different types of downstream tasks and different types of models to show the generalizability of our method, it could still fail on some certain situations. Morevoer, while we show that our method can be applied to partially-shared PLMs, its effectiveness in accelerating unshared PLMs remains to be fur- ther explored. We have only conducted some basic experiments on the unshared model to show the potential of the method. Further research is needed to determine if and how our method can be adapted for unshared models.
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. 2018. Neural ordinary differen- tial equations. In Advances in Neural Information Processing Systems 31: Annual Conference on Neu- ral Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pages 6572-6583.
Ethics Statement
This work focuses on the acceleration of inference in PSPLMs. While our research does not directly involve human subjects or sensitive data, it does have implications for the broader use of these mod- els in society. The primary potential ethical impact of our work involves the expanded use of PLMS. By providing methods for accelerating PSPLMs, we may enable wider deployment of these models, including in contexts with limited computational resources. While this has potential benefits, such as increased accessibility to advanced language pro- cessing technology, it may also have unintended
Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E, Gonzalez, Ion Stoica, and Eric P. Xing. 2023. Vicuna: An open- source chatbot impressing gpt-4 with 90%* chatgpt quality.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton. Sehastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Tay, Noam Shazeer, Vin- odkumar Prabhakaran, Emily Reif, Nan Du, Ben

---

Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, An- drew M. Dai, Thanumalayan Sankaranarayana Pil- lai, Marie Pellat, Aitor Lewkowyez, Erica Moreira, Rewon Child, Oleksandr Polozov. Katherine Lee, Zongwei Zhou, Xuezhi Wang. Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Hellee-Hllsttem, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. 2022. Palm: Scaling language mod- eling with pathways. CoRR, abs/2204.023111.
ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
Weijie Liu, Peng Zhou, Zhiruo Wang. Zhe Zhao. Haotang Deng, and Qi Ju. 2020. Fastbert: a self- distilling BERT with adaptive inference time, In Proceedings of the 58th Annual Meeting of the As- sociation for Computational Linguisties, ACI. 2020, Online, July 5-10, 2020, pages 6035 6044, Associa- tion for Computational Linguistics.
Ilya Loshchilov and Frank Hutter. 2019. Decoupled weight decay regularization. In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenRe- view.net.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals. Jakob Uszkoreit, and Lukasz Kaiser, 2019. Univer- sal transformers. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans. LA, USA, May 6-9, 2019. OpenReview.net.
Yiping Lu, Aoxiao Zhong. Quanzheng Li, and Bin Dong. 2018. Beyond finite layer neural networks: Bridging deep architectures and numerical differential equa- tions. In Proceedings of the 35th International Con- ference on Machine Learning, ICML 2018, Stock- holmsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Re- PMLR.
Jacob Devlin, Ming-Wir Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: pre-training of deep bidirectional transformers for language under- standing. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistiex; Human Language Tech- nologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pages 4171-4188 Association for Computational Linguistics.
search, pages
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2017. Pointer sentinel mixture mod- els. In Sth International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. Open- Review.net.
Weinan E. 2017. A proposal on machine learning via dynamical systems. Communications Mathemat- ics and Statistics, 5(1):1-11.
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary. Vijay Kor- thikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro. Amar Phanishayee. and Matei Zaharia. 2021. Efficient large-scale language model training on GPU clusters using megatron-Im. In International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2021, St. Louis, Missouri, USA, Novem- ber 14-19, 2021, page 58. ACM.
Leo Gao, Stella Biderman, Sid Black, Laurence Gold- ing. Travis Hoppe. Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. 2021. The pile: An 800gb dataset of diverse text for language modeling. CoRR, abs/2101.000271
David Ha, Andrew M. Dai, and Quoc V. Le. 2017. Hypernetworks. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Pro- ceedings. OpenReview.net.
Denis Paperno, Germán Kruszewski, Angeliki Lazari- dou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. 2016. The LAMBADA dataset: Word prediction requiring broad discourse context. In Proceedings of the 54th Annual Meeting of the As- sociation for Computational Linguistics, ACI. 2016, August 7-12, 2016, Berlin, Germany, Volume 1: Long Papers. The Association for Computer Linguistics.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang. and Eduard H. Hovy. 2017. RACE: large-scale read- ing comprehension dataset from examinations. In Proceedings of the 2017 Conference Empirical Methods in Natural Language Processing, EMNLP 2017, Copenhagen, Denmark, September 9-11, 2017, pages 785-794, Association for Computational Lin- guistics.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. 2019. Language models are unsupervised multitask learners.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. 2020. ALBERT: A lite BERT for self-supervised learning of language representations. In 8th Inter- national Conference on Learning Representations,
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020. Exploring the limits of transfer learning with a unified text-to-text trans- former. J. Mach. Learn. Res., 21:140:1-2 1-406677

---

Pranav Rajpurkar, Jian Zhang. Konstantin Lopyrev, and Percy Liang. 2016. Squad: 100, 000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas, USA, November 1-4, 2016, pages 2383-2392. The Association for Computational Linguisties.
Fuzhao Xue, Ziji Shi, Futao Wei, Yuxuan Lou, Yong Liu, and Yang You. 2022. Go wider instead of deeper. In Thirty-Sixth AAAT Conference gence, AAAI 2022, Thirty- Fourth Conference on In- novative Applications of Artificial Intelligence, IAAI 2022, The Twelveth Symposiun on Educational Ad- vances in Artificial Intelligence, EAAT 2022 Virtual Event, February 22 March 1, 2022, pages 8779- 8787. AAAI Press.
Artificial Inselli-
James Requeima, Jonathan Gordon, John Bronskill, Sebastian Nowozin, and Richard E. Turner 2019. Fast and flexible multi-task classification using con- ditional neural adaptive processes. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Sys- tems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 7957-7968.
Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian J. McAuley, Ke Xu, and Furu Wei. 2020. BERT loses patience: Fast and robust inference with early exit. In Advances in Neural Information Processing Sys- tems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler. 2022. Confident adaptive language modeling. In NeurlPS.
A Proof on Pre-Norm Transformer Architecture as Residual Network
To demonstrate that the pre-norm Transformer ar- chitecture can be considered a residual network, we will analyze the computation within each layer. Let us consider the hidden state h,, and examine the computation in the t-th layer of the pre-norm transformer, which can be represented as follows:
Sho Takase and Shun Kiyono. 2021. Lessons on pa- rameter sharing across layers in transformers. CoRR, abs/2104.06022.
Rohan Taori, Ishaan Guirajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. 2023. Stanford alpaca: An instruction Instruction-follloinng following Ilama model. https:// gi thub. cww/ http://ttsulaa//aae wat.-laabbssaaaaaaae sttaffrrddaaal __d_aapaca..
3g ATT, h 11 FFN, (LN₂ (2:)) + +zt
hi,
(8)
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023. Llama: Open and efficient foundation language models. CoRR. abs/2302.13971.
By substituting Eq. 7 into Eq. 8, we obtain:
hr+1 h, (LN...))(....) + (TT. =h+ f(.h)).
Hence, we observe that the pre-norm Transformer architecture functions as a residual network, where each layer parameterizes the derivatives of the hid- den states and updates the hidden states utilizing Buler's ' method.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. 2019. GLUE: multi-task benchmark and analysis plat- form for natural language understanding. In 7th In- ternational Conference on Learning Representations. ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net.
B Proof on the Errors of Euler's Method
Here we provide the proof in the case where the dimension of state is for simplicity. It can be easily generalized to high-dimensional setting. At time to. we denote the real state of the ODE as h((o), The Buler's method (or residual network) approximates the real state at f = to + + aas
Jue Wang, Ke Chen, Gang Chen, Lidan Shou, and Ju- lian J. McAuley. 2022. Skipbert: Efficient inference with shallow layer skipping. In Proceedings of the 60th Annual Meeting of the Association for Compu- tational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27. 2022, pages 7287- 7301. Association for Computational Linguistics.
he -h(lo) + 8- (()(
(9)
Ji Xin, Raphael Tang, Jaejun Lee, Yaoliang Yu, and Jimmy Lin. 2020. Deebert: Dynamic early exiting for accelerating BERT inference. In Proceedings of the 58th Annual Meeting of the Association for Com- putational Linguistics, ACL 2020, Online, July 5-10, 2020. pages 2246-225 Association for Computa- tional Linguistics.
here we move the notation of time / from the sub- script (Eq. 3) into the parentheses to explicitly de- note the dependency on time
The real state h((1) can be expanded by using the Taylor series expansion around to:

---

From Eq. 14, we have [Tyl S s By substitut- ing into Eq. 17 and applying the absolute value inequality, we can see that
$² 2
h(4)) = l((t)) sh (1o) + ((ot)),
(10)
where ig is some number between to and to + 8,
|et41l S let, + +s f(h(ts), (3) - f(he,, 4:)l
and h' htot)) = f(h(to),to)
M
(18)
2
Local Truncation Error. Now derive the lo- cal truncation error T of Euler's method, which is essentially the discrepancy between the real state h(tt) and the one-step approximated state hy start from hto)). It is formally written as:
Because the assumption of f and its derivative be- ing continuous and bounded, according to the mean value theorem, we have
(()(t)).))) ((h,,,))))) (hi
(19)
To = h(t)) (h(to) + to)(
where hi, is some number between htss) and hag Since fn. is bounded, we have
5²
2 ((o)).
(11)
We can further derive the (4o) by differentiating W (to) f(h(to); to):
[f(h(ta),t.) f(he, ti)l S Rleal
(20)
where R is some constant. By substituting Eq. 20 into Eq. 18, we obtain the relation between the errors of two consecutive steps:
+ fi(h(to), to) I (to) = (i((t(t)), to)h to) (o
+ fc(h(to),, to). (12)
Jet4il I +R)))11 + M
$² 2
(21)
Assume that f and its derivatives (h(tt), to) and fe(htt)), to) are all continuous and bounded, then there exists a constant M so that
For simplicity, let (1 + +R). We can then iterative apply the inequality starting from t = to and eto = h(to) - hto as
|h" (to) S M.
(13)
M
Taking Eq. 13 into Eq. 11, then we can establish the local truncation error of Euler's method:
le, s 2 8
M
lcra
S (1 + C)
T M
2
(14)
Since M is constant, the local truncation error of Euler's method is of order Oss2), ), i.c., the square of the step size.
M
le S (1+ (1+C+ 1m-
(22)
2
and since (1 + C + C [1+ R-)"=>"" Is we then obtain the error bound the approximated final result
1-C"
1-
Global Truncation Error. The global truncation error is the accumulated error from initial time to to final time T. To derive the global truncation error, we first define e,, = h((tt)= - fr, as the difference between the real state h(44) and approximated state hr, at time f.. Recall that we have
(1+ Rs)
1 M
le,,1 5
8
R
2
C Ilns 1M
R
l(((+1) h(tt) + 8- f(h(t)),4 + Tu,, h444 hu + th),,,,),
(15)
RT
(16)
(ns T). (23)
R.
2
where Eq. 16 is the Euler's method, and Eq. 15 is the Buler's method with local truncation error term. Substracting Eg.. 16 from Eq. 15, we have
RT -1 M
Denoting K =
R 2 then the global trunca-
tion error is
le,, S Ks,
(24)
e: Ee + +((()())) 4)) f(b,, 44) + T.. (17)
which of order 0(s), i.e., linear to the step size.

---

Hyper-Parameter
Value
Hyper-Paree Pareme-
Value
Ir
lc-4
Ir
le44
Ir decay style
linear
Ir decay style
cosine
min Ir
le-5
min Ir iterations
le-5
iterations
300,000
300,000
batch size
1024
batch size
512
weight decay
0.01
weight decay
0.01
warmup ratio
0.01
warmup ratio
0.01
gradient norm dropout
1.0
gradient norm dropout
1.0
0.0
0.1
Table 2: param--rrsesers for pre-training all the BERT models in this work.
Table 3: Hyper- parameters for pre-training all the GPT- models this work.
As for the difference between the final approx- imated state hy obtained using larger step size and the state by obtained using the original step size s, since lerl = |u - h(T) < Ks and ēel = 1b - h(T) < K,,s, therefore we have 1Pr herl = (Iur - -M(T) (Per - A(T))
A (Loshchilov and Hutter, 2019) as the opti- mizer. The model configurations for our BERT and GPT-2 are kept the same as BERT large and GPT- large respectively.
E Configurations for the Downstream Tasks
hvt - h(TT) + - h(T) (25)
< ((1 + B)s.
(26)
We adopt the approach employed by Megatron-LM framework for handling MNLI and RACE tasks. For the classification tasks MNLI and SST-2, we utilize the hidden state of the [CLS] token for clas- sification and report accuracy on the development set. In the RACE task, we predict the probability of each answer using the [CLS] token's 's representa- tion and report test set accuracies. Regarding the SQuAD v1.1 and v2.0 tasks, we adhere to BERT's training procedure, applying span extraction loss, and record the F1 score on the development set using the official evaluation script?.
Therefore the difference between ht and hr is also bounded, and linear to step size 8. When the step size at different time is different, one can casily derive that the global truncated error is bounded by the largest step size 8*
[ My <
(27)
Additional Figures for Section 4.3 and Section 4.4
Due to the length constraint, place the addi- tional plots for Section 4.3 Figs. 8 and 10, and plots for Section 4.4 at Fig. 9.
During fine-tuning BERT on all the downstream tasks, we use the linear learning rate warmup and decay schedule. The gradient norm constrained to 1.0. We apply a dropout rate of 0.1 and weight decay of 0.01. Further hyperparameter details are documented in Table 4.
Configurations for the Pre-training
We use Mcgatron-LM (Narayanan et al., 2021) the framework to pre-train our parameter- shared BERT and GPT-2. The OpenWetText dataset for pre-training GPT-2 is prepared fol- lowing the instructions in Megatron-LM repos- itory', and the preparation of the Pile dataset for pre-training BERT follows the instructions in the Megatron--Deepspeed repository?. We use
For GPT tasks, we adopt a zero-shot approach. The performance on the LAMBADA task is as- sessed using cloze accuracy, which involves pre- dicting the last word (not the last token) based the preceding tokens. Performance on Wikitext- 103 is measured using the perplexity metric on the test set.
httl::///ggguuub..c///////////// 'https://gi http///tthu...oommiirrs///t//// Megatron-DeepSpeed

---

1.0
0.0
0.4
3.4
3.
1.0
0.4
Layor
Layar Too:
(a) SST-2
(b) SQuAD
(c) SQuAD2.0
Figure The cosine similarity between the derivatives given by the model consecutive iterations.
Step Sux
Nart
M
21.0
0:23
93.5
M9.75
23.03
62.35
90.33
Dep diss
78
5.05
92.53
54
Petroturs netieaee
Paiter ene Inbeeste
SST-2
(b) SQuAD
SQuAD2.0
Figure 9: The inference performance of partially-shan shared models with n = 12 sets parameters on SST-2, SQuAD and SQuAD2.0. The red dashed line represents the performance of the unshared 12 layer model.
Configurations for the Early Exit
responding results are presented in Table 5.
On close examination, we note peculiar trend: as the iteration count reduces from 24 to 12, the zero-shot perplexity the Wikitext-100 103 dataset first increases, and then decreases. This anomaly could be attributed the model's inability learn the usage of linear interpolation between the end- point parameters for derivative calculation when 24. That is, when n = 24, the linear interpola- tion (Eq. 6) always returns the parameters on one of the endpoint. Consequently. during the inference phase, as we increase the step size, the parameters derived through linear interpolation start deviating from the parameters utilized by the model during training. This divergence is potentially responsible for a significant degradation in the model's perfor- mance.
We employ the methodology from DeeBERT (Xin et al., 2020) for the early exit experiment. During the training phase, we keep the BERT model and the original classification head fixed and only train the additional classifiers each layer. The training loss is computed as the sum of the losses from all additional classifiers.
In the inference phase, we calculate the entropy the output logits at each layer. If the entropy value falls below a predetermined threshold at a layer, we halt the computation and take the predic- tion at this layer as the final prediction. We record the average number of iterations at the point of out- put across all instances. The thresholds utilised in this experiment are follows: [0, 0.01, 0.05, 0.07, 0.1, 0.2, 0.3, 0.4, 0.5]. Higher thresholds induce an earlier exit.
Interestingly, when the scaling factor 3 for each iteration is adjusted to 2 and the iteration count is reduced to 12, the model yields a reasonable per- formance. We hypothesize this is due to the fact that when scaling factors integers, the parame- ters derived remain endpoint parameters, which the model has been accustomed to handle during the training phase. However, the experiment still high- light the potential of our method when applying to the unshared model.
G Performance under 24 Setting
In this section, extend partially-shared model include 21 sets of parameters, which renders it equivalent an unshared 24-layer model. The only difference is that we use the same initial- ization for the 24 sets of parameters. This attempt is only made on the GPT-2prge model, and the cor-

---

Syype-Parramete
MNLI
SST-2
RACE
SQuAD
SQuAD2
Step Size
Ir
2c-05
lc-05
le-15
3c-05
2c-05
epochs
5
5
3
5
batch size
64
32
16
32
32
warmup ratio
0.05
0.01
0.1
0.01
0.01
Step Size 0.1
Ir
lc-05
le-05
2c-15
3c-15
20-05
epochs
batch
16
warmup ratio
0.1
0.05
0.1
0.1
Step Size 0.05
Ir
lc-05
3c-05
2c-15
3c-05
2c-05
epochs
5
5
3
5
batch size
16
32
32
32
16
warmup ratio
0.1
0.05
0.1
0.01
0.01
Step Size 0.01
epochs 4c-05
Ir
2e-05
3c-05
2c-05
3c-05
batch
64
16
16
16
32
warmup ratio
0.05
0.01
0.05
Table 4: ppremeearsettrs for downstream tasks with different step sizes.
#Iters
PPL
24
19.74
20
34.67
16
168.68
14
77.21
12
57.93
Table 5: Zero-shot perplexity Wikitext- 103 of the partially-shan model with n 24.

---

10
E.01
3.03
66
64
12 Layer Absolute Difference
16 Layer Absolute Difference
20 Layer Absolute Difference
121
101
8802
50u40
LAMACC
2073
LAMAAA
S20422 LAMIAGA
Layer Relative Difference
16 Layer Relative Difference
(f) Layer Relative Difference
Figure 10: The absolute and relative difference between the final hidden states obtained with 24 iterations and 12, 16, 20 iterations.


Reviews:
[{"review_id": "6uhMlLHavg", "reviewer": "Reviewer_k26S", "summary": "", "questions": "Question A: How are the beta values tuned for each task ? Is it on the dev set ? \nQuestion B: For the results in Table 1, it seems like encoder only models are more robust to reduced iterations compared to zero shot performance from decoder only models. Is there any intuition on why that might be the case ?\nQuestion C: Is there an intuition on why models trained with a smaller step-size outperform their size=1 counterparts ? \nQuestion D: From Figure 10 (b), the absolute drop between h_{T} and \\tilde{h}_{T} between Race and SQuAD is about the same and quite less. Based on the hypothesis presented in the paper, this should imply that the downstream performance between the two should be quite similar too. However, the downstream performance is quite different for the case of RACE, with the 16 iteration model outperforming the 24 iteration model. Is there an intuition on why that might be the case ?\n", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate your comprehensive review and insightful feedback on our work. We are gratified by your positive assessment and will address each of your concerns in sequence.\n\n---\n\n**Question 1**: Tuning of beta values (scaling factor of the step size) and use of the dev set. \n\n**Answer 1**: We concur with your viewpoint on the importance of specifying the process. **For tasks with a dev set, such as RACE, we tune beta values using dev set and report the performance on test set. For GLUE datasets, we follow conventions of prior works, tuning on the dev set and reporting the dev result.** We utilized the Optuna toolkit (line 303), running 50 trials for each task. It is really fast since only one round of inference on the dev set is required in each trial, and it can also be easily parallelized. For example, when using one GPU, the search only takes around 14 minutes for SST-2. In our response to reviewer y9YL, we also show that we can still achieve good results without tuning the beta value and setting it directly to $\\beta = 24 / \\text{reduced-iteration}$. **We have shown the potential of varying the step size at each step in our paper, and we believe that continuing to explore how to quickly determine the beta values would be an important and practical direction.**\n\n---\n\n**Question 2**: Intuition on why encoder-only models are more robust than decoder-only models?\n\n**Answer 2**: **It's possible the perceived fragility of decoder-only models isn't intrinsic but rather a result of the evaluation metrics.** On wikitext-103, we used the perplexity metric which can degrade exponentially with diminishing model capability. The larger classification space in text generation (the entire vocabulary) versus encoder tasks might also explain the pronounced decline in LAMBADA accuracy.\n\n---\n\n**Question 3**: Intuition on why models trained with a smaller step size outperform their size=1 counterparts?\n\n**Answer 3**: We hypothesize that smaller step sizes produce smaller gradients, which might mitigate issues like large gradient variance during training. This could act in a manner reminiscent of gradient norm clipping. Your observation is astute and warrants further investigation.\n\n---\n\n**Question 4**: Is there an intuitive explanation for the difference in downstream performance between the 16 iteration and 24 iteration models for RACE, despite the similar absolute drop observed between $h_{T}$ and $\\tilde{h}_{T}$ in Figure 10 (b)?\n\n**Answer4**: We can truly feel your meticulousness in reviewing our paper from this question, thank you! **Comparing absolute drops in $h_{T}$ and $\\tilde{h}_{T}$ across distinct tasks might not provide meaningful insights due to different loss landscapes** (or accuracy landscapes if it's a classification task). For example, even minor deviations from optimum values can cause sharp performance declines in steeper landscapes, whereas in flatter landscapes, the same deviations might result in milder effects. The absolute difference only makes sense for the same task.\n\n---\n\nRegarding the suggestions mentioned in the reasons to reject, such as the performance and speed of the unshared 24-layer model, we will incorporate them into our paper for a more comprehensive presentation. And we will correct all the typos that you have patiently pointed out. Again, thank you for your comprehensive review and feedback! We hope our clarifications have addressed your concerns.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. Leveraging the connections of Pre-LayerNorm transformer networks to ODEs for reducing inference latency by dynamically scaling step sizes is, to the best of my knowledge, novel and of practical utility.\n2. The inference time step-size scaling results are quite compelling, especially for encoder only models\n3. The low step-size training experiments are theoretically motivated, and the paper presents strong experiments validating the claims (Section 4.3.2)\n4. The proposed method for partially shared PLMs is very interesting with compelling results, and demonstrates the possibility of additional future explorations.\n5. The integration with early exit strategies further demonstrates the utility of the presented techniques.", "reasons_to_reject": "1. Some results need additional context to better understand the implications, in my opinion. Specifically, for the section on partial sharing, in addition to Figure 6, it would also be good to \n1.1 Present the upper-bound performance using a 24 layer unshared model\n1.2 Present the iteration speed, as done in Table 1. \nSince the paper is mostly presented as a method for improving speed, an iso-parameter unshared model with lower performance somewhat deviates from the main message (since in this case, as per my understanding, the iso-parameter model would technically have lower latency). While line 495 partially addresses this (w.r.t comparing the performance of a post reduction model comparing to baseline), I think the above two data-points are valuable to get a full picture.\n\n2. The details about how the beta values are chosen per iteration budget are somewhat fuzzy. Specifically, what is the set that the beta values are tuned on ? For the language model setup, where metrics are zero-shot, are the beta values tuned at a per task level, or at an aggregate level ? This is important since for decoder only language models that operate on a zero-shot setup, one of the primary advantages is the task agnostic nature of the models. Thus it would be good to see how robust the approach is to task agnostic tuning (eg: tuning on perplexity on the validation set of the pre-training objective), and how that might translate to zero-shot setups. On a similar note, it would be good to understand how robust the approach is to different sets of beta values.\n\n3. Appendix G presents an analysis based on Table 5. However at a perplexity of >= 35, the model performance is really bad. So any conclusions / trends based on that might not be very robust. \n\n", "questions_to_authors": "Question A: How are the beta values tuned for each task ? Is it on the dev set ? \nQuestion B: For the results in Table 1, it seems like encoder only models are more robust to reduced iterations compared to zero shot performance from decoder only models. Is there any intuition on why that might be the case ?\nQuestion C: Is there an intuition on why models trained with a smaller step-size outperform their size=1 counterparts ? \nQuestion D: From Figure 10 (b), the absolute drop between h_{T} and \\tilde{h}_{T} between Race and SQuAD is about the same and quite less. Based on the hypothesis presented in the paper, this should imply that the downstream performance between the two should be quite similar too. However, the downstream performance is quite different for the case of RACE, with the 16 iteration model outperforming the 24 iteration model. Is there an intuition on why that might be the case ?\n", "paper_topic_and_main_contributions": "1. The paper presents a novel way of improving inference speed for parameter shared pre-trained language models (PSPLMs), by drawing a connection between Euler's method for solving ODEs and the per layer output of the PSPLMs. Specifically, the authors propose increasing the step size, thereby reducing the number of iterations as a way of trading off accuracy of estimation of the final layer output with speed. \n2. The authors show that for both encoder style models models (BERT) fine-tuned on downstream tasks, as well as decoder only models (GPT-2) on zero-shot task specific metrics, their proposed method achieves minimal loss in performance while improving inference cost ranging from 1.20x to 1.92x.\n3. The authors also upper-bound the approximation error based on largest scaling factor used across layers and step-size. This guides the intuition for the experiments demonstrating pre-training with smaller step sizes would potentially yield lower approximation errors at fewer iteration counts.\n4. The authors also present a piece-wise linear approximation for partially shared parameter based PLMs, and demonstrate it's effectiveness compared to models of a similar parameter count.\n5. The authors finally also demonstrate the effectiveness of their approach when combined with other latency reduction techniques like early exit methods.\n", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "1. It would be good to bring at least a part of Figure 10 (L422) into the main paper.\n2. The axis on Figure 4 seem to be incorrect (for RACE, SST-2 and MNLI)"}, "extra_text": {}}, {"review_id": "Q5hmEYgN0I", "reviewer": "Reviewer_T3zw", "summary": "", "questions": "A: What is the best practive for determining the step size and number of iterations? While it may vary depending on the model or data, even suggesting heuristic methods to find appropriate values would be valuable.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for the thorough review of our paper! We value the feedback provided and have made efforts to address your concerns. Here are our responses:\n\n---\n\n**Question 1**: Best practice for determining the step size and number of iterations? \n\n**Answer 1**: **The number of iterations is determined based on actual needs when using the model for inference**. The user of the model decides on the number of iterations based on the complexity of the problem, the resources they have, and their tolerance for errors. Fewer iterations lead to faster inference but potentially worse results. **As for the step size during pre-training, this is a hyperparameter, and it's generally challenging to systematically determine the optimal value of a hyperparameter.** Our empirical results suggest step sizes between 0.5 to 0.05 work well. For the step size scaling factor $\\beta$ during inference, the reported values in our paper are derived using Optuna toolkit searches on each task's dev set (line 303). We searched for 50 trials. **It is really fast since only one round of inference on the dev set is required in each trial, and it can be easily parallelized.** For example, when using one GPU, the search only takes around 14 minutes for SST-2. Additionally, as shown in our response to reviewer y9YL, we give evidence that setting the beta value directly to $\\beta = 24 / \\text{reduced-iteration}$ also yields satisfactory results.\n\n---\n\n**Question 2**: Effectiveness of this method across scales.\n\n**Answer 2**: Due to rebuttal time constraints, we don't have enough resources to train a larger model. As an alternative, we train two parameter-shared BERT base and reported their performance (without searching $\\beta$) on MNLI as follows:\n\n| #Iters | MNLI (step size=0.1) | MNLI (step size=0.05) |\n| ------ | -------------------- | --------------------- |\n| 12     | 81.7               | 81.7                |\n| 10     | 81.9 (+0.2\\%)    | 82.0 (+0.4\\%)     |\n| 8      | 81.9 (+0.2\\%)    | 82.1 (+0.5\\%)     |\n| 6      | 79.8 (-2.3\\%)    | 79.8 (-2.3\\%)     |\n\nAfter reducing the number of iterations, we directly set $\\beta$ to $12/\\text{reduced-iter}$. As can be seen, increasing the step size after reducing the number of iterations is also effective for a small model at the BERT base level . As for larger models, we believe our approach will still be applicable. Although we haven't tested our method on billion-parameter-level models, we believe our approach has been validated in practical models like BERT and has practical value, making it worthy of publication. This also paves a potential new way for accelerating model inference in the future.\n\n---\n\nWe hope our responses address your concern. Thank you again for your patient review!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- There is definite merit in reducing the inference cost without the need for additional modules.\n- The paper holds value in the presence of several insightful empirical results.", "reasons_to_reject": "- It is disappointing that there is no method provided to find appropriate values for important variables such as step size or number of iterations, leaving heavy reliance on heuristics.\n- Additionally, it would be beneficial to include an analysis of the model's scale. Given the recent attention towards LLMs that demonstrate diverse generalization abilities, showing the potential effectiveness of this methodology across various scales and applicability in LLMs would enhance the significance of the research.\n- As the authors also acknowledge in the limitations section, further validation is required for Partially-shared PLMs.", "questions_to_authors": "A: What is the best practive for determining the step size and number of iterations? While it may vary depending on the model or data, even suggesting heuristic methods to find appropriate values would be valuable.", "paper_topic_and_main_contributions": "- The paper presents a neural ordinary differential equation (ODE) inspired technique aimed at enhancing the inference efficiency of parameter-shared PLMs (PSPLM).\n- Furthermore, a novel pre-training method, involving a slightly altered forward propagation, is proposed to further accelerate the inference of PSPLM.\n- Additionally, the study showcases the potential of the proposed strategy in the context of partially-shared PLM, highlighting the flexibility and robustness of the approach.\n- Overall, the findings contribute to the advancement of PSPLM research and underscore the efficiency gains achievable through the proposed techniques.", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "PxEdoqDan6", "reviewer": "Reviewer_y9YL", "summary": "", "questions": "s=0.05 was consistently better than s=0.1 in the fully-shared setting but s=0.1 outperformed s=0.05 in most datasets in the partially shared setting. Do the authors have any intuition on this discrepancy?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough and constructive review of our paper. We sincerely appreciate the time and effort you've put into pointing out the weaknesses. Below, we will respond to each of your comments.\n\n---\n\n**Question 1**: The results of fixed $\\beta$ (the step-size scaling factor) are not reported.\n\n**Answer 1**: We initially aimed to demonstrate the maximum potential of acceleration via our method by exploring the best $\\beta$. **In fact, the search is very fast, for example, searching for the $\\beta$ value we reported on SST-2 on a single GPU takes only about 14 minutes.** Nonetheless, you rightly pointed out the importance of reporting results with a fixed $\\beta$. As such, we now present results using a direct setting of $\\beta=24/\\text{reduced-iteration}$ for various iteration counts:\n\n`fully-shared step size=1.0`\n\n| #Iters | MNLI               | SST-2             | RACE               | SQuAD              | SQuAD2             | Wiki-103             | LAMBADA            |\n| ------ | ------------------ | ----------------- | ------------------ | ------------------ | ------------------ | -------------------- | ------------------ |\n| 24     | 83.6             | 91.1            | 64.3             | 90.2             | 81.5             | 33.0               | 31.1             |\n| 20     | 83.2 (-0.5\\%)  | 90.6 (-0.5\\%) | 63.5 (-1.2\\%)  | 88.8 (-1.6\\%)  | 75.7 (-7.1\\%)  | 38.3 (+16.1\\%)   | 16.1 (-48.2\\%) |\n| 16     | 81.0 (-3.1\\%)  | 89.5 (-1.8\\%) | 60.1 (-6.5\\%)  | 82.6 (-8.4\\%)  | 57.2 (-29.8\\%) | 64.7 (+96.1\\%)   | 6.4 (-79.4\\%)  |\n| 12     | 71.9 (-14.0\\%) | 86.6 (-4.9\\%) | 52.3 (-18.7\\%) | 61.7 (-31.6\\%) | 33.7 (-58.7\\%) | 168.5 (+410.6\\%) | 1.6 (-94.9\\%)  |\n\n`fully-shared step size=0.1`\n\n| #Iters | MNLI               | SST-2              | RACE               | SQuAD              | SQuAD2             | Wiki-103           | LAMBADA            |\n| ------ | ------------------ | ------------------ | ------------------ | ------------------ | ------------------ | ------------------ | ------------------ |\n| 24     | 83.2             | 91.4             | 67.5             | 90.1             | 81.4             | 33.4             | 25.5             |\n| 20     | 83.6 (+0.5\\%)  | 91.6 (+0.2\\%)  | 67.3 (-0.3\\%)  | 89.5 (-0.7\\%)  | 79.8 (-2.0\\%)  | 34.9 (+4.5\\%)  | 23.5 (-7.8\\%)  |\n| 16     | 82.6 (-0.7\\%)  | 89.6 (-2.0\\%)  | 61.3 (-9.2\\%)  | 83.5 (-7.3\\%)  | 66.9 (-17.8\\%) | 35.4 (+6.0\\%)  | 23.7 (-7.1\\%)  |\n| 12     | 52.2 (-37.3\\%) | 79.4 (-13.1\\%) | 44.1 (-34.7\\%) | 33.5 (-62.8\\%) | 36.9 (-54.7\\%) | 45.4 (+35.9\\%) | 18.1 (-29.0\\%) |\n\n`fully-shared step size=0.05`\n\n| #Iters | MNLI              | SST-2             | RACE              | SQuAD             | SQuAD2            | Wiki-103           | LAMBADA            |\n| ------ | ----------------- | ----------------- | ----------------- | ----------------- | ----------------- | ------------------ | ------------------ |\n| 24     | 82.6            | 90.9            | 68.1            | 90.3            | 81.3            | 33.9             | 26.5             |\n| 20     | 82.8 (+0.2\\%) | 91.2 (+0.3\\%) | 68.3 (+0.3\\%) | 90.2 (-0.1\\%) | 81.3 (+0.0\\%) | 35.2 (+3.8\\%)  | 27.4 (+3.4\\%)  |\n| 16     | 83.2 (+0.7\\%) | 91.1 (+0.2\\%) | 68.2 (+0.1\\%) | 90.0 (-0.3\\%) | 80.8 (-0.6\\%) | 35.5 (+4.7\\%)  | 27.5 (+3.8\\%)  |\n| 12     | 81.7 (-1.1\\%) | 89.5 (-1.5\\%) | 62.2 (-8.7\\%) | 85.5 (-5.3\\%) | 74.5 (-8.4\\%) | 45.6 (+34.5\\%) | 21.9 (-17.4\\%) |\n\n`fully-shared step size=0.01`\n\n| #Iters | MNLI              | SST-2             | RACE              | SQuAD             | SQuAD2            | Wiki-103           | LAMBADA           |\n| ------ | ----------------- | ----------------- | ----------------- | ----------------- | ----------------- | ------------------ | ----------------- |\n| 24     | 81.7            | 90.1            | 64.4            | 89.3            | 79.5            | 35.0             | 27.1            |\n| 20     | 81.8 (+0.1\\%) | 90.5 (+0.4\\%) | 64.7 (+0.5\\%) | 89.3 (+0.0\\%) | 79.6 (+0.1\\%) | 35.9 (+2.6\\%)  | 27.0 (-0.4\\%) |\n| 16     | 82.0 (+0.4\\%) | 90.3 (+0.2\\%) | 64.0 (-0.6\\%) | 89.2 (-0.1\\%) | 79.5 (+0.0\\%) | 36.0 (+2.9\\%)  | 27.8 (+2.6\\%) |\n| 12     | 82.0 (+0.4\\%) | 89.7 (-0.4\\%) | 60.8 (-5.6\\%) | 83.2 (-6.8\\%) | 73.1 (-8.1\\%) | 45.1 (+28.9\\%) | 25.6 (-5.5\\%) |\n\nFor the fully-shared model, the results indicate that directly setting $beta$ to $24/\\text{reduced-iteration}$ can produce favorable outcomes without searching. However, for models with a step size of 0.1, this approach seems to be not so effective. **Nonetheless, the significance of the zero-shot acceleration achieved in models with step sizes of 1, 0.05, and 0.01 should not be overlooked. We anticipate that these insights might pave the way for more in-depth research in model inference acceleration.** Also, with a simple and quick search, we can yield much better $\\beta$ values as reported in our paper. These observations are consistent for the fully-shared model as well:\n\n`fully-shared step size=0.1`\n\n| #Iters | MNLI              | SST-2             | RACE               | SQuAD              | SQuAD2             | Wiki-103            | LAMBADA            |\n| ------ | ----------------- | ----------------- | ------------------ | ------------------ | ------------------ | ------------------- | ------------------ |\n| 24     | 85.7            | 93.1            | 71.7             | 91.2             | 82.0             | 21.3              | 39.0             |\n| 20     | 84.9 (-0.9\\%) | 93.2 (+0.1\\%) | 70.5 (-1.7\\%)  | 90.6 (-0.7\\%)  | 80.2 (-2.2\\%)  | 31.1 (+46.0\\%)  | 31.3 (-19.7\\%) |\n| 16     | 84.7 (-1.2\\%) | 92.2 (-1.0\\%) | 69.1 (-3.6\\%)  | 89.7 (-1.6\\%)  | 78.2 (-4.6\\%)  | 24.8 (+16.4\\%)  | 27.9 (-28.5\\%) |\n| 12     | 80.7 (-5.8\\%) | 90.6 (-2.7\\%) | 59.8 (-16.6\\%) | 78.3 (-14.1\\%) | 67.5 (-17.7\\%) | 47.4 (+122.5\\%) | 28.0 (-28.2\\%) |\n\n`fully-shared step size=0.05` \n\n| #Iters | MNLI              | SST-2             | RACE               | SQuAD              | SQuAD2             | Wiki-103            | LAMBADA            |\n| ------ | ----------------- | ----------------- | ------------------ | ------------------ | ------------------ | ------------------- | ------------------ |\n| 24     | 85.2            | 93.3            | 71.4             | 91.4             | 82.6             | 22.1              | 39.2             |\n| 20     | 84.7 (-0.6\\%) | 92.9 (-0.4\\%) | 68.4 (-4.2\\%)  | 90.1 (-1.4\\%)  | 80.7 (-2.3\\%)  | 35.6 (+61.1\\%)  | 35.8 (-8.7\\%)  |\n| 16     | 84.4 (-0.9\\%) | 92.7 (-0.6\\%) | 68.2 (-4.5\\%)  | 89.5 (-2.1\\%)  | 76.2 (-7.7\\%)  | 26.5 (+19.9\\%)  | 38.0 (-3.1\\%)  |\n| 12     | 78.7 (-7.6\\%) | 87.5 (-6.2\\%) | 54.9 (-23.1\\%) | 73.6 (-19.5\\%) | 63.8 (-22.8\\%) | 52.9 (+139.4\\%) | 25.3 (-35.5\\%) |\n\nFor the fully-shared model, directly setting the beta value yields slightly inferior results, but at 16 layers, the performance remains quite good. **We believe our findings, particularly the zero-shot acceleration displayed in our experiments, will stimulate further interest and research in the realm of model inference acceleration.**\n\n---\n\n**Question 2**: Lack of comparison with early exit techniques.\n\n**Answer 2**: We discuss the relationship between our method and early exit techniques in Section 4.5. **Our approach is orthogonal to early exit techniques, complementing each other.** Our experiments have proven that when they are combined, even superior inference acceleration can be attained (Figure 7). Considering our method's zero-shot nature, it can be viewed as an enhancement to pre-existing early exit techniques. Moreover, **since early exit techniques require training on the training dataset and introduce additional new parameters, directly comparing our method with them might not be fair.**\n\n---\n\n**Question 3**: Notational inconsistencies and paper readability issues.\n\n**Answer 3**: Thank you for your meticulous feedback regarding the inconsistencies and errors. We are committed to resolving all the highlighted issues to offer a clearer and more reader-friendly paper. Regarding the issue where we referenced figures from the appendix in the main text, we will adjust the layout to move essential content to the main body. Thank you once again for your patient correction!\n\n---\n\n**Question 4**: In the fully-shared model, $s=0.05$ performs better, but in the partially-shared model, $s=0.1$ is better. Is there any explanation?\n\n**Answer 4**: The observed discrepancy indeed presents an intriguing aspect of our experiments. A potential hypothesis could be that the fully-shared model, with each iteration influencing the same parameters, has a larger gradient norm than the partially-shared model. As such, using a smaller step size might be analogous to gradient norm clipping, thereby reducing the gradient variation. However, this remains a hypothesis, and we are actively exploring more insights into this phenomenon.\n\n---\n\nThe proposed revisions, guided by your invaluable feedback, will significantly enhance the paper. We're grateful for your insights and look forward to furthering this discussion.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "This paper addresses an important problem, i.e., the efficiency of Transformers. The proposed method, which applies weights to sub-layers during the inference, is an interesting approach to accelerates the inference speed.", "reasons_to_reject": "This paper is inadequate to publish due to some reasons.\n\nIn my understanding, the claims of this paper have not been proved yet. The authors varied beta during the experiments but did not report the results in fixed beta such as beta=1. Thus, it is difficult to agree with the claim \"we can accelerate the inference speed by using appropriate beta.\" In addition, I recommend reporting the portability of beta. Can we use beta which is adjusted to a specific task for other tasks?\n\nIn addition, the authors should compare the proposed method with the early exit techniques. These method might need extra modules and/or computational costs (especially, during the training) but the proposed method also need additional computational costs to search the appropriate beta before the inference. The proposed method is more efficient in the total cost?\n\nThis paper is hard to follow. At least, the authors should reduce the ambiguity in notations. For example, in Equation (3), f_t -> f_i in my understanding. In addition, the authors defined the number of layers as T in Section 2 but use L in Section 3.3. Moreover, the description contains several errors. For example, in Figure 4, vertical axes of (a), (b), and (c) represent perplexity but the correct representation is Accuracy. Moreover, the authors mentioned figures in Appendices for the main discussion such as Section 4.3.2. The authors should include them into the main part if the authors use them to support the main claims.", "questions_to_authors": "s=0.05 was consistently better than s=0.1 in the fully-shared setting but s=0.1 outperformed s=0.05 in most datasets in the partially shared setting. Do the authors have any intuition on this discrepancy?", "paper_topic_and_main_contributions": "In these days, the large-scale pre-trained language models such as GPT series have succeeded in various NLP applications. In this paradigm, to save the computational costs, we have to explore an efficient method. As a parameter-efficient method, the parameter sharing approach is widely used. In this paper, the authors propose a method to accelerate inference speed of the parameter sharing Transformers. The authors indicated that we can accelerate the inference speed of Transformers whose parameters are shared with all layers by modifying weights to each sub-layers such as FFNs and attention modules. The authors focused on BERT and GPT-2, and conducted experiments on several tasks.", "excitement": 3, "reproducibility": 2, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents methods for improving the inference speed for parameter shared pre-trained language models, presenting both theoretical and empirical evidence. The reviewers appreciate the importance of the problem, the novelty, elegance, practicality and theoretical foundations, the richness of the results, and their quality. In contrast, there were initially claims about missing experiments and comparisons, which resulted in unjustified claims, and clarity issues. Post-rebuttal, many of these claims seem to be resolved (as evident in the internal discussion amongst the reviewers), but the clarity issues are still very much concerning to some of them. The scaling concerns also still remain, though I have discounted them due to the ACL policy (https://www.aclweb.org/portal/content/efficient-nlp-policy-document) which discourages unjustified requests for larger experiments.
Decision: Accept-Findings