Diffusion Reward Models
Abstractの概要
本論文では、報酬モデリングを単一のスカラー値の予測ではなく、プロンプトと応答のペアが与えられた際の報酬に関する条件付き密度推定として再構築するDiffusion Reward Models（DRM）を提案します。この手法は、凍結されたLLMエンコーダと軽量なDiffusion Transformerヘッドを用いて、特定のパラメトリックな分布族を仮定することなく報酬分布をモデル化し、同一のフレームワークで多属性回帰とペアワイズ嗜好学習の両方をサポートします。推論時において、DRMは平均、分散、または分位数へと集約可能な複数の報酬サンプルを生成し、標準的なスカラー値スコアリングと分布を考慮した決定ルールの両方を可能にします。評価は、5つの公開報酬モデルベンチマーク、人間の評価のばらつきに関する分析、および下流のRLHF実験にわたって行われています。
新規性
主な新規性は、従来のスカラー型や固定分布族を仮定したパラメトリックな報酬ヘッドを、完全な条件付き報酬分布をノンパラメトリックにモデル化する拡散ベースのヘッドに置き換えた点にあります。本フレームワークは、多属性回帰とペアワイズ嗜好データの両方に統一的なアーキテクチャを提供するとともに、逆拡散サンプリングの反復による報酬軸でのテスト時スケーリング機構を導入しています。
成果
同一のデータおよびバックボーン条件下において、DRM-Multi-8Bはベンチマーク平均スコア66.2を達成し、ArmoRM（62.3）やQRM-Llama3.1-8B-v2（64.1）を上回りました。分布評価において、DRMはベースラインよりも人間の評価分布と密接に一致しており、不確実性を考慮した棄却や下限信頼限界による集約が意思決定精度を向上させました。下流のRLHFの学習報酬として使用した場合、DRM-MultiはFsfairXベースラインと比較して、Arena-Hard v2のポリシー性能を1.3から2.0へ、MT-Benchを73.6から74.8へと改善しました。
論文の注目点
- DRMは、凍結されたLLMエンコーダ上のDiffusion Transformerヘッドを用い、報酬をノンパラメトリックな条件付き分布としてモデル化する。
- 多属性の教師あり学習とペアワイズ嗜好学習の両方の目的に対応し、複数の逆拡散サンプルを通じた報酬軸でのテスト時スケーリングを可能にする。
- 5つのRMベンチマーク、人間の評価不一致の分析、および下流のRLHFにわたる評価において、同一条件のベースラインに対する優位性と効果的なリスク考慮型集約が実証された。
参考リンク
- arXiv: https://arxiv.org/abs/2609.33803v1
- Fugu-MT: https://fugumt.com/fugumt/paper_check/2609.33803v1
- GitHub: https://github.com/thunlp/DRM
- Hugging Face: https://huggingface.co/Teburile/DRM