R ev MUX : Data Multiplexing with Reversible Adapters for Efficient LLM Batch Inference (2024)

Abstract: Large language models (LLMs) have brought a great breakthrough to the natural language processing (NLP) community, while leading the challenge of handling concurrent customer queries due to their high throughput demands. Data multiplexing addresses this by merging multiple inputs into a single composite input, allowing more efficient inference through a shared forward pass. However, as distinguishing individuals from a composite input is challenging, conventional methods typically require training the entire backbone, yet still suffer from performance degradation. In this paper, we introduce RevMUX, a parameter-efficient data multiplexing framework that incorporates a reversible design in the multiplexer, which can be reused by the demultiplexer to perform reverse operations and restore individual samples for classification. Extensive experiments on four datasets and three types of LLM backbones demonstrate the effectiveness of RevMUX for enhancing LLM inference efficiency while retaining a satisfactory classification performance.

## introduction in recent years, large language models (llms), such as gpt-3 (175b) @xcite , palm (540b) @xcite , and gpt-4 (1.7t) @xcite have emerged as a cornerstone in natural language processing @xcite . the field has witnessed a dramatic increase in model sizes, which, although improving downstream performance, also poses considerable challenges. inference with these llms has become increasingly resource-intensive, often confronting users with capacity limits @xcite . with the rise of "language model as a service" @xcite , improving inference efficiency has become a key focus to accommodate these growing model sizes. to explore efficient inference for llms, the community has mainly focused on model-centric or algorithm-centric approaches @xcite . model-centric approaches, including quantization @xcite and knowledge distillation @xcite , aim to compress llms into smaller models while retaining the capabilities of the vanilla models. in contrast, algorithm-centric approaches, such as speculative decoding @xcite and kv-cache optimization @xcite , aim to reduce latency and memory usage in sequence generation tasks. however, when processing batch queries in a single forward pass, these methods generally result in a significant increase in computational load, e.g., flops, linear with the number of inputs. the multi-input multi-output (mimo) architecture @xcite @xcite has emerged as an effective approach for predicting multiple samples simultaneously in a single forward pass, incurring the computational cost of only a single input. this architecture requires jointly training a multiplexer and a demultiplexer alongside the entire base model: the multiplexer combines multiple inputs into one, and the demultiplexer separates the base model's outputs back into individual ones. subsequent research @xcite applied mimo-style training to enhance inference with llms, exemplified by bert @xcite . however, this method not only necessitates training the multiplexer and demultiplexer during pre-training phase but also requires fine-tuning the entire model, including bert's parameters, thereby limiting its applicability to increasingly larger language models. moreover, updating the backbone model's parameters necessitates maintaining multiple copies of the backbone model to accommodate dynamic inference budgets, further constraining its scalability. in this paper, we explore mimo training on a fixed llm to improve its inference efficiency with-large language model out additional pre-training. a major challenge in implementing mimo for fixed llms is to trace and preserve the uniqueness of each input, as the fixed llms can struggle to differentiate individual instances within the consolidated inputs, resulting in performance degradation @xcite . inspired by reversible neural networks @xcite , which split the input into two halves for parallel processing and enable reconstruction of lower-layer inputs from upper-layer outputs, we propose reversible adapters to achieve data multiplexing, dubbed revmux. the reversible multiplexer leans to map different samples into distinct feature spaces and the mapping function is shared with the demultiplexer to perform a reverse operation that dis-aggregates the output from the llm for classification. we train these reversible adapters in a parameter-efficient manner @xcite on downstream tasks and then apply them for batch inference. through extensive experiments on four datasets and three types of llm backbones, we demonstrate the effectiveness of revmux in enhancing llm inference efficiency while maintaining satisfactory classification performance. notably, our revmux method, which freezes the entire backbone llm, achieves performance comparable to or surpassing that of two fine-tuned baselines on bert base . we also extended our method to encoder-decoder architectures such as t5 and decoder-only architectures like llama3-8b @xcite . results on all three architectures across five scales consistently show that the proposed reversible adapters significantly contribute to performance retention during data multiplexing. moreover, the combined use of reversible adapters in both the multiplexing and demultiplexing processes creates a synergistic effect, amplifying performance benefits beyond those achieved by individual components. ## conclusion in this paper, we introduce revmux, a parameterefficient data multiplexing framework designed to enhance the batch inference efficiency of llms. revmux features a reversible multiplexer that combines multiple samples, allowing the demultiplexer to reverse this process and restore individual outputs for classification. we train revmux on downstream tasks while keeping the backbone llm frozen, and apply it for batch inference. extensive experiments on bert-base, t5 across three scales, and llama3-8b demonstrate the effectiveness of revmux in enhancing both accuracy and efficiency. ablation studies confirm the synergistic function of the reversible multiplexer and the reverse demultiplexer.