Adaptive dense reward:Understanding the Gap Between Action and Reward Space in Alignment

We would like to sincerely thank Reviewer 7YEj for providing a detailed review with insightful suggestions.

W1: in lines 272 - 274, the authors claim "Bradley terry model can be represented as an MDP", this is incorrect.

Thank you for pointing that out. The intent here is to convey that autoregressive generative models can be viewed as MDPs, which is widely considered the foundation for applying RL to LLMs. We have revised the wording to: can be represented as a Markov Decision Process (MDP) due to their autoregressive generative Transformer structure.

W2-8:

We truly apologize for writing quality.

• We have rechecked the formatting of the entire paper. Improvements include a comprehensive check and correction of spacing, formulas, and citations.
• The introduction now contains the research objectives, significance, and main contributions of this work. Parts that duplicated the literature review have been removed or relocated to Section 5.
• The literature review has been significantly condensed, focusing on a comparison of the most relevant work. Section 5.1 now covers alignment（As reviewer GTLV's opinion）and RLHF, while Section 5.2 discusses research on fine-grained reward signals.

Q1-2: Why are some experiments truncated? Can the authors explain their efforts that could guarantee the comparisons are fair to all algorithms?

Some experiments were stopped early when both algorithms reached a performance plateau, indicating further training offered minimal improvement. This was done to conserve computational resources.

To ensure fairness in comparisons, we took several measures:

• Consistent Data and Code: All algorithms were trained and evaluated on the same dataset using the same codebase.
• Controlled Random Seed: A fixed random seed (42) was used across all experiments to minimize the impact of randomness and ensure reproducibility.
• Multiple Runs: Each algorithm was run twice, with results consistent within 0.1%, except for minor variations due to multi-machine communication and flash-attention. This consistency reinforces the reliability of the reported results and minimizes the potential bias introduced by early stopping. We are happy to provide the results from all runs if requested.

Q3-4: How do the authors determine the values of delta? Is the b learnable?

Following the online-DPO training procedure, δ was set to 0.1, consistent with the threshold.

We have two methods for setting the baseline(b).

1. Offline sampling to calculate the mean reward score, and then using 0 as the baseline after subtracting this mean from the reward model's scores.

2. The second method starts by using 0 as the baseline to distinguish positive and negative rewards at step 0 of the experiment. Simultaneously, we initialize an array to store the current step and the mean reward. As training progresses, we dynamically update this mean based on the reward scores given by the reward model, and use this updated mean as the baseline.

In our experiments, we found the first method to be more effective. All experiments in this paper utilize the first method.

Q5: Equation 11 and equation 16 present different possible values for M.

The meanings of 'M' in Equation 11 and Equation 16 are distinct. In Equation 11, 'M' represents a loss mask, taking binary values of 0 and 1. Conversely, 'M' in Equation 16 signifies an evaluation of corpus quality, categorized as positive, negative, or neutral.

Q6: In Table 1, could the authors also compare with other methods like step-dpo, token-dpo?

I've attempted a comparison, but I haven't been able to find a fair comparison method. Both methods' authors only provide baselines on specific tasks (like math and summarization). Step-DPO's performance is significantly affected by the choice of step divisions, and token-DPO fails to converge on general tasks. Comparing them may not be essential to the core contribution of this paper.

Q1: explain the concept of the gap between action and reward space alignment.

In autoregressive generative transformers, each token represents an action. The size of the vocabulary is the size of the action space. As mentioned in Section 1, in RLHF the action space is typically more sparse than the reward space. Therefore, we need to increase the density of the reward space to align them.

Q2: How is the threshold b in Eq. 11 determined

We have two methods for setting the b.

1. Offline sampling to calculate the mean reward score, and then using 0 as the baseline after subtracting this mean from the reward model's scores.

2. The second method starts by using 0 as the baseline to distinguish positive and negative rewards at step 0 of the experiment. Simultaneously, we initialize an array to store the current step and the mean reward. As training progresses, we dynamically update this mean based on the reward scores given by the reward model, and use this updated mean as the baseline.

