Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Chenhan Yuan, Fei Huang, Ru Peng, Keming Lu, Bowen Yu, Chang Zhou, Jingren Zhou

TL;DR

Otter is proposed, inserting extra parameters into the transformer architecture to predict calibration signals along with the original LLM output, offering state-of-the-art performance on multiple demanding tasks while saving up to 86.5% extra space and 98.5% extra time.

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. Our code is publicly available at \url{https://github.com/chenhan97/Otter}

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

TL;DR

Otter is proposed, inserting extra parameters into the transformer architecture to predict calibration signals along with the original LLM output, offering state-of-the-art performance on multiple demanding tasks while saving up to 86.5% extra space and 98.5% extra time.

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. Our code is publicly available at \url{https://github.com/chenhan97/Otter}
Paper Structure (40 sections, 11 equations, 3 figures, 9 tables)

This paper contains 40 sections, 11 equations, 3 figures, 9 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Inference Intervention for LLM
  4. Parameter-Efficient Finetuning for LLM
  5. Method
  6. FFN Layer Adaption
  7. Multi-head Attention Layer Adaption
  8. Layer Norm Regularization
  9. Experiments
  10. Reward-guided Search for Helpful and Harmless Alignment
  11. Setup
  12. Results
  13. Controlled Bi-Experts Generation for Reducing Toxicity
  14. Setup
  15. Results
  16. ...and 25 more sections

Figures (3)

  • Figure 1: Comparison of inference intervention methods with and without Otter for harmless response generation. By inserting parameters into the frozen LLM, Otter significantly reduces space and time costs, while enabling seamless online deployment.
  • Figure 2: The Otter architecture. Grey denotes frozen parameters while blue is trainable.
  • Figure 3: The comparisons of initialization methods' effectiveness on speculative decoding and preference alignment. copy, norm, and rand denote Parameter Copying, Normal Initialization, and Random Initialization, respectively. Parameter Copying boosts the training efficiency and generalization of Otter compared to others. The loss in the preference alignment task reflects the training of the reward model. Higher average reward values in this task indicate better alignment. In speculative decoding, the Top-1 medusa head Acc. measures the average next token prediction accuracy of the first medusa decoding head during training. Higher accuracy corresponds to a higher acceleration ratio. The speed-up ratio in this task quantifies the acceleration achieved against the base model, evaluated at each training epoch.