Kati8/ICLR_2023_2025 huggingface dataset:
authors: Rui Zheng
authors: Wei Shen
authors: Yuan Hua
authors: Wenbin Lai
authors: Shihan Dou
authors: Yuhao Zhou
authors: Zhiheng Xi
authors: Xiao Wang
authors: Haoran Huang
authors: Tao Gui
authors: Qi Zhang
authors: Xuanjing Huang
reviews: 
Comments: This paper suggests a method for Improving RLHF by preventing it to learn shortcuts/reward-hacking using a concept data groups where groups are classified rather than given beforehand. At the training time, the suggested approach puts more focus on learning the best behavior for challenging groups.

The grouping objective is adversarial to find groups with maximum discrepancy in reward, while the RL objective is encouraging reduction of loss variance across groups

-- Updated my evaluation after reading the responses and revisions
Strengths: - The paper is well written and clear
- The method is generally sound and intuitive
- Promising experiment results
Weaknesses: 1- It is unclear how to find the optimal number of groups? The paper seems to miss discussion of this and sharing what settings used in the presented results, and how authors reached that setting.

2- The adversarial objective may pickup noisy or outliers and divert the optimization. I think this deserves more investigations. 

3- The learning objective is encouraging reduction of loss variance across groups, I was wondering why authors didn’t directly go for optimization macro average of loss which is a bit more intuitive and has less chance to have side-effects on the overall performance as mean is generally more stable than variance?

4- Robust optimization has been extensively studied in the general optimization context. Many of such methods could be applicable to the RLHF/LLM problem and the method proposed in this paper is also also applicable to other settings. I do not see any comparisons to support this method is optimal for RLHF compared to existing robust optimization methods.

5- Related to the previous two comments, we we should have compared with other robust optimization methods and variations of applying the group loss

6- In Section 3, “The training of an AI assistant consists of three main stages...“ is not necessarily the case for all AI Assistants. I suggest revising this statement and connecting the presented method to more broader usecases as it is not really limited to this case.
Recommendation: 8: accept, good paper
Reviewer confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness/correctness: 3 good
Clarity: 3 good
reviews: 
Comments: The paper introduces a method for improving LLMs' response quality on rare or difficult examples in the training data. Rather than treating all examples as equally important, the paper proposes to separate the data into groups so as to maximize the variance in quality among groups, then train the model to minimize that variance so that the model performs at roughly the same quality for all groups. Essentially this automatically adds additional training weight to difficult examples. The paper then presents empirical evidence that the proposed approach substantially improves response quality.
Strengths: The paper is well motivated, proposes a novel idea, describes it clearly, and presents compelling evidence that it works. In short, this is an excellent paper!

It also has immediate practical significance for training LLMs to produce higher quality outputs, and initial evidence suggests this would be more helpful and less harmful for users.

Figure 1 is an excellent visualization of the problem and the proposed solution.

Figure 4 clearly demonstrates that the proposed method is much more stable than vanilla PPO.
Weaknesses: My concerns are minor.

- Sec 4.3 para 1: "the changes *to which* the algorithm needs to remain invariant..."
- p6, Stage 2:
    - "binary groups": Is there an advantage to using just 2 groups? I would have expected more. Variance is really just a simple difference unless there are more than two groups
    - "narrows the performance gap between the two groups": Relative to PPO, yes, but it still stays flat or increases over time. Why is that?
