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

We thank Reviewer S14E for the detailed feedback. We are encouraged that the reviewer finds our writing clear, our approach novel and our experiments comprehensive and informative. We address concerns below.

---

Weakness 1:.. 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.

To clarify, the reward model training costs of our method and other test-time alignment baselines, such as ARGS, are almost the same, ensuring a fair comparison. In all of our experiments, both the Autoregressive RM (used by GenARM) and trajectory-level RM (used by baselines like ARGS) are initialized from the same 7B LLM. Importantly, other test-time alignment methods like ARGS also require training a 7B trajectory-level RM.

For a more detailed analysis, including training time comparisons, please refer to our response to Weakness 2.

the training cost of the reward model is basically the same as training-time alignment methods like DPO.

This is not accurate: in our weak-to-strong guidance experiments (Section 6.2), where test-time alignment methods use small RM to guide larger LLM, the training costs of the 7B Autoregressive RM and 7B trajectory-level RM are significantly lower than those of training-time baselines, which require training much larger LLMs, such as a 70B model.

Weakness 2: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.

Thank you for the suggestion! We have updated our paper to include a detailed training time comparison in Appendix C.1 (line 803–817). To summarize, training Autoregressive RM only introduces negligible increase in training time, compared with trajectory-level RM: We compared them under identical conditions. Both models were trained on the filtered HH-RLHF dataset using LoRA fine-tuning for one epoch, and both were initialized from the same LLaMA-7B model.

The results are:

• Training trajectory-level RM took 11,896 seconds
• Training Autoregressive RM took 12,162 seconds (only 2% more than above)

This demonstrates that training Autoregressive RM introduces a minimal increase in time compared to a trajectory-level RM of the same size. This is expected, as both models are 7B, with the only difference being in the last layer, which adds just 1% more parameters (trajectory-level RM: 6.61B vs. Autoregressive RM: 6.72B).

Weakness 3: 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?

Thank you for pointing this out! In Section 6.1, we have updated our paper to include results for BoN (N=16). The results show that GenARM achieves a 57.22% win rate and a 60.55% win + 0.5 tie rate against BoN. This demonstrates that GenARM outperforms BoN while being significantly more efficient, as BoN requires generating N samples from the base LLM.

We did not include Transfer-Q in Section 6.2 (Weak-to-Strong Guidance) due to its prohibitive inference time. Transfer-Q generates the full response before predicting the next token, making it extremely slow for scenarios with longer response lengths and larger models. On the AlpacaEval2 benchmark, responses are lengthy (e.g., over 1000 tokens), and in our weak-to-strong guidance experiments, we use larger base LLMs, such as 70B models. Based on our estimates, running Transfer-Q in this setting would take over 2000 hours, making it infeasible for evaluation.

Question: it would be great if the authors can explain the connection of the proposed method with DPO objective

Thank you for the question! When training an LLM with DPO, the reference policy (i.e., the base LLM) must be pre-specified during training. This means the alignment is tightly coupled to the specific base LLM used during the training phase.

In contrast, when training the Autoregressive RM in our method, no base LLM is used during training. Instead, the trained Autoregressive RM can be flexibly combined with different base LLMs during test-time, offering significant configurability. For example:

• Pairing a smaller Autoregressive RM with a larger base LLM for weak-to-strong guidance.
• Using multiple Autoregressive RMs to guide a single base LLM for multi-objective alignment.

The key distinction lies in test-time flexibility: DPO ties the alignment process to a specific base LLM during training, whereas our method decouples the RM training from the base LLM, enabling diverse test-time applications and configurations.

---

Thank you again for your 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.

We thank Reviewer Dm7h for the detailed feedback. We are encouraged that the reviewer finds our work on addressing the mismatch between trajectory-level reward models (RM) and autoregressive text generation both promising and timely. We also appreciate the reviewer's recognition of our method's versatility, enabling a wide range of applications, including standard text generation, weak-to-strong guidance, and multi-objective alignment. Below we address the concerns in details.

---

Weakness 1.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.

Thank you for highlighting these relevant works on dense rewards! We have updated our related work section (Section 2, line 114-137) to incorporate and discuss these contributions, as detailed below:

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). These approaches tackle the well-known challenges of learning from sparse and delayed reward signals in RL [4]. 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.

Weakness 1.2: Further, the idea of per-step-reward guided generation has been presented in [3] and the reference therein.

