AdaJudge: Adaptive Multi-Perspective Judging for Reward Modeling

AI summary

Overview

Research area: Natural Language Processing, specifically reward modeling for large language model (LLM) alignment with human preferences.

Technical level: Advanced. The paper assumes familiarity with transformer architectures, Bradley-Terry preference modeling, pooling operations, mixture-of-experts routing, and RLHF-style training objectives.

Scope: The paper proposes AdaJudge, a two-stage reward modeling framework that replaces the fixed pooling head used by standard reward models with learned representational refinement plus a prompt-conditioned mixture of three pooling "experts," evaluated on the RM-Bench and JudgeBench benchmarks.

What This Paper Is About

Standard reward models compress a prompt-response sequence into a single scalar score using one fixed pooling rule, usually the last token's hidden state. The authors argue this creates two mismatches: the pooling rule's fixed spatial bias cannot fit the fact that different tasks place their verifying evidence in different places (e.g., a final answer versus a holistic style judgment), and the backbone's hidden states were optimized for next-token generation rather than fine-grained preference discrimination. AdaJudge's goal is to jointly fix both problems by refining the backbone representations and then adaptively combining multiple pooling perspectives per input.

Key Contributions

1. Identification of two structural mismatches in reward modeling: generative-focused backbones lack the fine-grained features needed for pairwise discrimination, and fixed pooling cannot capture the spatially diverse, multi-granular distribution of preference evidence.

2. The AdaJudge framework, a unified two-stage design that combines depth-gated representation refinement (Stage I) with a domain-aware gated mixture-of-pooling head (Stage II), aligning backbone representations with task-specific evaluation criteria without token-level or process supervision.

3. Empirical validation on RM-Bench and JudgeBench showing AdaJudge outperforms both strong off-the-shelf reward models and controlled same-backbone baselines across domains and model scales.

4. Internal mechanism analysis showing that refinement increases the cosine alignment between the model's scoring direction and true preference-difference directions, and that different domains elicit systematically different routing weights over the pooling experts.

Main Findings

• Beating a 27B model with an 8B backbone: On RM-Bench overall, Qwen3-8B with traditional fixed pooling scores 67.3, while the same backbone with AdaJudge scores 71.1, surpassing the 27B-parameter Skywork-Reward-Gemma-2-27B at 70.5. AdaJudge also reaches 66.0 overall on JudgeBench with Qwen3-8B, versus 63.1 for last-token pooling and 63.4 for mean pooling on that backbone.

• Relieving the bottleneck in compact models: For Phi-3.5-mini-instruct, mean pooling scores 48.7 on RM-Bench Math and 50.7 on Code, while last-token pooling scores 55.1 overall on JudgeBench. AdaJudge lifts the same backbone to 59.8 on RM-Bench and 59.4 on JudgeBench.

• Resolving conflicting inductive biases in Qwen3-8B: Mean pooling leads on Safety (82.6 vs. 79.9 for last-token), while last-token pooling scores only 34.8 on the RM-Bench Hard subset. AdaJudge reaches 80.4 on Math, 87.5 on Safety, and 43.0 on the Hard split within a single architecture.

• Adaptive aggregation beats every static readout (Qwen3-4B, Table 2): Last-token pooling gets the best Safety score (84.7%) but weaker Math (65.0%); mean pooling improves Math (68.0%) but drops Safety (77.0%); attention pooling scores 38.7% on Hard. AdaJudge achieves the best overall results on RM-Bench (70.8%) and JudgeBench (66.0%), and on the Hard subset scores 43.7%, exceeding the best static alternative by 5.0%.

• Refinement matters most on reasoning-heavy tasks (Qwen3-4B, Table 3): Removing Stage I refinement drops JudgeBench overall from 66.0% to 62.6% (a 3.4% improvement from refinement), with Math gaining 5.4% and Code gaining 4.8%. On RM-Bench, both variants are comparable on Easy (approximately 92%), but refinement raises Hard accuracy from 39.6% to 43.7%.

• Refinement makes evidence globally accessible: Cosine alignment between the scoring direction and the preference difference improves across domains after refinement, with the largest gains on Chat and Safety (RM-Bench) and Math and Code (JudgeBench). Gains concentrate on mean and attention pooling views, while last-token alignment remains flat or degrades.

• Routing is domain-dependent: On Code and Math, AdaJudge assigns a larger proportion of weight to the Attention pooling expert; the Safety domain shifts toward Mean pooling; the Last-token expert receives a relatively small average weight across all domains.

• Training overhead is bounded: Training AdaJudge on Qwen3-8B peaks at 32.99 GB of memory and 10.80 hours on a single NVIDIA RTX Pro 6000 GPU, versus 24.52 GB and 7.51 hours for last-token and mean pooling baselines. Throughput drops to 4.46 × 10^5 TFLOP/h from 5.30 × 10^5 for the static baselines.

• Setup: All reward models are trained on the preference split of HelpSteer3, containing about 40.5K preference pairs, using LoRA on Phi-3.5-mini-instruct, Qwen3-4B, and Qwen3-8B. RM-Bench contains 1.33k test instances; the paper's truncated content does not report JudgeBench's dataset size.

Methodology in Plain English

The authors keep the existing reward model backbone and change what happens on top of it.

Stage I — Refining the representations. After the LLM produces its usual token-level hidden states, AdaJudge stacks K small extra transformer blocks on top (K=2 for Phi-3.5-mini-instruct, K=3 for Qwen3-4B and Qwen3-8B). Rather than just using the last block's output, it computes a global context vector from the sequence by mean pooling, feeds that through a small network, and gets a set of mixture weights via softmax. The final refined representation is a convex combination of all the intermediate block states, so an easy input can lean on shallow features while a hard one can lean on deeper reasoning layers.

