---
title: 'ArmoRM+MoE: Interpretable Reward Modeling'
url: https://www.emergentmind.com/topics/armorm-moe
type: topic
---

# ArmoRM+MoE: Interpretable Reward Modeling

ArmoRM+MoE is a two-stage framework for interpretable, prompt-conditioned reward modeling in large language model (LLM) alignment. This approach addresses the opacity of conventional black-box reward models (RMs) in RLHF by making sub-scores for human-interpretable objectives explicit and enabling dynamic weight adjustment per prompt context. The method achieves near-state-of-the-art performance on the RewardBench benchmark using only an 8B-parameter Llama-3 backbone, while providing granular insight into model preferences and choices [2406.12845].

## 1. Absolute-Rating Multi-Objective Reward Model Construction

ArmoRM is trained on multi-dimensional absolute ratings instead of traditional pairwise preferences. For each prompt-response tuple $(x, y)$, a vector of absolute scores $r = [r_1, ..., r_k]^\top$ is provided, with each $r_i \in [0,1]$ representing a human-interpretable objective such as honesty, correctness, verbosity, safety, instruction-following, or code-readability. The datasets incorporated for label acquisition include HelpSteer, UltraFeedback, BeaverTails, CodeUltraFeedback, Prometheus, Argilla‐Capybara, Argilla-Math, among others. Collectively, ArmoRM sees 19 objectives with approximately 600,000 absolute-rating examples, where each task has its own rubric.

For model architecture, a frozen Llama-3 8B transformer is used as a feature extractor $f_\theta$. The concatenated prompt and response $(x \oplus y)$ pass through $f_\theta$, producing the final token hidden state $h \in \mathbb{R}^d$. A regression head $W \in \mathbb{R}^{d \times k}$ maps $h$ to sub-score predictions $\hat{y} = W^\top h \in \mathbb{R}^k$. The regression head $W$ is trained via mean squared error only on present labels for each example:
$$
\min_{W} \ \mathbb{E}_{(x, y, r)} \| W^\top f_\theta(x \oplus y) - r \|_2^2.
$$
This design enables explicit supervision for each interpretable axis available per example.

## 2. Mixture-of-Experts Gating and Prompt-Specific Scalarization

To integrate the k-dimensional output of ArmoRM into a scalar reward $R$ suitable for ranking or PPO, a fixed linear aggregation is insufficiently flexible. Instead, a Mixture-of-Experts (MoE) gating network $g_\phi$ dynamically selects a convex combination $v(x) \in \Delta^{k-1}$ over objectives based solely on the prompt feature $f_\theta(x)$. The gating network is a three-layer ReLU MLP with 1024 units per layer and a terminal softmax layer:
$$
v(x) = \mathrm{softmax}( \mathrm{MLP}_\phi(f_\theta(x)) ) \in \mathbb{R}^k
$$
with $\sum_i v_i = 1,\ v_i \geq 0$.

Substantial attention is given to verbosity debiasing, as many objectives correlate strongly with response length. For each target objective, the linear correlation with verbosity is removed:
$$
r'_i = r_i - \lambda_i \cdot r_\text{verbose}
$$
via $\lambda_i$ such that the debiased objectives are uncorrelated with verbosity on a held-out UltraFeedback reference set.

The scalar reward becomes
$$
R(x, y) = v(x)^\top r'(x, y)
$$
where $r'(x, y)$ contains the debiased sub-scores.

## 3. Two-Stage Training Protocol

Training proceeds in two decoupled phases:

