URM LLaMa 3.1 8B

🚀 URM-LLaMa-3.1-8B

URM-LLaMa-3.1-8Bは、不確実性を考慮した報酬モデルです。このモデルは、ベースモデルと不確実性を考慮した属性固有のバリューヘッドから構成されており、ベースモデルはSkywork-Reward-Llama-3.1-8Bを使用しています。

🚀 クイックスタート

このモデルは、不確実性を考慮した報酬モデルであり、特定のデータセットを用いて2段階のトレーニングを行っています。以下に、モデルの概要や使用方法を説明します。

✨ 主な機能

• 不確実性を考慮した報酬モデルで、不確実性の低い報酬はより信頼性が高く、より良いアライメント結果をもたらします。
• 2段階のトレーニングを行うことで、属性回帰とゲーティングレイヤーの学習を行っています。

📦 インストール

本READMEにはインストール手順が記載されていないため、このセクションをスキップします。

💻 使用例

基本的な使用法

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_name = "LxzGordon/URM-LLaMa-3.1-8B"
model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    device_map='auto',
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "What is the range of the numeric output of a sigmoid node in a neural network?"
response1 = "The output of a sigmoid node is bounded between -1 and 1."
response2 = "The output of a sigmoid node is bounded between 0 and 1."

resp1 = [{"role": "user", "content": prompt}, {"role": "assistant", "content": response1}]
resp2 = [{"role": "user", "content": prompt}, {"role": "assistant", "content": response2}]

# Format and tokenize the conversations
resp1 = tokenizer.apply_chat_template(resp1, tokenize=False)
resp2 = tokenizer.apply_chat_template(resp2, tokenize=False)
resp1 = tokenizer(resp1, return_tensors="pt").to(model.device)
resp2 = tokenizer(resp2, return_tensors="pt").to(model.device)

with torch.no_grad():
    score1 = model(resp1['input_ids'],attention_mask=resp1['attention_mask']).logits[0][0].item()
    score2 = model(resp2['input_ids'],attention_mask=resp2['attention_mask']).logits[0][0].item()
print(score1,score2)

# Response 1 score: 2.3285412788391113, Response 2 score: 12.438033103942871

📚 ドキュメント

データセット

Property	Details
データセット1	nvidia/HelpSteer2
データセット2	Skywork/Skywork-Reward-Preference-80K-v0.1

論文とモデル

• 論文: https://arxiv.org/pdf/2410.00847
• モデル: URM-LLaMa-3.1-8B
  • ファインチューニング元: Skywork-Reward-Llama-3.1-8B

アーキテクチャ

URMは、図中のRMの1つです。

アライメント結果

不確実性の推定を用いて大規模言語モデルのアライメントを改善した結果です。不確実性の低い報酬はより信頼性が高く、より良いアライメント結果をもたらします。

概要

URM-LLaMa-3.1-8Bは、不確実性を考慮した報酬モデルです。このRMは、ベースモデルと不確実性を考慮した属性固有のバリューヘッドから構成されています。ベースモデルはSkywork-Reward-Llama-3.1-8Bを使用しています。

URMは2段階のトレーニングを行っています。

1. 属性回帰
2. ゲーティングレイヤーの学習

属性回帰

データセット: HelpSteer2

トレーニング中は、多属性スコアではなく、不確実性を考慮したバリューヘッドの出力を正規分布のパラメータとし、そこからスコアをサンプリングします。その後、出力とラベルを用いて回帰を行い、バリューヘッドをトレーニングします。勾配の逆伝播を可能にするために、再パラメータ化技術を使用しています。

ゲーティングレイヤーの学習

データセット: Skywork-Reward-Preference-80K-v0.1

ArmoRMにインスパイアされ、SteerLM-RMの固定重みではなく、ゲーティングレイヤーを学習して多属性スコアを結合しています。ゲーティングレイヤーの学習目標は、BT損失を通じて選択された応答を拒否された応答よりも優先することです。このプロセスでは、HelpSteer2の5つの属性（有用性、正確性、一貫性、複雑性、冗長性）のみを使用しています。この間、バリューヘッドとベースモデルは凍結されたままです。

🔧 技術詳細

本READMEには十分な技術詳細が記載されていないため、このセクションをスキップします。

📄 ライセンス

本READMEにはライセンス情報が記載されていないため、このセクションをスキップします。

参考文献

以下の文献を引用してください。

@article{lou2024uncertainty,
  title={Uncertainty-aware Reward Model: Teaching Reward Models to Know What is Unknown},
  author={Lou, Xingzhou and Yan, Dong and Shen, Wei and Yan, Yuzi and Xie, Jian and Zhang, Junge},
  journal={arXiv preprint arXiv:2410.00847},
  year={2024}
}