Predicting Rewards Alongside Tokens: Non-disruptive Parameter Insertion for Efficient Inference Intervention in Large Language Model (2024)

Abstract: Transformer-based large language models (LLMs) exhibit limitations such as generating unsafe responses, unreliable reasoning, etc. Existing inference intervention approaches attempt to mitigate these issues by finetuning additional models to produce calibration signals (such as rewards) that guide the LLM’s decoding process. However, this solution introduces substantial time and space overhead due to the separate models required. This work proposes Non-disruptive parameters insertion (Otter), inserting extra parameters into the transformer architecture to predict calibration signals along with the original LLM output. Otter offers state-of-the-art performance on multiple demanding tasks while saving up to 86.5% extra space and 98.5% extra time. Furthermore, Otter seamlessly integrates with existing inference engines, requiring only a one-line code change, and the original model response remains accessible after the parameter insertion.

## introduction transformer-based large language models (llm) have demonstrated remarkable prowess in a wide range of natural language processing tasks with human-like proficiency, including text generation @xcite , translation @xcite , summarization @xcite @xcite , etc. however, extensive inspections from multiple directions have revealed llms are not flawless. for instance, llm occasionally produces unsafe or toxic outputs @xcite , and in reasoning capabilities, yielding unreliable results for mathematical proofs @xcite ). * work done during internship at alibaba group † corresponding author otter-1b llm-7b calibration model-7b i read through the project document. that idea is... llm-7b i read through the project document. that idea is perfect otter-llm-8b (online deployable) non-disruptive parameter insertion inference intervention i read through the project document. that idea is ... i read through the project document. that idea is perfect calibrate 😊 😊 0.78 terrible 0.19 bad ... calibrate 😞 0.78 terrible 0.19 bad ... 😞 0.92 perfect 0.05 good ... i read through the project document. that idea is terrible 0.92 perfect 0.05 good ... to address the above limitations of llms, a series of approaches have been proposed, which can be broadly categorized into two directionsfinetuning and inference intervention. finetuning methods leverage techniques such as knowledge graphs @xcite and data augmentation @xcite to curate vast amounts of high-quality data to align llms with preferred responses. nevertheless, prior research has highlighted that continual pretraining or finetuning can lead to unintended consequences, such as catastrophic forgetting @xcite @xcite , and hallucination @xcite . these side effects can severely undermine the overall performance and reliability of llms. to avoid associated side effects, plug-and-play inference intervention as an alternative has been proposed. the method employs a calibration model to produce signals (such as rewards), which are used to calibrate the decoding tokens during llm inference. when the intervention is not required, the calibration model can be unplugged, allowing the llm to return to its original output. benefiting from this, inference intervention achieves reliable reasoning, such as process-supervised reward models in mathematical reasoning @xcite , reduces bias and harmful responses with rewardguided search @xcite . therefore, inference intervention paves a promising way to enhance llms' capabilities. regretfully, the success of inference intervention heavily relies on training additional calibration models, which are typically large-scale models, and causes the space complexity to increase drastically @xcite . in addition, as shown in fig. 1 , the inference time is generally longer because the original model's output is re-used as input for the calibration model, increasing the times of transformer forward passes @xcite . this introduces substantial space and time overhead, making it challenging to deploy these methods online. motivated by the powerful knowledge inherent in llms, training an entirely new model is not necessary. instead, it is reasonable to utilize this knowledge by re-using the original parameters. specifically, our objective is to introduce a small set of additional parameters into the original llm, allowing it to simultaneously generate accompanying output (either reward prediction or language modeling) alongside its primary output. this accompanying output should behave the same as the additional models in inference intervention methods. introducing additional parameters into llms has been indeed widely used in parameter-efficient finetuning/adapting methods (peft), such as qlora and lora @xcite . these methods work by adding adapted parameters to the original frozen parameters, i.e., @xmath0, where only the added ∆w is trainable. however, by directly adding parameters to the original parameters, lora alters all model outputs. consequently, lora is not suitable for inference intervention methods, where the original model outputs (i.e., logits) and a calibration signal are required simultaneously. in this work, we propose a non-disruptive parameters insertion (otter) method for inference intervention in llms. the key idea behind otter is to concatenate the added trainable parameters across all components of the transformer, including the multi-head attention layer and the feed-forward neural network layer. by passing through all layers, the final extended hidden state can be mapped and utilized as an inference intervention signal. compared with existing work, this study offers several contributions: otter saves up to 86.5% extra space and 98.5% extra time while obtaining comparable performance to state-of-the-art inference intervention methods on three high-demanding tasks: generation detoxification, preference alignment, and inference acceleration. ## method we introduce otter for predicting calibration signals alongside tokens in transformer-based llms without disrupting their original output and knowledge. our method involves inserting new trainable parameters into the feed-forward network (ffn) and multi-head attention (mha) layers of the transformer architecture. specifically, for i-th block of the transformer, we expand the original hidden state h i into @xmath1, h ′ i ], thereby making it possible to predict the inference intervention signals based on the last layer h ′ n . ## experiments to comprehensively assess otter's capabilities, we consider three tasks: human preference alignment @xcite , detoxification @xcite , and inference speed-up @xcite . the evaluation involves various language models, including llama @xcite , vicuna @xcite , and gpt-2 @xcite . in addition to task-specific evaluation metrics, we analyze the computational overhead in terms of space and time complexity. we define the space overhead as the ratio of gpu memory consumed by the modified model to that of the original model during inference. the time overhead is defined as the ratio of the time consumed by the modified model to the time consumed by the base model. all inference interventions are performed on a single a100 gpu with the huggingface transformers library @xcite . in every task, we introduce an ablation of otter, denoted as task head only, which solely adds a trainable linear layer on the top of the frozen llm. this linear layer is then fine-tuned to predict the intervention signals used in each individual task. ## conclusion otter concatenates a small set of trainable parameters to the transformer architecture, enabling additional output for inference intervention while preserving the original model's behavior and computational efficiency. experiments demonstrate otter's efficacy in aligning outputs with human preferences, mitigating toxicity in text generation, and improving inference speed-achieving these enhancements with significantly less overhead than existing methods. as an extension of the same model architecture, otter can seamlessly integrate with existing infrastructure, providing an appealing solution for enhancing llm performance efficiently for the community.