#1minPapers “Critique-out-loud Reward Models” — by Zachary Ankner, Mansheej Paul, Brandon Cui, Jonathan D. Chang, Prithviraj Ammanabrolu

I’m down the rabbit hole of optimized reward models. This paper is still in preprint.

Key takeaway: Adding a self-consistency, generated critique, to a classic reward model does improve pairwise preference classification accuracy, particularly for safety, reasoning, and shallow-step tasks. Simple inference compute to slightly lift accuracy.

Context: Augment variable inference compute in CoT (Chain-of-Thought) — combine classic reward models (usually with linear or shallow multi-layer perceptrons MLPs) and LLM-as-a-Judge.

Architecture: (1) Supervise finetune SFT a pretrained LLM to generate critiques for the responses by finetuning on the oracle critiques (used Llama-3.1–405B-Instruct). (2) Rebuild the dataset by replacing oracle critiques with critiques generated by the finetuned model. (3) Predict a scalar reward head (Bradley-Terry preference model) on top of the finetuned model, and train on the new dataset composed of self-generated critiques, to minimize both a language modeling and a preference modeling loss.

Model: Used both Llama-3–8B and 70B, each trained with 2 random seeds. Hyperparamater sweep of learning rate and number of training epochs, additional reward model for the SFT loss weight sweep. Warmup duration 5%. Final decay factor 1%. Adam optimizer with decoupled weight decay.

Data: UltraFeedback and UltraInteract. Pairwise preference classifications via RewardBench. BoN win rate performance viaArenaHard.

Result: Improved pairwise preference classification accuracy on RewardBench by 4.65% and 5.84% for the 8B and 70B base models.

Optimizations: Explored stochasticity to improve critique generation via multi-sample inference techniques.

Ablation showed that on-policy training (vs off-policy) is essential for the success of Critique-out-loud CLoud reward models for both preference classification and for Best-of-N BoN. For both model sizes, training in an off-policy manner leads to a significant drop in performance. Importance of matching the distribution of critiques seen by the reward head during training and inference.

Self-consistency (Wang et al., 2023) is an inference technique for computing the maximum marginal likelihood answer, by sampling many reasoning-answer-tuples. Simple method to improve performance at the cost of added inference compute.

On the reasoning category, adding self-consistency is an effective method to trade added inference compute for increased performance.

Self-consistency leads to an improvement in preference classification accuracy of 0.70% and 0.49% for the 8B and 70B base models. No gain in self-consistency for non-reasoning categories.

Prompts that require only 1–2 reasoning steps (shallow-steps) saw a significant lift in performance as the number of critiques increased. All other groups of reasoning steps actually degrade in performance past 8 critiques. Self-consistency with CLoud reward models are best for short horizon tasks.

Resource:

Github code and paper https://github.com/zankner/CLoud