- p7, Table 1: It would be helpful to know how statistically significant these results are. Is there some way to do that? It would also be helpful to add a calibration where a method is evaluated against itself. It should be a 100% tie, but obviously it will be more like 33/34/33.
- p9, ablation study: "it enables the model to explore a larger action space on challenging samples, leading to the discovery of a superior policy". How does it enable that?
- p9, reward distribution: "which will aid in its generalization to unseen data." Why will it do that? Did you show this somewhere?
Recommendation: 10: strong accept, should be highlighted at the conference
Reviewer confidence: 4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness/correctness: 4 excellent
Clarity: 4 excellent
reviews: 
Comments: The paper is motivated by the notion that maximizing only the expected return of a policy can be suboptimal when the return distribution has high variance or long tails. In Reinforcement Learning, the goal is to learn a policy that performs consistently across different data groups or scenarios. Here the group refers to different data distributions with varying performance characteristics. To achieve this, the authors propose incorporating a group inference classifier in the critic model to softly assign data points to groups. They maximize a variance objective to amplify the differences between groups. Additionally, they introduce an adaptive KL regularization scheme that allows different groups to have varied regularization strengths. This balances exploration and exploitation based on the group's difficulty. In summary, the paper aims to improve policy generalization by minimizing performance disparities between automatically identified data groups, while dynamically adapting the regularization to enable optimized exploration.
Strengths: pros: 
1. This paper introduced a framework for policy invariant learning that does not reply on prior domain or group knowledge, where labeling is inefficient.
2. They propose a novel dynamic KL penalty based on group labels discovered before. Easier groups get stricter constraints, potentially preventing overoptimization.
3. Good experimental results. 	The method demonstrates good experimental results, outperforming baselines on in-distribution data. This highlights its capabilities in known domains. On out-of-distribution data, it shows even greater performance gains, underscoring its generalization abilities.
4. Ablation studies validate that both core components - group invariant learning and adaptive KL - contribute to the overall performance gains.
Weaknesses: Cons: 
1. The experiments seem to imply a static number of groups, binary in the experiments,(best and challenging groups) which is not ideal with multi-modal data in dynamics environments. Experimenting with varying number of groups would be better to see how well does the method perform when the number of data groups increases?
2. The approach relies on accurate group label inference. Incorrect initial grouping could potentially lead to improper KL penalties. How robust is the method to errors in initial group assignment?
3. question: how do the newly introduced hyperparameters add complexity to tuning the algorithm? as training PPO is already known to be hard to tune.
Recommendation: 6: marginally above the acceptance threshold
Reviewer confidence: 3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness/correctness: 2 fair
Clarity: 3 good
reviews: 
Comments: The paper presents an innovative method addressing the challenge of 'reward hacking' and neglect of complex samples in AI assistants powered by language models, through Reinforcement Learning from Human Feedback (RLHF). The proposed technique advances a consistent policy learning across different data groups, enhancing the AI’s performance evenly across domains. It innovatively classifies data into groups to highlight performance variances, optimizes policy for difficult groups, and adaptively regulates exploration space, thus boosting training stability and generalization.
Strengths: 1. **Organization and Accessibility**: The paper is well-structured and the content is presented in a manner that is accessible to readers.

2. **Significance of Addressed Problems**: The authors tackle critical issues within RLHF, such as reward hacking and the overlooking of complex samples, which are pertinent for the advancement of universal AI assistants.

3. **Innovative Concept**: The application of group-invariant learning to the alignment problem is both a novel and a promising idea.
Weaknesses: 1. The explanation in Section 4.4 regarding the probability of assigning samples to the highest-performing group lacks clarity and warrants further detail.

2. The term $R_{g}(\theta)$ is not clearly defined within the paper. It’s assumed to represent the expected return of group $g$, yet its relationship to the last term in Equation 6 is ambiguous and needs clarification.

3. The first term of the final learning objective in Equation (8) does not appear to be directly related to group $g$, raising questions about its role in achieving invariant learning.

4. The paper would benefit from the inclusion of code or detailed pseudocode to clearly convey the training process, as the current description does not sufficiently outline the methodology.

5. The connection between group invariant learning and its ability to address issues of shortcut exploitation and neglect of challenging samples is not clearly articulated. More intuitive explanations or empirical evidence would strengthen the understanding of this relationship.
Recommendation: 6: marginally above the acceptance threshold
Reviewer confidence: 3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness/correctness: 3 good
Clarity: 3 good