GenARM: Reward Guided Generation with Autoregressive Reward Model for Test-time Alignment (2410.08193v4)

Abstract: LLMs exhibit impressive capabilities but require careful alignment with human preferences. Traditional training-time methods finetune LLMs using human preference datasets but incur significant training costs and require repeated training to handle diverse user preferences. Test-time alignment methods address this by using reward models (RMs) to guide frozen LLMs without retraining. However, existing test-time approaches rely on trajectory-level RMs which are designed to evaluate complete responses, making them unsuitable for autoregressive text generation that requires computing next-token rewards from partial responses. To address this, we introduce GenARM, a test-time alignment approach that leverages the Autoregressive Reward Model--a novel reward parametrization designed to predict next-token rewards for efficient and effective autoregressive generation. Theoretically, we demonstrate that this parametrization can provably guide frozen LLMs toward any distribution achievable by traditional RMs within the KL-regularized reinforcement learning framework. Experimental results show that GenARM significantly outperforms prior test-time alignment baselines and matches the performance of training-time methods. Additionally, GenARM enables efficient weak-to-strong guidance, aligning larger LLMs with smaller RMs without the high costs of training larger models. Furthermore, GenARM supports multi-objective alignment, allowing real-time trade-offs between preference dimensions and catering to diverse user preferences without retraining. Our project page is available at: https://genarm.github.io.

• The paper introduces an autoregressive reward model (ARM) that guides token-level decisions for real-time text generation without retraining.
• It integrates ARM with frozen LLMs under a KL-regularized framework, achieving efficient next-token sampling and cost-effective weak-to-strong guidance.
• Experimental results demonstrate that GenARM rivals training-time methods, enabling effective multi-objective alignment across varied user preferences.

Analyzing GenARM: A Novel Approach for Test-Time Alignment in LLMs

The paper presents "GenARM," an innovative test-time alignment method leveraging the Autoregressive Reward Model (ARM) to improve the adaptability and efficiency of LLMs without requiring repeated retraining. This approach addresses challenges associated with aligning LLMs to human preferences, a crucial task given the computational cost and inefficiency of traditional training-time methods like Reinforcement Learning from Human Feedback (RLHF).

Key Contributions and Methodology

1. Autoregressive Reward Model (ARM): The paper introduces a unique ARM that directly predicts next-token rewards, which is crucial for autoregressive text generation. This contrasts with existing trajectory-level Reward Models (RMs) which evaluate completed responses and are ill-suited for token-level decisions required in real-time text generation.
2. Theoretical Guarantee: The ARM's parameterization allows it to guide frozen LLMs under a KL-regularized reinforcement learning framework, ensuring that its expressiveness aligns with distributions achievable by traditional RMs.
3. GenARM Methodology: GenARM utilizes the ARM to integrate token-level reward predictions with the logits of a frozen LLM during text generation. This integration permits efficient next-token sampling and reduces computational inefficiency by forgoing full trajectory evaluations, unlike some alternative methods.
4. Superior Performance: Empirically, GenARM outperforms existing test-time alignment techniques and rivals training-time efficiency—specifically matching the performance of the DPO method. It distinguishes itself by providing a feasible pathway for weak-to-strong guidance, where smaller ARM models guide much larger LLMs, which is particularly cost-effective given the high resource demand for training larger models.
5. Multi-Objective Alignment: GenARM supports multi-objective alignment by accommodating various user preference dimensions, allowing real-time trade-offs without necessitating retraining. This flexibility is demonstrated through its application to multiple dimensions, such as helpfulness and harmlessness, by using multiple ARMs calibrated for different objectives.

Experimental Results

Several experiments illustrate the efficacy and efficiency of GenARM:

• Performance Comparison: Against several baseline methods, GenARM substantially reduced inference time while maintaining high performance, showcasing its capability to handle diverse preferences efficiently.
• Weak-to-Strong Generalization: GenARM successfully utilized a smaller ARM trained on a 7B parameter scale to guide larger models like 70B, achieving significant alignment without direct training of the larger models.
• Efficient Multi-Objective Trade-Offs: By dynamically adjusting reward weights, GenARM enables efficient tailoring of LLM outputs to varied user preferences, as evidenced in its handling of both helpfulness and harmlessness dimensions.

Implications and Future Directions

The implications of GenARM's methodology are significant for scaling the alignment of large-scale AI models. By bypassing the need for re-training through its adaptable test-time alignment, GenARM can potentially pave the way for more scalable and cost-effective model adaptations. Future research could extend GenARM's architecture to other complex tasks such as coding and reasoning, and further refine the ARM's predictive capabilities across diverse LLM architectures and languages.

In summary, GenARM represents a substantial contribution to efficient LLM alignment, setting a pioneering benchmark for adapting large models to complex human values dynamically and efficiently. By facilitating weak-to-strong guidance and multi-objective consideration, it offers a robust approach to managing the growing computational demands in the evolving AI landscape.