DM o ERM : Recipes of Mixture-of-Experts for Effective Reward Modeling (2024)

Abstract: The performance of the reward model (RM) is a critical factor in improving the effectiveness of the large language model (LLM) during alignment fine-tuning. There remain two challenges in RM training: 1) training the same RM using various categories of data may cause its generalization performance to suffer from multi-task disturbance, and 2) the human annotation consistency rate is generally only 60% to 75%, causing training data to contain a lot of noise. To tackle these two challenges, we introduced the idea of Mixture-of-Experts (MoE) into the field of RM for the first time. We propose the Double-Layer MoE RM (DMoERM). The outer layer MoE is a sparse model. After classifying an input into task categories, we route it to the corresponding inner layer task-specific model. The inner layer MoE is a dense model. We decompose the specific task into multiple capability dimensions and individually fine-tune a LoRA expert on each one. Their outputs are then synthesized by an MLP to compute the final rewards. To minimize costs, we call a public LLM API to obtain the capability preference labels. The validation on manually labeled datasets confirms that our model attains superior consistency with human preference and outstrips advanced generative approaches. Meanwhile, through BoN sampling and RL experiments, we demonstrate that our model outperforms state-of-the-art ensemble methods of RM and mitigates the overoptimization problem. Our code is available at: https://github.com/quanshr/DMoERM.

## introduction after an initial stage of pre-training and subsequent instruction fine-tuning @xcite , large language models (llms) undergo a crucial stage of high-quality alignment fine-tuning based on reinforcement learning with human feedback (rlhf) to improve their abilities @xcite . during the rlhf process, a reward model (rm) often needs to be trained, which acts as a proxy of human preferences and assigns scores to the outputs of the llm. the scores are then used as reward signals to optimize the llm through reinforcement learning (rl). in this process, the llm and rm are interdependent and iteratively optimized, and the rm is expected to be highly consistent with human preferences. in addition, during the inference stage, the rm can also be augmented with best-of-n (bon) sampling strategies to further enhance the quality of the outputs of the llm @xcite . training of reward models relies on the data derived from human annotators who manually rank the varying outputs under a single input by their preferences. however, many studies have found that agreement rates among human annotators typically only range between 60-75% @xcite @xcite , thereby introducing a minimum of 25% noise within the labeled dataset. one important reason for this phenomenon is the multifaceted nature of evaluation-it is often observed that one response may excel in one aspect while simultaneously falling short in another. this multifaceted evaluation conundrum has been exemplified in previous studies @xcite @xcite , which illustrate the inherent tensions between enhancing helpfulness and harmlessness. as these attributes can at times be inversely related, adjudicating between a response that is more helpful yet potentially less harmless poses a significant challenge for comparative assessment. we further validate this perspective through empirical studies. in this study, we pioneer the integration of the mixture-of-experts (moe) framework @xcite @xcite into reward modeling. our approach employs a double-layer moe architecture. the outer layer comprises a sparse moe model specifically designed to avoid multi-task disturbance @xcite . as shown in figure 1 , we categorize inputs into several distinct tasks and use a pre-trained router to route the inputs to their corresponding task-specific expert. this strategy can also facilitate distributed deployment and enhance the model's capacity and capabilities without a commensurate increase in computational demands @xcite . subsequently, within each inner layer lies a dense moe model, which is tailored to the specific capabilities set required for its category. for instance, in roleplay scenarios, we divide it into six core capabilities, including personality and emotional investment, conversational sense, empathy ability and so on (details are in appendix ## result 6.1 training phases in this section, we use roleplay as an example to demonstrate the progress of the model at different training phases. by recording the accuracy of the reward model's ranking on the pairs of responses in training and testing sets, we can intuitively display the training results. the results are shown in figure 4 . we find: compared with the 56% accuracy in phase 1, it can be proven that our method of training multiple experts based on different capability points and aggregating those experts can significantly improve the model's performance.