Stage II — Combining multiple views. From the refined representation, AdaJudge extracts three different summaries: the hidden state at the final response token, the masked mean over all response tokens, and an attention-weighted sum over all tokens with learned attention scores. It also computes a prompt-only context vector by mean pooling over the prompt tokens. Each of the three response summaries goes through its own MLP to produce a scalar score. A routing network takes the concatenation of all four vectors and outputs three mixture weights via softmax, and the final reward is the weighted sum of the three scalar scores. Because the routing is conditioned on the prompt context, the model can infer the task's evaluation intent and pick the pooling bias that fits it.

Training. The objective is a focal Bradley-Terry loss: the standard pairwise preference probability is weighted by a term derived from the preference magnitude (square-root scaling) and a focal factor that emphasizes hard samples. An entropy regularizer penalizes routing distributions that collapse toward a single expert. The paper's training configuration uses a maximum sequence length of 4096 tokens, LoRA rank 96, LoRA scaling factor 128, LoRA dropout 0.05, gradient norm clipping at 2.0, AdamW with β1=0.9 and β2=0.95, weight decay 0.1, and a constant learning rate schedule with linear warmup covering 3% of total training steps, with checkpoints every 25 optimization steps. Per-backbone settings vary: Qwen3-4B uses learning rate 4e-5, focal temperature 1.3, focusing parameter 0.9, entropy coefficient 0.01, and target gate entropy 0.7; Phi-3.5-mini-instruct uses 2e-5, temperature 1.2, focusing parameter 0.5, entropy coefficient 0.01, and target entropy 0.7; Qwen3-8B uses 3e-5, temperature 1.2, focusing parameter 0.7, entropy coefficient 0.003, and target gate entropy 0.65. All runs use an effective batch size of 8 preference pairs.

Why This Matters

Impact on research. The paper argues that the scoring head of a reward model — long treated as a rigid afterthought while effort went into data scale and backbone size — is a real structural bottleneck. It provides an alternative to expensive inter-model ensembles (which the paper notes incur prohibitive computational overhead) by internalizing multiple complementary scoring perspectives inside a single model, and it offers a diagnostic methodology (measuring alignment between the scoring gradient and the preference difference) for studying representation quality in reward models.

Real-world applications:

• RLHF and LLM alignment pipelines, where a more discriminative reward model directly improves the preference signal used to fine-tune policy models.
• LLM-as-a-judge evaluation systems, since JudgeBench-style benchmarks directly measure how reliably a model can adjudicate between competing responses.
• Safety filtering and content moderation, where the paper shows adaptive routing can emphasize the holistic, distributed evidence that safety judgments require rather than focusing on terminal tokens.
• Reasoning-task verification in math and code, where the paper shows the attention pooling expert receives more weight and refinement produces the largest gains.
• Resource-constrained deployment, since the method improves compact models rather than only large ones — the paper notes peak memory remains within modern GPU capacity, though latency overhead is acknowledged.

Industry relevance. The result that an 8B backbone with AdaJudge (71.1 on RM-Bench overall) surpasses a 27B reward model (70.5) is a direct cost argument: it suggests parameter efficiency through architectural design can substitute for brute-force scaling or massive data curation. The method trains on a single dataset of about 40.5K preference pairs, far smaller than the large-scale preference corpora used by the off-the-shelf baselines.

Future Directions

• Extending to multi-objective reward modeling, letting the routing mechanism better disentangle trade-offs between criteria such as factuality and helpfulness.
• Scaling validation beyond 8B parameters, with the authors specifically calling out extreme-scale models (e.g., 70B+) to test whether dynamic routing benefits persist as backbone capacity grows.
• Reducing latency and parameter overhead, potentially by distilling the adaptive capabilities into linear heads, since the iterative refinement blocks and routing network add cost relative to fixed-pooling baselines.
• Broadening the expert library and input modality, including learned convolutional filters, hierarchical pooling, task-specific experts, automatic expert discovery, multimodal inputs, cross-modal reward discrimination, and multi-objective reward modeling. The paper also notes AdaJudge has so far only been evaluated on text-only preference modeling.

Target Audience

Researchers and engineers working on LLM alignment, reward modeling, and RLHF who want to understand where the scoring head's architecture limits performance and how adaptive pooling can be designed. It is most useful to readers already comfortable with transformer internals, Bradley-Terry preference objectives, and benchmark-based evaluation — and to practitioners building reward models under compute or parameter constraints who want evidence that head design can matter as much as backbone scale. Readers looking for an introductory treatment of preference learning will find the paper assumes substantial background.

Authors’ abstract

Reward modeling is essential for aligning large language models with human preferences, yet predominant architectures rely on a static pooling strategy to condense sequences into scalar scores. This paradigm, however, suffers from two key limitations: a static inductive bias that misaligns with task-dependent preference signals, and a representational mismatch, as the backbone's optimization for generation leaves its representations ill-suited to fine-grained discrimination. To address this, we propose AdaJudge, a unified framework that jointly adapts representation and aggregation. AdaJudge first improves backbone representations into a discrimination-oriented space via gated refinement blocks. It then replaces the static readout with an adaptive multi-view pooling module, which dynamically routes and combines evidence. Extensive experiments on RM-Bench and JudgeBench show that AdaJudge outperforms strong off-the-shelf reward models and traditional pooling baselines.