Exploring Domain Robust Lightweight Reward Models based on Router Mechanism (2024)

Abstract: Recent advancements in large language models have heavily relied on the large reward model from reinforcement learning from human feedback for fine-tuning. However, the use of a single reward model across various domains may not always be optimal, often requiring retraining from scratch when new domain data is introduced. To address these challenges, we explore the utilization of small language models operating in a domain-specific manner based on router mechanisms. Our three approaches are: 1) utilize mixture of experts to form a single reward model by modularizing an internal router and experts, 2) employing external router to select the appropriate reward model from multiple domain-specific models, and 3) the framework reduces parameter size by loading reward models and router adapters onto a single small language model using adapters. Experimental validation underscores the effectiveness of our approach, demonstrating performance comparable to baseline methods while also reducing the total parameter size.

## introduction most widely adopted large language models (llms) have used the reward model of reinforcement learning from human feedback (rlhf) @xcite for fine-tuning. these reward models are trained from various human feedback domains and are subsequently utilized as evaluation metrics during llm fine-tuning processes. however, training a single reward model across various domains to serve multiple purposes may lead to situations where the model is not fit for specific domains. additionally, there is a challenge of retraining the reward model from scratch when new dataset from a new domain is introduced. in this paper, we explore various router methods to address these challenges, as summarized in to validate our methodologies, we conducted experiments with five different domains of reward datasets. in this experiment, our methods generally outperform the baseline, while rodos shows the best performance. more showcases a size reduction of about 52%, while arliss achieves a reduction of approximately 55% compared to the baseline. ## related works recent research focuses on improving llms @xcite @xcite training efficiency. introducing the reward model serves to evaluate llm performance in the rlhf fine-tuning method. in @xcite , the reward model spans various domains, while @xcite applies the rlhf method to image generation. research has explored methods for routing language models, such as routing llms @xcite @xcite . furthermore, various studies are underway to modularize and utilize routers within models @xcite @xcite , with mixture of experts (moe) @xcite being one. research on efficient fine-tuning of language models is ongoing. low-rank adaptation (lora) @xcite attaches adapters to each layer and updates only the adapter parameters, enabling efficient learning. building upon lora, further research explores efficiency improvements @xcite @xcite and additional tasks @xcite @xcite . we do not train a single reward model across diverse domains. instead, we utilize adapters to construct multi-reward models and routers, employing a small language model with lora, thereby reducing training time and parameters. ## router based switching reward models the reward model assigns rewards to prompt and response. in rlhf, the reward model's loss function calculates the difference between the rewards for the chosen and rejected responses. reward model dataset has the structure of one input prompt and least two of responses. these reward models cover diverse domains like human preferences and toxic responses, using largescale models. however, relying solely on one large model may not suit specific domains, and training from scratch for new domains takes time. ## experimental results our study investigates the effectiveness of the proposed router methods through experimental analyses, focusing on key aspects: evaluating the router's impact on application performance, analyzing training time across methods, and comparing total parameters with and without arliss integration. ## conclusion in addressing the limitations of a single large reward model, which can be unsuitable for specific domains and requires retraining when new domain data is introduced, we have implemented router methods. more features an internal router alongside a single small reward model, while rodos incorporates an external router and domain-specific reward models. these methods effectively mitigate challenges related to domain specificity and the need for retraining when new domain data is introduced. moreover, the arliss framework, with adapters for routers and multi-reward models, shows potential for gpu memory optimization by reducing model size. further research will focus on optimizing the arliss framework. additionally, we plan to investigate the integration of the arliss framework into more.