P ru MUX : Augmenting Data Multiplexing with Model Compression (2023)

Abstract: As language models increase in size by the day, methods for efficient inference are critical to leveraging their capabilities for various applications. Prior work has investigated techniques like model pruning, knowledge distillation, and data multiplexing to increase model throughput without sacrificing accuracy. In this paper, we combine two such methods – structured pruning and data multiplexing – to compound the speedup gains obtained by either method. Our approach, PruMUX, obtains up to 7.5-29.5X throughput improvement over BERT-base model with accuracy threshold from 80% to 74%. We further study various combinations of parameters (such as sparsity and multiplexing factor) in the two techniques to provide a comprehensive analysis of the tradeoff between accuracy and throughput in the resulting models. We then propose Auto-PruMUX, a meta-level model that can predict the high-performance parameters for pruning and multiplexing given a desired accuracy loss budget, providing a practical method to leverage the combination effectively.

## introduction large language models (llms) have achieved state-of-the-art performance across various nlp tasks and resulted in impressive user-facing demonstrations such as chatgpt. 2 however, their large size necessitates the use of enormous amounts of compute and memory at inference time, which limits their widespread use. two types of techniques have been explored to reduce the cost of model inference. the first is model compression including network pruning @xcite @xcite , quantization @xcite , knowledge distillation @xcite , combinations of multiple methods @xcite . the figure 1 : throughput improvements (×) of cofi, data-mux, and prumux over the bert-base model @xcite on the mnli task @xcite . the sparsity for a cofi's data point is labeled as s. the width of multiplexing for a datamux's data point is labeled as n . the parameter pair for a prumux's data point is labeled as (n , s). while both types of methods leverage the overparameterization effect @xcite in modern deep neural networks to improve the throughput-to-compute cost ratio, the manner in which they do so is different. model compression aims at reducing the number of parameters in the model, hence reducing the overall compute cost (denominator) to improve the ratio. data multiplexing, on the other hand, compresses multiple inputs into one to improve throughput (numerator) while keeping the model size fixed. this observation naturally leads us to hypothesize that the two types of methods could be complementary and can be combined for maximal gain in the throughput-to-compute cost ratio. there are two challenges to this hypothesis. the first is that both model compression and data multiplexing aim at trading a small accuracy loss for large throughput improvement. intuitively, the combination may incur an accuracy loss larger than either method and it is not clear how they interact with each other when combining them together. a research question is how to combine the two methods such that the combination achieves better throughput than each type of method individually, given any accuracy loss budget or accuracy threshold. the second challenge is to efficiently find the best parameters pair (n, s) where n is the width of the data multiplexing and s is the sparsity of the model compression method. training and testing with each parameter combination is costly and time-consuming. a research question is how to automatically predict and find top parameters based on the model's performance on one set of parameters. to address the first research question, we present prumux, a combination of model compression and data multiplexing. our method is simple and consists of three phases -multiplexed model pretraining, task-specific fine-tuning and task-specific model compression. in our implementation, we make use of cofi @xcite , a state-ofthe-art model compression method that includes intermediate knowledge distillation steps that help minimize accuracy hits, and datamux @xcite , which performs vector-based input multiplexing over instances. our results over four datasets (mnli, qnli, qqp and sst-2) demonstrate that prumux achieves significantly higher throughput over cofi and datamux individually for a large range of accuracy thresholds. as an example, figure 1 shows the throughput improvements over the bert-base model on task mnli, providing a more optimal pareto frontier in the tradeoff between accuracy and throughput. to address the second research question, we propose auto-prumux, a meta-model to automatically predict and find the high-performance parameter combinations for a desired accuracy loss budget on a task based on the model's performance on one set of parameters without running additional experiments. we use interpolation and estimation models over a set of data points to predict the accuracy and throughput of a prumux model based on sparsity and multiplexing factor. we show promise in modeling the tradeoffs accurately and auto-prumux can find high-performance combinations of known parameters as well as unknown parameters, pro-viding a practical method for choosing a highperformance prumux model for a downstream task. our key insight for why prumux can achieve better throughput than model compression and data multiplexing individually is that they improve the throughput of a model in two different dimensions: reducing the latency of an inference and compressing multiple inferences. in addition, both methods lead to non-linear drops in model accuracy at some points. prumux can achieve high throughput while avoiding each method's limitations. ## auto-prumux to address the question above, we propose auto-prumux, a method to search for top (n, s) parameters, to help practitioners balance the performance vs throughput trade-off. our research question is: suppose we have some experimental data of prumux and the experimental data of datamux and cofi, how can we find and predict the top parameters (n, s) given an accuracy loss budget? our approach is to develop performance models for the accuracy and throughput of prumux. we first train prumux models for a set of (n, s) combinations and measure both the accuracy and the throughput improvement. we then use this data to fit a throughput model and an accuracy model to predict throughput and accuracy respectively given (n, s) parameters. we first discuss how we fit the accuracy and throughput models with a set of sparse data points. given that we are working with a limited set of data points, we opt to use a simple class of interpolation models for modeling prumux accuracy and use an estimation model for modeling throughput. we then outline how we leverage these models to predict top (n, s) parameters, given an accuracy loss budget. we then demonstrate the effectiveness of the auto-prumux in predicting the top parameters across a wide range of accuracy loss budgets. ## conclusion we propose prumux, a method to combine model compression and data multiplexing to build high throughput transformers. our implementation of prumux makes use of cofi and datamux and we show that it achieves substantial throughput improvement over either cofi or datamux for a large range of accuracy thresholds. we conclude that the reason that prumux performs well in certain range of accuracy loss budgets is that cofi and datamux improve the throughput of a model in two different dimensions: reducing the latency of an inference and compressing multiple inferences. when the accuracy loss budget is large, both methods lead to non-linear drops in model accuracy, prumux can achieve much better performance than either approach because it uses more conservative parameters for cofi and data-mux before each reaches its bad trade-off point. we also present auto-prumux, a meta-model to automatically predict high-performance parameter combinations for a desired accuracy on a task. we show it is promising in predicting parameters without individual data points and additional training.