Boosting Inference Efficiency: Unleashing the Power of Parameter-Shared Pre-trained Language Models (2023)

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.

## introduction in recent years, there has been a significant increase in the number of parameters of plms. this began with the advent of bert @xcite , containing 340 million parameters, and has escalated to models like t5 @xcite , @xcite , and palm @xcite , with the latter reaching an astounding 540 billion parameters. the trend of plm expansion has, undeniably, improved performance across numerous tasks. nonetheless, the corresponding increase in computation and storage requirements has raised substantial barriers for scenarios characterized by stringent latency requirements or resource limitations. while plms encompassing merely a * equal contribution † corresponding author. in contrast to the monumental advances in plms, the real-world applications often still favor more established models such as @xcite . these models, despite their relatively fewer parameters, deliver satisfactory performance across many tasks while requiring significantly less resources. this balance offers an appealing trade-off between performance and cost. moreover, parameter sharing techniques have successfully demonstrated that model size can be greatly reduced without significant performance degradation, mitigating the storage burden and yielding better cost-effectiveness. this has sparked an interest in parameter-shared plms (psplms) like al-bert @xcite , a derivative of the bert architecture that shares parameters across all layers, 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 it resolves the storage issue, it does not address the computational challenge. early exit techniques promise to reduce the number of layers processed during inference by halting computation at early layers @xcite @xcite . while effective, 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 it be complementary to early exit techniques, allowing their combined use for further acceleration? in this study, we show that the problem can be well addressed in psplms. specifically, we illustrate how significant acceleration in psplm inference can be achieved through our straightforward yet effective technique. this technique, inspired by the principles of neural odes, accelerate the inference without necessitating the addition of modules or calculations to the model. hence, in addition to the inherent storage efficiency, our method notably makes psplms to be computational efficient. we also introduce a 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 acceleration during inference, and we give theoretical explanation to aptly support our method. we further extend the application of our method beyond the domain of fully shared psplms. our research demonstrates the potential of our acceleration 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, our method is not in competition with other acceleration strategies. we demonstrate that our method can be combined orthogonally with early exit techniques, thereby facilitating further acceleration. remarkably, this synergy of methods makes the partially-shared model surpass its unshared equivalent within an equivalent computational budget. in essence, our work offers a novel route to accelerate the inference for psplms, and lay a novel foundation for unleashing the psplms' potential, offering critical insights for deployment of psplms in resource-constrained settings. ## ode perspective on residual networks we begin by providing a brief overview of the relationship between residual networks and odes, which forms the fundamental basis of our research. in a t -layer residual network, we denote the layer t as f θt . the update formulation for the hidden state h can be expressed as: remarkably, this update scheme aligns with euler's method for solving odes. consider an ode: euler's method approximates the solution of this ode from @xmath0, each with step size s i , such that n-1 @xmath1. the method iteratively computes the following formula from @xmath2: the final value y n serves as an approximation of the solution to eq. 2 at time t . the correspondence between eq. 1 and eq. 3 is evident. a t -layer residual network can be interpreted as parameterizing the vector field f that characterizes the derivative along the path from the input space to the final output space. the ode perspective generalizes the concept of depth in residual networks to the continuous domain, where the notion of progression from input to output is captured by the continuous time rather than discrete depth or layer index. during the inference process of a trained model, the vector field remains fixed because the parameters are frozen. as a result, the model's inference can be seen as solving an ode within this vector field using euler's method, where the initial value corresponds to the input embedding, and the solution time is t . furthermore, the pre-norm transformer architecture is also a 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. ## conclusion in this study, we draw inspiration from the ode perspective on residual networks. our research proposes straightforward strategies to expedite the inference 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 technique, the partially-shared plms demonstrate superior performance compared to unshared models under the same computational budget. we believe that our methodology harbors substantial potential, particularly in the acceleration of inference in unshared plms -a promising avenue for future research. we anticipate the extension of our tech-niques to the acceleration of large language models encompassing billions of parameters, and look forward to further explorations in this field. local truncation error. now we derive the local truncation error t of euler's method, which is essentially the discrepancy between the real state h(t 1 ) and the one-step approximated state h t 1 start from h(t 0 ). it is formally written as: assume that f and its derivatives f h (h(t 0 ), t 0 ) and f t (h(t 0 ), t 0 ) are all continuous and bounded, then there exists a constant m so that taking eq. 13 into eq. 11, then we can establish the local truncation error of euler's method: since m is a constant, the local truncation error of euler's method is of order o(s 2 ), i.e., the square of the step size.