Alternating Reinforcement Learning with Contextual Rubric Rewards

1. Problem Statement

The paper addresses a critical limitation in Reinforcement Learning with Rubric Rewards (RLRR). While RLRR improves upon standard RLHF and RLVR by using structured, multi-dimensional evaluations (rubrics) rather than single scalar scores, existing methods suffer from linear scalarization.

• The Issue: Current approaches compress vector rewards (e.g., accuracy, completeness, instruction following) into a single scalar value using fixed linear weightings.
• Consequences:
  • Information Loss: This aggregation dampens the variance and entropy of feedback signals, obscuring the nuanced differences between high-quality and mediocre responses.
  • Reward Hacking: The policy may optimize for "easier" metrics (e.g., completeness) to maximize the scalar score while neglecting harder, critical metrics (e.g., accuracy).
  • Sensitivity: The performance is highly sensitive to the artificial design of weights and fails to capture non-linear correlations between different reward dimensions.

2. Methodology: ARL-RR

The authors propose Alternating Reinforcement Learning with Rubric Rewards (ARL-RR), a framework that eliminates fixed scalarization by optimizing one semantic rubric meta-class at a time.

Core Components:

1. Meta-Class Decomposition:

   • Instead of treating all rubric criteria as a single vector, the framework partitions criteria into fixed semantic meta-classes: Accuracy, Completeness, Instruction Following, Context Awareness, and Communication Quality.
   • For a given query $q$ and response $a$, the reward $r_m$ for a specific meta-class $m$ is calculated as the weighted average of only the criteria belonging to that class.

2. Alternating Optimization Process:

   • The training proceeds in stages (epochs). In each stage, the policy $\pi_\theta$ is optimized using the reward signal from only one active meta-class.
   • The advantage estimation in the objective function (based on GRPO) is calculated solely based on the distribution of rewards within that specific meta-class.
   • This forces the model to learn the nuances of each dimension explicitly, preventing the masking of failures in one area by success in another.

3. Search-Based Adaptation (Meta-Class Searching):

   • To determine the optimal order of meta-classes, the authors introduce a lightweight search strategy.
   • The model is trained on a subset of data ($p\%$) across different meta-class orders in parallel.
   • The checkpoint with the best performance on the subset dictates the next meta-class to be trained on the full dataset. This dynamically prioritizes critical objectives based on current task performance.

4. Synthetic Meta-Classification:

   • To handle scenarios where expert-labeled meta-classes are unavailable, the paper introduces a prompt-based method to automatically classify rubric criteria into the fixed meta-classes using an LLM.

3. Theoretical Analysis

The paper provides a theoretical justification for why scalarization is detrimental, termed Variance Contraction.

• Theorem 3.1: Under mild correlation assumptions (where meta-class rewards are not perfectly correlated, $\rho < 1$), the variance of a scalarized reward $R$ is strictly less than the variance of any individual meta-class reward $R_m$.
• Implication: Scalarization acts as a "compressor" that smooths out fluctuations. In RL, high variance in rollout rewards is crucial for distinguishing strong responses from weak ones (providing a strong learning signal). By reducing variance, scalarization weakens the advantage estimation, leading to slower or less effective learning. ARL-RR preserves this variance by isolating dimensions.

4. Key Contributions

1. Novel Framework (ARL-RR): A training method that bypasses fixed linear weighting, preserving the structure of multi-dimensional evaluation criteria.
2. Meta-Class Pipeline: A systematic approach to map contextual rubrics into fixed semantic categories, enabling consistent training across diverse tasks.
3. Dynamic Adaptation: A search-based strategy to dynamically select the active meta-class, transforming a multi-objective problem into a sequence of focused sub-problems.
4. Theoretical Insight: Proof that linear scalarization induces variance contraction, explaining the performance degradation in scalarized RL.
5. Empirical Validation: Comprehensive experiments demonstrating superior performance and efficiency across multiple model scales.

5. Experimental Results

Experiments were conducted on the HealthBench dataset (5K samples, expert-annotated rubrics) using Qwen3 models ranging from 1.7B to 14B parameters.

• Performance Gains: ARL-RR uniformly outperformed scalarized RL (SRL) baselines across all model sizes.
  • Example: On Qwen3-14B, ARL-RR achieved a score of 0.788 vs. 0.762 for SRL.
  • Example: On Qwen3-1.7B, ARL-RR achieved 0.576 vs. 0.556 for SRL.
• Training Efficiency: ARL-RR reached higher evaluation scores in fewer epochs compared to scalarized methods.
• Reward Variance: Empirical analysis confirmed that scalarized rewards had significantly lower variance (0.10) compared to individual meta-class rewards (e.g., 0.20 for "Instruction Following"), validating the variance contraction theory.
• Robustness:
  • The method remained effective across different reward models (Qwen3-4B to 32B).
  • Synthetic Meta-Classification showed only a marginal performance drop compared to expert labels, proving the method's applicability in real-world scenarios without ground-truth meta-class labels.
• Search Strategy: The search-based order selection improved performance, with a data subset of ~20% ($p=20\%$) being sufficient to find an effective training order.

6. Significance

This work fundamentally shifts the paradigm of post-training alignment for Large Language Models (LLMs).

• Beyond Scalarization: It demonstrates that moving away from scalar aggregation is not only practical but necessary for handling complex, context-dependent human expectations.
• Safety and Reliability: By preventing the "masking" of critical failures (e.g., hallucinations or safety violations) via easy-to-optimize metrics, ARL-RR fosters the development of more robust and trustworthy AI systems.
• Generalizability: The framework offers a principled approach to multi-objective optimization that can be applied to various domains beyond healthcare, provided rubric-based evaluations are available.

In conclusion, ARL-RR offers a more flexible, principled, and effective method for aligning LLMs with complex rubric-based rewards, achieving state-of-the-art results in both model performance and training efficiency.