In our limited experiments, we found the first method to be more effective. All experiments in this paper utilize the first method.

Q3&W5: Have you experimented with alternative model architectures?

Yes. We've observed similar results on both Llama 3 and our closed-source model. If you believe it's absolutely necessary, I can supplement with experiments on Llama 3.

Q4: What are the additional computational costs associated with this method?

It requires fine-grained annotations (compared to the sentence-level preference annotations in RLHF), or a well-trained reward model.

We are grateful to the reviewer BZ8y for careful reading of our manuscript and valuable feedback.

W1: Limited Contribution: The main contribution, calculating the mask value for each instance in Eq. 11, seems insufficiently significant.

That's an excellent suggestion. We've therefore expanded the presentation of our theoretical analysis.

• The introduction now contains the research objectives, significance, and main contributions of this work. Parts that duplicated the literature review have been removed or relocated to Section 5.
• To provide a clearer problem definition, we've added Section 2.3. In Section 2.3, we provide a detailed definition of the problem. In Section 2.3, we define and differentiate between the two types of alignment discussed in this paper. We clarify that the alignment referred to in the title is preference alignment (Section 2.3.1) and also explain representational alignment (Section 2.3.2, noting the need for a denser reward signal to align the action space).
• In Section 3.1, we identify potential sources of error when the reward signal is sparse, clarifying why our approach of identifying key information within the sequence can effectively divide subsequence and reduce errors.

In Section 3.1, we define the supervisory signal error. Traditional methods assign the reward of a step to all tokens within that step/sequence. The error in this method stems from the discrepancy between the overall reward for the sequence $r_s$ and the reward assigned to individual actions/tokens $r_t$. The error is formulated as:$\text{err}_{\text{sequence level}} = \sum (r_t - r_s)^2$ At the token level, the error from this method arises from random noise. Assuming we can whiten the rewards to have a mean of 0 and variance σ², the token-level error is:

$\text{err}_{\text{token level}} = \sum {c_i}^2 = \sigma^2N$

where $c_i$ represents the random noise, and $N$ is the total sequence length. The total error is the sum of sequence-level and token-level errors:

$\text{err} = err_s + err_t = \sum_{k=1}^{K} \sum_{t \in S_k} (r_t - r_k)^2 + \sigma^2 K$

where $K$ is the number of steps/sequences, and $S_k$ represents the set of tokens in the k-th step/sequence. $r_k$ represents the reward assigned to the k-th sequence. This formula effectively illustrates two key issues:

1. Error: Steps must be effectively divided based on crucial information within the context to minimize the overall error across all tokens.

2. Variance in Token-Level Rewards: Why the reward signal at the token level exhibits greater variance.

W2: Questionable Effectiveness.

Your suggestion provided a valuable alternative perspective, which led us to add Section 3.2.1 and appendix D.

We would like to clarify that the mask used here is a loss mask (as opposed to an attention mask). This technique is widely adopted in the community, for instance, in the HuggingFace TRL framework and others' OpenRLHF frameworks.

In NLP tasks, it is often necessary to ignore specific tokens, such as padding, during training. Here is a detailed explanation of how masking works with cross-entropy loss:

Cross-Entropy Loss Definition: $L = -\sum_{i} y_i \log(p_i)$

Applying Mask to the Loss: Define a mask $m_i$, where $m_i = 0$ if the token at position $i$ is to be ignored, and $m_i = 1$ if it should be included in the loss. To ignore tokens, the masked loss is calculated as: $L = -\sum_{i} m_i \, y_i \log(p_i)$ This ensures that positions where $m_i = 0$ contribute zero to the loss, effectively ignoring those tokens.

Effect on Gradients:

$$m_i \, y_i \log(p_i) = 0 \quad \text{if} \quad m_i = 0$$

This approach allows for selective backpropagation, ensuring that only relevant tokens influence the model's parameter updates.

You can refer to the following works：

