Decoding-time Realignment of Language Models (2402.02992v2)

Abstract: Aligning LLMs with human preferences is crucial for reducing errors and biases in these models. Alignment techniques, such as reinforcement learning from human feedback (RLHF), are typically cast as optimizing a tradeoff between human preference rewards and a proximity regularization term that encourages staying close to the unaligned model. Selecting an appropriate level of regularization is critical: insufficient regularization can lead to reduced model capabilities due to reward hacking, whereas excessive regularization hinders alignment. Traditional methods for finding the optimal regularization level require retraining multiple models with varying regularization strengths. This process, however, is resource-intensive, especially for large models. To address this challenge, we propose decoding-time realignment (DeRa), a simple method to explore and evaluate different regularization strengths in aligned models without retraining. DeRa enables control over the degree of alignment, allowing users to smoothly transition between unaligned and aligned models. It also enhances the efficiency of hyperparameter tuning by enabling the identification of effective regularization strengths using a validation dataset.

• The paper introduces Decoding-time Realignment (DeRa), a novel method to adjust the tradeoff between reward maximization and regularization at decoding time without retraining.
• The paper demonstrates that models aligned under different regularization strengths are geometric mixtures of a reference model and a single aligned model, enabling efficient exploration of alignment tradeoffs.
• The paper validates DeRa across tasks such as summarization, hallucination mitigation, and chatbot performance, highlighting significant improvements with minimal computational overhead.

Decoding-time Realignment of LLMs

The paper introduces an innovative approach to LLM (LM) alignment called Decoding-time Realignment (DeRa), which offers an efficient method for managing the tradeoff between reward maximization and regularization strength during LM optimization. This tradeoff is crucial because it impacts how well a model can maintain alignment with human preferences without suffering from reward hacking or performance regression due to excessive or insufficient regularization.

Key Contributions

1. Geometric Mixtures and Model Alignment: The authors demonstrate that models aligned under different regularization strengths are geometric mixtures of a reference model (such as an SFT model) and a single aligned model. This insight allows for the exploration of alignment without needing to retrain multiple models under different hyperparameter settings.
2. Decoding-time Realignment (DeRa): The proposed DeRa method allows users to explore different regularization strengths at decoding time rather than during model training. This addresses a significant inefficiency in traditional methods, which often require retraining large models to determine optimal regularization levels, saving computational resources.
3. Practical Implementation: The authors propose an autoregressive approximation that efficiently computes these geometric mixtures at the token level during the decoding process. This is achieved by linearly combining the logits of the reference and aligned models, modulated by a user-specified parameter, $\lambda$. This flexibility enables fine-grained control of the tradeoff between alignment and regularization strength on-the-fly.

Experimental Evaluation

The paper evaluates DeRa across several settings:

• Toy Example with Length Reward: Using a controlled task where the reward is tied to the generated text length, DeRa shows consistent performance improvements similar to those achieved by retrained models, validating its utility as an approximation of full retraining.
• Summarization Task: On the Reddit TL;DR dataset, DeRa is shown to effectively identify optimal KL regularization strengths, comparable to outcomes achieved through the standard retraining approach. The experiments indicate that while the base regularization level might be too strong, DeRa was able to identify better configurations.
• Hallucination Mitigation: In alignment tasks such as hallucination reduction in retrieval augmented generation, DeRa effectively tuned alignment strengths to balance task performance with hallucination control.
• Chatbots and Real-world Applications: Applying DeRa to Zephyr-7b models demonstrated its capacity to enhance general-purpose chat models, enabling performance adjustments for various downstream tasks such as open-domain conversation.

Theoretical and Practical Implications

DeRa provides a significant advancement in adaptive model deployment, especially in scenarios requiring rapid, responsive adjustments to alignment parameters without extensive computational overhead. The model's inherent flexibility suggests its potential applicability not just across a spectrum of tasks and model architectures, but also in dynamic environments where deployment conditions change frequently.

From a theoretical perspective, DeRa's foundation on the mixture of probabilistic distributions offers an elegant solution that incorporates established principles of model optimization while circumventing traditional computational limitations. This approach could have significant influences on future research directions concerning real-time model adaptability and continuous learning.

Conclusion

The Decoding-time Realignment approach represents a crucial step forward in efficient model alignment and optimization, reducing computational inefficiencies and offering nuanced control over model behaviors at deployment. This positions DeRa as a compelling tool in both academic research and industry applications, where adaptive and responsive models are increasingly in demand. Furthermore, by addressing the longstanding challenges associated with reward-regularity tradeoffs in LLMs, this work opens avenues for more effective and efficient model quality tuning processes.