Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Exploring Domain Robust Lightweight Reward Models based on Router Mechanism

Hyuk Namgoong, Jeesu Jung, Sangkeun Jung, Yoonhyung Roh

TL;DR

Problem: RLHF reward models trained for multiple domains can underfit domain-specific preferences and require retraining for new domains. Approach: three router-based schemes—MoRE, RODOS, and ARLISS—to construct domain-robust, lightweight reward models; Key contributions: (i) MoRE with an internal $top-k$ sparse MoE router, (ii) RODOS with an external domain router, (iii) ARLISS with adapter-based switching on a single LM using LoRA/PEFT; Findings: on five reward domains, the methods generally outperform a single baseline, with RODOS achieving the best accuracy, and MoRE/ARLISS reducing total parameters by about 52–55%; Impact: enables efficient, scalable RLHF fine-tuning on constrained hardware.

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.

Exploring Domain Robust Lightweight Reward Models based on Router Mechanism

TL;DR

Problem: RLHF reward models trained for multiple domains can underfit domain-specific preferences and require retraining for new domains. Approach: three router-based schemes—MoRE, RODOS, and ARLISS—to construct domain-robust, lightweight reward models; Key contributions: (i) MoRE with an internal sparse MoE router, (ii) RODOS with an external domain router, (iii) ARLISS with adapter-based switching on a single LM using LoRA/PEFT; Findings: on five reward domains, the methods generally outperform a single baseline, with RODOS achieving the best accuracy, and MoRE/ARLISS reducing total parameters by about 52–55%; Impact: enables efficient, scalable RLHF fine-tuning on constrained hardware.

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.
Paper Structure (21 sections, 1 figure, 7 tables)

This paper contains 21 sections, 1 figure, 7 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Router Based Switching Reward Models
  4. Mixture of Reward Experts
  5. Router for Domain-Specific Reward Models
  6. Adapter Router Lightweight Integrated Rewards Switching Framework
  7. Experiments Setup
  8. Datasets
  9. Language Models
  10. Baseline Methods
  11. Evaluation Metric for Reward Model
  12. Experimental Results
  13. Reward Models Performance
  14. Training Time
  15. Total Parameter Size
  16. ...and 6 more sections

Figures (1)

  • Figure 1: Illustration of each method. RM represents the reward model, and reward is scalar. $n$ denotes the number of domains, and Mixture of Reward Experts (MoRE) refers to Sparse Mixture of Experts with $k$ equal to 2. Router for DOmain-specific reward modelS (RODOS) involves loading all models for use, while the Adapter Router Lightweight Integrated rewardS Switching (ARLISS) framework loads only router and reward model adapters and a single language model, using adapter switching within the same language model.