• OpenRLHF: https://github.com/OpenRLHF/OpenRLHF/blob/78e1fbb7f34cb313fe63cc0eb0a6ba5b7ed764a9/openrlhf/trainer/dpo_trainer.py#L388
• TRL: https://github.com/huggingface/trl/blob/066fc37bd3381e5de3ac0bb988ce834a050c459f/trl/trainer/dpo_trainer.py#L1150

W3-4&W6: Poor Writing Quality.

We truly apologize for this issue. The literature review has been significantly condensed, focusing on a comparison of the most relevant work. Section 5.1 now covers alignment（As reviewer S1Kp'opinion）and RLHF, while Section 5.2 discusses research on fine-grained reward signals. We have rechecked the formatting of the entire paper. Improvements include a comprehensive check and correction of spacing, formulas, and citations.

We sincerely thank Reviewer GTLV for providing the insightful review, and we have revised our paper based on your review carefully.

W1&W3: Attention to writing quality is needed & Consistent errors in subscript formatting were observed in equations.

We truly apologize for this issue. We have rechecked the formatting of the entire paper. Improvements include a comprehensive check and correction of spacing, formulas, and citations.

W2: The paper could benefit from separate sections for literature review and problem definition.

Thank you for your insightful suggestion.

• The introduction now contains the research objectives, significance, and main contributions of this work. Parts that duplicated the literature review have been removed or relocated to Section 5.
• To provide a clearer problem definition, we've added Section 2.3. In Section 2.3, we provide a detailed definition of the problem. In Section 2.3, we define and differentiate between the two types of alignment discussed in this paper. We clarify that the alignment referred to in the title is preference alignment (Section 2.3.1) and also explain representational alignment (Section 2.3.2, noting the need for a denser reward signal to align the action space).
• In Section 3.1, we identify potential sources of error when the reward signal is sparse, clarifying why our approach of identifying key information within the sequence can effectively divide subsequence and reduce errors.

W3: The theoretical foundation appears weak.

We appreciate your insightful suggestion and have enhanced the theoretical analysis section.

We've added Section 3.2, which provides a problem definition, and Section 3.1, which theoretically derives the optimal solutions for error and variance. In Section 3.1, we define the supervisory signal error. Traditional methods assign the reward of a step to all tokens within that step/sequence. The error in this method stems from the discrepancy between the overall reward for the sequence $r_s$ and the reward assigned to individual actions/tokens $r_t$. The error is formulated as:$\text{err}_{\text{sequence level}} = \sum (r_t - r_s)^2$ At the token level, the error from this method arises from random noise. Assuming we can whiten the rewards to have a mean of 0 and variance σ², the token-level error is:

$\text{err}_{\text{token level}} = \sum {c_i}^2 = \sigma^2N$

where $c_i$ represents the random noise, and $N$ is the total sequence length. The total error is the sum of sequence-level and token-level errors:

$\text{err} = err_s + err_t = \sum_{k=1}^{K} \sum_{t \in S_k} (r_t - r_k)^2 + \sigma^2 K$

where $K$ is the number of steps/sequences, and $S_k$ represents the set of tokens in the k-th step/sequence. $r_k$ represents the reward assigned to the k-th sequence. This formula effectively illustrates two key issues:

1. Error: Steps must be effectively divided based on crucial information within the context to minimize the overall error across all tokens.

2. Variance in Token-Level Rewards: Why the reward signal at the token level exhibits greater variance.

W4: Visual content preparation is lacking.

We've adjusted the size of the images and legends.

Q1-2: A claim about dense supervision signals improve learning efficiency and generation quality.

Better allocation of reward signals can improve efficiency and training outcomes compared to relying on the model to perform adaptive credit assignment on tasks where it struggles. This is widely acknowledged within the RL community, for example:

• Andrychowicz, Marcin, et al. "Hindsight experience replay." Advances in neural information processing systems 30 (2017).

But your question is excellent. This paper doesn't offer a theoretical analysis or experimental results on this point, so these statements have been removed.

Q3: About dynamically updating the baseline.

We have two methods for setting the baseline.