We have revised our paper to include CARDS [3] in the related work section (Section 2, line 114-137) and compare it with our method in the experiments (Section 6.1: Table 1 and Table 2). Unlike our approach, [3] does not use a token-level RM to guide text generation. Instead, it repeatedly generates small semantic segments (a few tokens) and evaluates them using a trajectory-level RM to select the most satisfactory option. Similar to ARGS [5] (which we already cited and compared with in the original version), [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 [3], as summarized below:

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.

Therefore, our conclusion remains unchanged: our method achieves the highest win rate and the lowest inference cost among all test-time alignment approaches.

Weakness 1.3: The authors ought to have a adequate citation, discussion, and ideally comparison with these prior works. 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.

Thank you for the suggestions! We have added these works on dense rewards to our related work section (Section 2, line 114-137) and included [3] in our experimental evaluation (Section 6.1: Table 1 and Table 2). To accommodate these additions, as suggested, we have moved part of the related work to Appendix A.

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. [3] Li, Bolian, et al. "Cascade reward sampling for efficient decoding-time alignment." arXiv preprint arXiv:2406.16306 (2024). [4.] Ng, Andrew Y., Daishi Harada, and Stuart Russell. "Policy invariance under reward transformations: Theory and application to reward shaping." Icml. Vol. 99. 1999. [5] Khanov, Maxim, Jirayu Burapacheep, and Yixuan Li. "ARGS: Alignment as Reward-Guided Search." ICLR (2024).

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

Test-time alignment methods, as demonstrated in our paper, offer unique advantages such as convenient multi-objective alignment and weak-to-strong guidance, which training-time alignment methods struggle to achieve. Therefore, when evaluating inference time, the comparison should primarily be made against other test-time alignment methods. Importantly, our method achieves the lowest inference time among all test-time methods, as shown in Table 2.

will roughly double the decoding time

This is not accurate: when the reward model (RM) is significantly smaller than the LLM (e.g., RM is 7B and LLM is 70B in our weak-to-strong guidance experiment), our method introduces only marginal inference overhead (e.g, 10%) compared to training-time methods. Furthermore, for very large LLMs (e.g., 70B), where training becomes prohibitively expensive, training-time methods are no longer viable. In such cases, our method remains an efficient solution since it only requires training the smaller RM, making it both practical and cost-effective.

---

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

We have uploaded the code for review purpose.

---

Question 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?

It is not true that longer responses automatically yield a higher or lower aggregate sequence-level reward. To illustrate, consider the following two responses:

• I love apples [EOS] (shorter response)
• I love apples for dinner [EOS] (longer response)

Denote token-level reward as r (in our case it is a log probability.) There is no guarantee that r(EOS|I love apples) is higher or lower than r(for|I love apples) + r(dinner|I love apples for) + r(EOS|I love apples for dinner). Furthermore, our training objective in eq. (5) forces the model to assign higher score on prefered response, regardless of the length.

This principle also holds in traditional text generation: While the log probability of a text sequence is computed as the sum of individual token log probabilities, next-token sampling does not inherently favor longer or shorter responses.

---

Question 2: In Figure 3 right, why does the word "words" also receives a low reward?

Good question!

This RM trained for harmlessness is designed to distinguish harmful tokens from harmless ones. The reward values are meaningful only when comparing the next tokens conditioned on the same prefix; absolute reward values are not directly interpretable.

To clarify, we conducted a simple quantitative test with the following examples:

• An effecitve way ... use cruel words (original sentence in Fig. 3)
• An effecitve way ... use cruel defamation (new sentence)

While "words" has a seemingly low absolute reward of -1.94, the more harmful token "defamation" receives an even lower reward of -3.21. This demonstrates that during guided generation, the RM correctly assigns a lower probability to generating "defamation," which aligns with its goal of discouraging harmful language. This behavior is exactly what we expect from a token-level RM focused on harmlessness.

---

Thank you again for your 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.

We thank Reviewer Jvd2 for the detailed feedback and for recognizing our key insight into the limitations of existing methods that use reward models not explicitly trained for token-level decoding. We also appreciate the positive comments on our paper's clarity and theoretical analysis. Below, we address the concerns in detail.

---

Weakness 1: 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.

Thank you for the insightful suggestions! We have updated our paper (lines 342 and 398) to include the reward accuracies of both the Autoregressive RM and the traditional trajectory-level RM on the test splits of HH-RLHF and UltraFeedback. Our results show that the proposed Autoregressive RM achieves the same level of accuracy to the trajectory-level RM:

• HH-RLHF: Autoregresive RM = 77.6% and trajectory-level RM = 78.3%
• UltraFeedBack: Autoregresive RM = 75.1% and trajectory-level RM = 75.4%

It is important to note, however, that the best way to assess the quality of an RM for test-time alignment is to evaluate the quality of the generated responses. This has already been thoroughly analyzed in our paper.

---

Weakness 2: 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.

Although it requires an additional training step it is not a real downside

One potential downside of distilling a trajectory-level RM into a value function, as done in [2], is the risk of error accumulation. The trajectory-level RM itself introduces some error (e.g, accuracy = 75%), and the additional step of distilling it into a value function can compound these errors. In contrast, GenARM avoids this issue by directly learning token-level reward signals from the data, eliminating the need for a separate distillation step and reducing the potential for cumulative errors.

I believe it should added as a baseline.

We attempted to run the code provided by [2] but were unable to do so successfully. After contacting the authors, they informed us that VAS is in the process of being integrated into the official TRL codebase and suggested we try an example that is still under development. Despite our efforts, we encountered issues with the code. We are currently awaiting the authors' response and appreciate their assistance. Once the complete code becomes available, we will include VAS as a baseline in our work.

[2] Han, Seungwook, et al. "Value Augmented Sampling for Language Model Alignment and Personalization." 2024.

---

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

Yes, reinforcement learning algorithms like PPO can benefit from dense, token-level rewards, as highlighted in prior work [1,2]. We have updated our related work section (lines 127–134) to include a discussion on these approaches and their use of dense rewards for training LLMs with RL.

[1] Ng, Andrew Y., Daishi Harada, and Stuart Russell. "Policy invariance under reward transformations: Theory and application to reward shaping." Icml. Vol. 99. 1999. [2] Yang, Shentao, et al. "Preference-grounded token-level guidance for language model fine-tuning." Neurips 2023.

---

Question 2: [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.

Thank you for highlighting this interesting paper on token-level information from a representation perspective! We have updated our related work section to include it.

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

---

Question 3: 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.

Thank you for the suggestions! We have updated Table 1 to include confidence intervals (based on three experiments), and the conclusions remain consistent with our original results.

---

Thank you again for your 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.

We thank Reviewer 9RAr for the detailed and thoughtful feedback. We are pleased that the reviewer finds our method's design simple and effective, our test-time approach more accurate and cost-efficient than existing test-time alignment baselines, and our weak-to-strong guidance results promising. We also appreciate reviewer's recognition of our method's simplicity in handling multiple objectives. Below, we address the reviewer's concerns in detail.

---

Weakness 1: 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;

To clarify, when the reward model (RM) and the base LLM are of the same size, our proposed GenARM matches the performance of the training-time method DPO (there is no performance gap). In contrast, other test-time alignment baselines, such as ARGS and TransferQ, perform worse than DPO in this setting. Notably, GenARM effectively closes this performance gap.

When the RM is smaller than the base LLM (e.g., RM is 7B while the base LLM is 70B), GenARM does exhibit a performance drop compared to DPO. However, it is important to emphasize that GenARM requires training only the smaller RM (e.g., 7B), whereas training the much larger LLM (e.g., 70B) in DPO may be computationally prohibitive or entirely infeasible. This makes GenARM a practical and scalable alternative in such scenarios.

As for the theoretical results: We show that the proposed Autoregressive RM is expressive enough as a function class to represent any decoding distribution. This expressiveness is crucial; using a parametrization that cannot capture the target distribution would undermine the entire test-time alignment approach. However, this result does not guarantee how well such a function can be learned from the data. Empirically, larger LLMs with more pretrained knowledge provide a better initialization for reward models, leading to more accurate RMs after fine-tuning. Due to computational constraints, we only trained 7B RMs in this work, but we anticipate that training a larger RM could further enhance performance.

---

Weakness 2: The evaluation experiments are relatively simple, only conducted on simple benchmarks;

Our evaluation experiments are sufficiently comprehensive to demonstrate the effectiveness of GenARM in various scenarios, including comparisons with test-time and training-time baselines, standard reward-guided decoding, and practical applications like weak-to-strong guidance and multi-objective alignment. Specifically:

1. Standard Reward-Guided Generation on general preference (Section 6.1):
   • Conducted on the widely used HH-RLHF dataset with RM and LLM of the same size.
   • Results show that:
     • GenARM outperforms other test-time alignment baselines.
     • GenARM is more efficient during inference.
     • It achieves comparable performance to training-time alignment methods.
2. Weak-to-Strong Guidance (Section 6.2):
   • Evaluated on the widely used AlpacaEval2 benchmark.
   • Results show that:
     • GenARM outperforms test-time alignment baselines.
     • It performs comparably to training-time methods when RM and LLM are of the same size.
     • GenARM effectively provides weak-to-strong guidance, recovering most of the performance gap with training-time methods while requiring only the smaller RM. This is crucial in scenarios where training-time methods are computationally infeasible.
3. Multi-Objective Alignment (Section 6.3):
   • Conducted on the widely used multi-objective dataset PKU-SafeRLHF.
   • Results highlight that:
     • GenARM is efficient and effective in aligning with multi-dimensional preferences.
     • GenARM enables weak-to-strong guidance in multi-objective alignment.

These experiments, leveraging widely recognized benchmarks, effectively showcase the versatility and practicality of GenARM across real-world alignment challenges.