1. **ArmoRM linear probing:** The 8B Llama-3 backbone is frozen and features $h = f_\theta(x \oplus y)$ are precomputed. The regression head $W$ is optimized using a multi-output least-squares regression solver (e.g., scikit-learn CPU backend).
2. **MoE gating optimization:** $f_\theta$ and $W$ remain fixed. The gating network $g_\phi$ is trained with a Bradley-Terry loss using pairwise comparisons from ten preference datasets. Loss is given by
$$
\min_{\phi, \beta} -\mathbb{E}\left[ \log\left( \frac{\exp(\beta R_\text{chosen})}{\exp(\beta R_\text{chosen}) + \exp(\beta R_\text{rejected})} \right) \right]
$$
with scalar temperature $\beta$ initialized to 100. Training proceeds for 10,000 steps on an A6000 GPU using AdamW (learning rate $1 \times 10^{-3}$, batch size 1024, cosine learning rate decay).

## 4. Empirical Performance on RewardBench

RewardBench tests reward models on their ability to correctly rank preferred over rejected responses in five categories: Chat, Chat-Hard, Safety, Reasoning (each with weight 1.0), and Prior-Sets (weight 0.5). Major comparative results:

| Model                                      | RewardBench Weighted Accuracy (%) |
|--------------------------------------------|-----------------------------------|
| Nemotron-4 340B (HelpSteer2 RM)            | 89.3                              |
| ArmoRM + MoE (Llama-3 8B)                  | 89.0                              |
| HelpSteer2 RM on Llama-3 70B               | 86.3                              |
| Bradley-Terry RM on Llama-3 8B (backbone)  | 83.6                              |
| LLM-as-a-judge (GPT-4 Turbo)               | 84.2                              |
| LLM-as-a-judge (GPT-4o)                    | 83.3                              |

ArmoRM+MoE matches the 340B baseline on Safety and Prior-Sets and excels at Chat and Reasoning sub-tasks. Notably, the performance of ArmoRM+MoE (8B) significantly exceeds the LLM-as-a-judge paradigm as implemented with GPT-4 judges.

## 5. Interpretability and Practical Auditing

The MoE gating network outputs a per-prompt vector $v(x)$ revealing the weight assigned to each human-centric objective. This design enables direct auditing and interpretability:

- On safety-critical prompts, $v(x)$ assigns 70% or more weight to "is-safe," with minimal allocation to verbosity.
- For mathematical questions, the mass shifts to correctness, truthfulness, and instruction-following.

Practitioners can inspect $v(x)$ to diagnose unintended emphasis (such as over-weighting verbosity), and optionally steer model behavior by clamping specific weights to zero. This provides a direct mechanism to check alignments, investigate model errors, and mitigate reward hacking.

## 6. Limitations and Open Questions

Several limitations and areas for future work are identified:

- **Non-joint training:** The regression head $W$ and gating network $g_\phi$ are learned sequentially; joint fine-tuning might improve overall alignment.
- **Handling of missing labels:** Many data points lack complete coverage of all $k$ objectives; current training ignores missing dimensions without imputation. A principled approach to missingness may enhance robustness and interpretability.
- **Prompt-only gating:** The gating network considers only the prompt, not the response. For some objectives (e.g., factuality), conditioning on the response may be beneficial.
- **Static debiasing:** Verbosity de-correlation uses fixed $\lambda_i$ and a reference distribution $D$; adaptivity to novel domains is not addressed.
- **Human oversight effect:** While $v(x)$'s transparency makes the model auditable, no user studies confirm improvement in human oversight or reduction of reward hacking.

A plausible implication is that future research incorporating joint fine-tuning, dynamic debiasing strategies, or response-aware gating may further close the performance gap to larger models and enhance trust in RM outputs.

## 7. Summary and Context

ArmoRM+MoE operationalizes interpretable multi-objective reward modeling by decoupling axes of human preference and introducing prompt-specific, auditable scalarization. The framework demonstrates that a relatively compact 8B model suffices for high reward modeling accuracy, rivaling massive baselines while supporting fine-grained scrutiny of alignment. These results position ArmoRM+MoE as a candidate methodology for transparent and trustworthy LLM reward modeling, encouraging further work on joint learning, adaptive debiasing, and human-centered evaluation [2406.12845].

Source: https://www.emergentmind.com/topics/armorm-moe