1. Offline sampling to calculate the mean reward score, and then using 0 as the baseline after subtracting this mean from the reward model's scores.

2. The second method starts by using 0 as the baseline to distinguish positive and negative rewards at step 0 of the experiment. Simultaneously, we initialize an array to store the current step and the mean reward. As training progresses, we dynamically update this mean based on the reward scores given by the reward model, and use this updated mean as the baseline.

In our experiments, we found the first method to be more effective. All experiments in this paper utilize the first method.

Q4: selection on datasets

1. HelpSteer's richer data domain, encompassing longer sequences, makes it better suited to our needs than other datasets. The inclusion of PRM800K further ensures comprehensive coverage across a diverse range of tasks.

2. Our annotation resources are limited. Step-level annotation of the dataset is prohibitively expensive.

Thank you again for your patient review and guidance.

Problem Definition

The term "alignment" in the title refers to aligning model outputs with human preferences, a notion we believe is generally familiar to readers. For example:

• Tang, Yunhao, et al. "Understanding the performance gap between online and offline alignment algorithms." arXiv preprint arXiv:2405.08448 (2024).
• Guo, Shangmin, et al. "Direct language model alignment from online ai feedback." arXiv preprint arXiv:2402.04792 (2024).
• Xu, Shusheng, et al. "Is dpo superior to ppo for llm alignment? a comprehensive study." arXiv preprint arXiv:2404.10719 (2024).
• Zhu, Banghua, et al. "Fine-tuning language models with advantage-induced policy alignment." arXiv preprint arXiv:2306.02231 (2023).

Section 3.1

This section simply identifies a source of error（from reward model）. A detailed explanation and derivation of the masking method are provided in Section 3.2.1 and Appendix D.

This technique is widely adopted in the community, for instance, in the HuggingFace TRL and OpenRLHF.

In NLP tasks, it is often necessary to ignore specific tokens, such as padding (<pad>), during training. Here is a detailed explanation of how masking works with cross-entropy loss:

Cross-Entropy Loss Definition: $L = -\sum_{i} y_i \log(p_i)$

Applying Mask to the Loss: Define a mask $m_i$, where $m_i = 0$ if the token at position $i$ is to be ignored, and $m_i = 1$ if it should be included in the loss. To ignore tokens, the masked loss is calculated as: $L = -\sum_{i} m_i \, y_i \log(p_i)$ This ensures that positions where $m_i = 0$ contribute zero to the loss, effectively ignoring those tokens.

Effect on Gradients:

$$m_i \, y_i \log(p_i) = 0 \quad \text{if} \quad m_i = 0$$

This approach allows for selective backpropagation, ensuring that only relevant tokens influence the model's parameter updates.

You can refer to the following works：

• OpenRLHF: https://github.com/OpenRLHF/OpenRLHF/blob/78e1fbb7f34cb313fe63cc0eb0a6ba5b7ed764a9/openrlhf/trainer/dpo_trainer.py#L388
• TRL: https://github.com/huggingface/trl/blob/066fc37bd3381e5de3ac0bb988ce834a050c459f/trl/trainer/dpo_trainer.py#L1150

The masking mechanism is not the key point. Our work here uses it to explore the potential benefits of segmenting steps based on rewards from the reward model, suggesting that this approach may be more effective than using punctuation like periods or semicolons.

Dear Authors,

Thank you for your efforts in addressing my concerns.

I’m glad to see that W1, W3, and W5 have been successfully resolved. However, W2 and W4 still require further attention.

Problem Definition: It would be helpful to present the problem definition in a more formalized manner, as is standard in the ML community. A formalized approach would make the concept of 'alignment' more mathematically precise and persuasive for readers. Unfortunately, the revised version still lacks this level of clarity.

Section 3.1: This section still falls short in providing theoretical insights into the masking-based approach. While it explains how to quantify the mismatching error between token-level rewards and sequence-level rewards, it does not offer a theoretical explanation of how masking certain subsequences can help reduce this error. For example, it is well known that masking-based loss function can be formalized as a composite likelihood function, and therefore, mathematically showing that minimization of the composite (or pseudo) likelihood objective (in the setting of interest) leads to the variance reduction in rewards, compared to the general likelihood function, can be one promising approach.

