GenARM: Reward Guided Generation with Autoregressive Reward Model for Test-Time Alignment

Authors: Yuancheng Xu, Udari Madhushani Sehwag, Alec Koppel, Sicheng Zhu, Bang An, Furong Huang, Sumitra Ganesh
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: large language model alignment; controlled decoding;

Abstract:
Large Language Models (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.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1729646459041

thread_type:
official_review

<thread_object>
creator: reviewer_s14e

<content>
summary:
This paper presents GenARM, a novel reward model parameterization for test-time alignment. Comparing with prior works in the domain of test-time alignment, this paper proposes to use the sum of auto-regressive log probability of each token to be the reward model for the entire trajectory. The training objective function and usage for constrained decoding remains the same. Theoretical analysis has been provided to justify this parameterization of the reward model. Empirically, strong performances have been observed for GenARM against other popular baselines such as ARGS and Transfer-Q over three different settings including human preference alignment, weak-to-strong guidance, and multi-objective alignment.

strengths:
The clarity of the writing is worth mentioning, I enjoyed following the reading. Each component is properly motivated. And the insights from the experiments are properly organized.

The approach is, as far as I am aware of, novel. It also seems easy to train and does not incur too much additional computational overhead.

Theoretical analysis has been provided for completeness of the paper.

The effectiveness of the proposed method is validated thoroughly in the experiments section over three different settings including human preference alignment, weak-to-strong guidance, and multi-objective alignment. The section of weak-to-strong guidance is particularly informative.

weaknesses:
The training objective of the reward model (Equation [5]) draws an interesting connection with DPO objective (Equation [7] in DPO). This might result in an unfair comparison with other test-time alighment methods such as ARGS because now that the training cost of the reward model is basically the same as training-time alignment methods like DPO. 

In order to provide a more fair comparison, it might be helpful to dedicate additional experiments/paragraphs illustrating the number of samples/training time comparison with other baselines such as ARGS and Transfer-Q

There is an inconsistency of baselines across different settings. In particular, why is BoN of included in Section 6.1 and Transfer-Q not included in Section 6.2?

questions:
See weakness, it would be great if the authors can explain the connection of the proposed method with DPO objective
</content>
</thread_object>

</thread>
<thread>
creation date:
1729902778005

thread_type:
official_review

<thread_object>
creator: reviewer_dm7h

<content>
summary:
Test time alignment, which guide frozen LLMs during test time text generation, can align LLMs with preference without incurring training cost.
To address the mismatch between trajectory-level RMs and autoregressive text generation, this paper proposes GenARM, a test-time alignment approach that leverages the Autoregressive Reward Model to predict next-token rewards for efficient and effective autoregressive generation.
Experimental results and theoretical arguments are provided to justify the proposed approach, especially compared to the prior test-time alignment baselines.

strengths:
1. The idea of addressing the mismatch between trajectory-level RMs and autoregressive text generation is promising, and have been gaining attention in the community.
2. The proposed method is relatively versatile in experiments, allowing both standard generation, weak-to-strong guidance, and multi-objective alignment.

weaknesses:
1. Autoregressive Reward Model, trained by aligning the accumulated token-level rewards over a full response with the ground truth preference ordering, has been clearly proposed in the literature, e.g., [1,2] and the reference therein. Further, the idea of per-step-reward guided generation has been presented in [3] and the reference therein. The authors ought to have a adequate citation, discussion, and ideally comparison with these prior works. Otherwise, the contribution of this work will be significantly impaired. In particular, I strongly suggest the authors to move part of the current Related Work to the appendix for discussing in details these prior efforts for dense rewards in LLMs' RLHF/generation.

[1] Feng, Yihao, et al. "Fantastic Rewards and How to Tame Them: A Case Study on Reward Learning for Task-oriented Dialogue Systems." The Eleventh International Conference on Learning Representations.

[2] Yang, Shentao, et al. "Preference-grounded token-level guidance for language model fine-tuning." Advances in Neural Information Processing Systems 36 (2023).

[3] Li, Bolian, et al. "Cascade reward sampling for efficient decoding-time alignment." arXiv preprint arXiv:2406.16306 (2024).

2. The calculation of token sampling probability eq. (7) will roughly double the decoding time, especially compared to train-time alignment methods.

3. Source code seems not anonymously released for review purpose.

questions:
1. How would the per-step reward parametrization, eq. (4), mitigate the length bias that longer responses have more term to be sum up, and hence may automatically have a higher aggregate sequence-level reward?
2. In Figure 3 right, why does the word "words" also receives a low reward?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730330335623

thread_type:
official_review

<thread_object>
creator: reviewer_jvd2

<content>
summary:
The paper looks at reward models for RLHF and suggests training a token-level reward model from human preferences by parametrizing it as a sum of log probabilities. They prove that although it constrains the reward models to belong to a specific model class, this class is rich enough to model any reward under the Bradley-Terry preference framework. Empirically, the authors use the token-level reward to perform inference-time alignment.

strengths:
- The authors identify an important pain point of existing work, such as ARGS, that performs inference-time alignment on a token level without being trained for it.
- The theoretical analysis provides an important insight into the capacity of the parameterization the authors chose.
- The paper is written in a clear and easy-to-follow way.

weaknesses:
- While the focus on the downstream application of inference-time alignment is nice and leads to powerful results, I'm expecting work that suggests a new way to train a reward model to first evaluate how well the reward model itself is. Do the proposed reward models achieve better prediction capabilities than trajectory-level reward models? same level? I'm suggesting using rewardbench [1] for such evaluation, but even a classic train-test split on some preference dataset will be interesting.

- The author mentioned some prior work that distilled the reward function to a value function to be used in inference-time alignment [2]. Although these methods require an additional training step it is not a real downside, and I believe it should added as a baseline.


[1] Lambert, Nathan, et al. "Rewardbench: Evaluating reward models for language modeling." arXiv preprint arXiv:2403.13787 (2024).
[2] Han, Seungwook, et al. "Value Augmented Sampling for Language Model Alignment and Personalization." arXiv preprint arXiv:2405.06639 (2024).

questions:
- Can training-time alignment algorithms, like PPO, also benefit from the token-level reward? It makes credit assignment easier than trajectory-level rewards.

- [1] is another work that also deals with learning token-level information as part of the reward function. Although it uses different parameterizations, I believe it should be added to the related work section.

- Can you add confidence intervals to the results in Table 1? It will help the reader understand if something like a 48% win rate is significant or not.

[1] Nath, Vaskar, et al. "Learning Goal-Conditioned Representations for Language Reward Models." arXiv preprint arXiv:2407.13887 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730682512678

thread_type:
official_review

<thread_object>
creator: reviewer_9rar

<content>
summary:
This paper introduces GenARM, a testing-time alignment method that leverages the Autoregressive reward model, which is designed by the authors to predict next-token rewards for efficient and effective autoregressive generation. The authors show that this new design can allow us to efficiently guide frozen LLMs toward the target optimal distribution that is usually learned by training time alignment methods such as DPO or PPO. Experiments show that GenARM can outperform previous test-time alignment methods significantly without high costs of training larger models. Further, GenARM also supports multi-objective alignment, which allows users to balance trade-offs of different objectives without retraining.

strengths:
- The GenARM's design is relatively simple and straightforward, which allows the model to learn the reward effectively. 
- Compared with previous testing time alignment methods that naively used trajectory-based RM to guide the frozen LLMs to generate samples of target distributions, GenARM can provide relatively more accurate token-level guidance and signals with only a small costs to train the GenARM model;
- Empirical experiments shows that GenARM achieves much better performance without training the larger model, while the performance is close to high cost training time-based methods such as DPO;
- The simplicity of the GenARM method also allows users to align the frozen LLMs with multiple objectives, which provide flexibility for real-world alignment problems.

weaknesses:
- Compared with training time-based method, there is still gap of the current methods and DPO, while theoretically the authors show that the method can approximate the target optimal distribution;
- The evaluation experiments are relatively simple, only conducted on simple benchmarks;
- It would be good to compare [1] since this is also closely related to testing-time alignment.
- It would be good to add some discussion of previous token-level reward methods or literature, e.g. [2]




[1] Cascade Reward Sampling for Efficient Decoding-Time Alignment. https://arxiv.org/pdf/2406.16306
[2] Preference-grounded Token-level Guidance for Language Model Fine-tuning

questions:
- There are two ways to predict or output token level rewards, one is to directly view the problem as token-level prediction, and another method is to use an additional head to predict the real-values of the reward, I wonder if authors tried and compare these two methods?

- Regarding the gap of training time methods, I wonder if we leverage more advanced sampling methods such as MCTS to generate more trajectories, will the gap be eliminated?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732210308460

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
We thank all reviewers for their valuable feedback. The reviewers recognized our approach as novel (S14E), effective in addressing prior limitations (Jvd2), and paper well-written (Jvd2, S14E). They appreciated our comprehensive experiments (S14E), theoretical analysis (Jvd2), and the simplicity and effectiveness of our method (9RAr). Additionally, they commended its versatility across applications like standard text generation, weak-to-strong guidance, and multi-objective alignment (Dm7h, 9RAr), and found our work timely, promising (Dm7h, Jvd2), and cost-efficient compared to test-time baselines (9RAr).

Below, we address common concerns and misunderstandings and provide additional experimental results. We also outline the updates made to the paper.

---

### Concern 1: Related work on dense rewards

Reviewers Jvd2, 9RAr, and Dm7h highlighted several prior works on dense rewards in training LLMs that were missing from our related work section. We have updated Section 2 (lines 114–137) to incorporate these papers. Below, we summarize the most relevant related works:

We acknowledge that [1,2], highlighted by the reviewers, are pivotal works leveraging dense, token-level reward signals for **training** LLM policies with reinforcement learning (RL). However, our work differs in several key aspects:

1. Our paper focuses on **test-time** alignment, specifically without requiring any training of the LLM policy network. Unlike [1,2], which use dense reward signals for RL training, we leverage a token-level RM to enable efficient next-token prediction directly at test time.
2. Theoretically, we have proven that the parametrization of our proposed Autoregressive RM can guide frozen LLMs toward any distribution achievable by traditional RMs within the KL-regularized RL framework. This result is **fundamental for enabling test-time alignment** and is a capability **not addressed** in prior works that focus on providing dense rewards for training.

References:
[1] Feng, Yihao, et al. "Fantastic Rewards and How to Tame Them: A Case Study on Reward Learning for Task-oriented Dialogue Systems." ICLR 2023
[2] Yang, Shentao, et al. "Preference-grounded token-level guidance for language model fine-tuning." Neurips 2023.

---

### Concern 2: Comparison with another test-time alignment method

Reviewers 9RAr and Dm7h highlighted another test-time alignment method, CARDS [3], and Reviewer Dm7h expressed concern that it might use a similar token-level reward approach. We have revised our paper to include [3] in the related work section (Section 2, line 114-137) and compare it in the experiments (Section 6.1: Table 1 and Table 2). 

To clarify, unlike our method, **CARDS does not use token-level RM and suffers from significant inference costs**. Instead, it repeatedly generates small semantic segments and evaluates them using a **trajectory-level RM** to select the most satisfactory option. Similar to ARGS [4] (which we already compared with), [3] relies on trajectory-level RMs to evaluate partial responses, which can be inaccurate. Moreover, it induces significant inference costs as it needs to generate several candidate segments before choosing one. In contrast, our method is more efficient during inference, as the Autoregressive RM directly predicts the next-token reward for next-token generation. 

Experimentally, GPT-4 evaluation on the HH-RLHF dataset demonstrates that our proposed GenARM achieves a **higher win rate** compared to CARDS, as summarized below:

| Method | Win | Tie | Lose | Win + $\frac{1}{2}$Tie |
|-|-|-|-|-|
|GenARM VS CARDS|54.67|5.22|40.11|57.27|

Moreover, GenARM is significantly **more efficient** than CARDS during inference, as evidenced by the inference time required to generate 128 tokens using a 7B base LLM and a 7B RM, as detailed below.

|  | GenARM | CARDS |
|-|-|-|
|Inference time (s) |7.28|87.09|

References:
[3] Li, Bolian, et al. "Cascade reward sampling for efficient decoding-time alignment." arXiv preprint arXiv:2406.16306 (2024).
[4] Khanov, Maxim, Jirayu Burapacheep, and Yixuan Li. "ARGS: Alignment as Reward-Guided Search." ICLR (2024).

----
### Paper updates

All the updates in the paper are in blue.

* (JVd2, 9RAr, DM7h) Additional discussion on prior literature on dense rewards in Section 2 (lines 114–137). 
* (9RAr, DM7h) Adding experimental comparison with CARDS in Section 6.1 (Table 1, Table 2.)
* (JVd2) Adding confidence intervals to Table 1.
* (JVd2) Adding results on reward model accuracies in Section 6.1 (lines 342) and Section 6.2 (lines 398).
* (S14E) Adding results on comparison with BoN in Section 6.1 (Table 1)
* (DM7h) Moving parts of related work sections to Appendix A (lines 717-738).
* (S14E) Adding training-time comparison in Appendix C.1 (lines 803-817).


---
We thank all reviewers again for the time and effort in reviewing our paper! Please let us know if the above explanations do not address your concerns. We are happy to answer any further questions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes autoregressive reward models for test-time alignment. Given preference data, it trains a reward model that can predict a reward for each token as $log p_r(y_t \mid x_t, y_{<t})$, and the total reward is given by the sum of these individual rewards. It then trains the reward model using the standard approach of maximizing the log-likelihood assuming the Bradley-Terry distribution. At test time, we generate tokens using a product of the base model's probability multiplied by an adjusted reward model $p_r$ score.

However, once the reward model is trained, the base LLM isn't trained but instead is guided at test time using the learned token-level reward model. So, why is this better than training the base LLM directly using DPO? Authors claim that when the reward model is a weaker LLM, this is more computationally efficient for training.

Strengths:
1. Novel approach for test-time alignment. The most useful fact is that many people are not releasing trained reward models. This can allow someone to quickly try these reward models with any LLM to get quick numbers without having to train the reward model -- something that not everyone can do. This is a highly useful fact!

2. Authors were able to add experiments with the CARDS baseline and show that their approach does better. 

3. Reviewers found the approach simple and easy to use. This matters for real-world usage.

Weakness:
1. When using a weaker reward model, the performance does not match that of using a DPO trained policy. However, the authors do show that when using the same reward model class as the LLM, the performance does match. But when you do this, you might as well train the LLM directly using DPO as the computational benefit of avoiding training a big model is lost. Further, there is some overhead at inference due to computing rewards as well. For this reason, I think labs that can afford to train a model and have fixed preference data will not find this approach useful. 

2. The reviewer suggested experiments with the value distillation approach that the authors were unable to add. This would have been nice.

Overall, a good approach. There are some trade-offs here that I mentioned above. The real audience of this paper will be people with fewer resources who can only train small reward models, or who can benefit from large GenARM reward models being made available and who can quickly try these reward models with **any** LLM, that has same vocabulary, without having to do expensive RLHF. This is an important advantage and which is why I recommend accepting.