Given these points, I will keep my score unchanged.

Once again, I sincerely appreciate the authors’ efforts in addressing my concerns.

We thank the reviewer wBaN for insightful comments and constructive suggestions, which have significantly improved the quality of this manuscript.

W1: Lack of depth in the proposed algorithm:

Thank you for bringing this matter to our attention.

• The introduction now contains the research objectives, significance, and main contributions of this work. Parts that duplicated the literature review have been removed or relocated to Section 5.
• To provide a clearer problem definition, we've added Section 2.3. In Section 2.3, we provide a detailed definition of the problem. In Section 2.3, we define and differentiate between the two types of alignment discussed in this paper. We clarify that the alignment referred to in the title is preference alignment (Section 2.3.1) and also explain representational alignment (Section 2.3.2, noting the need for a denser reward signal to align the action space).
• In Section 3.1, we identify potential sources of error when the reward signal is sparse, clarifying why our approach of identifying key information within the sequence can effectively divide subsequence and reduce errors.

In Section 3.1, we define the supervisory signal error. Traditional methods assign the reward of a step to all tokens within that step/sequence. The error in this method stems from the discrepancy between the overall reward for the sequence $r_s$ and the reward assigned to individual actions/tokens $r_t$. The error is formulated as:$\text{err}_{\text{sequence level}} = \sum (r_t - r_s)^2$ At the token level, the error from this method arises from random noise. Assuming we can whiten the rewards to have a mean of 0 and variance σ², the token-level error is:

$\text{err}_{\text{token level}} = \sum {c_i}^2 = \sigma^2N$

where $c_i$ represents the random noise, and $N$ is the total sequence length. The total error is the sum of sequence-level and token-level errors:

$\text{err} = err_s + err_t = \sum_{k=1}^{K} \sum_{t \in S_k} (r_t - r_k)^2 + \sigma^2 K$

where $K$ is the number of steps/sequences, and $S_k$ represents the set of tokens in the k-th step/sequence. $r_k$ represents the reward assigned to the k-th sequence. This formula effectively illustrates two key issues:

1. Error: Steps must be effectively divided based on crucial information within the context to minimize the overall error across all tokens.

2. Variance in Token-Level Rewards: Why the reward signal at the token level exhibits greater variance.

W2: Statistical significance of the results

Thank you for raising this concern. While the robustness of the experiments could indeed be further improved, we believe this does not invalidate the core contributions and the correctness of the proposed approach. As other reviewers mentioned, extensive experimentation and testing may provide similar validation to running multiple experiments with different random seeds.

W3&Additional suggestions: The obscurity in the writing of the paper

• sample level supervision: The terminology in this paper has been standardized: all instances of "sample level" have been replaced with "sequence level (wise)."
• We have added content related to Process-supervised Reward Models in Section 2.1.
• We've enhanced the figures with more detailed explanations and ensured that each figure is referenced within the main text.
• The part of the Schmitt trigger has been moved to the experimental settings section and the appendix, and corresponding citations have been added.
• We have rechecked the formatting of the entire paper, including a comprehensive check and correction of spacing, formulas, and citations.
• The introduction now contains the research objectives, significance, and main contributions of this work. Parts that duplicated the literature review have been removed or relocated to Section 5.
• Section 2.3 on contextual dueling bandits: This section, which discussed the challenges of sparse reward signals requiring adaptive credit assignment, has been removed for conciseness and practicality.

Q1: What are the authors referring to by "understanding gap between action and reward space in alignment" in the title?

In autoregressive generative transformers, each token represents an action. The size of the vocabulary is the size of the action space. As mentioned in Section 1, in RLHF the action space is typically more sparse than the reward space. Therefore, we need to increase the density of the reward space to align them.

Q2: "Human Eval et al." do you mean "Human Eval etc."?

Yes. Thank you for your suggestion. We have made this revision.