URM LLaMa 3.1 8B

🚀 [URM-LLaMa-3.1-8B 不確定感知獎勵模型]

[URM-LLaMa-3.1-8B 是一個不確定感知獎勵模型，由基礎模型和不確定感知且特定屬性的值頭組成。該模型通過兩階段訓練，能有效提升大語言模型的對齊效果。]

🚀 快速開始

本項目主要圍繞 URM-LLaMa-3.1-8B 不確定感知獎勵模型展開，下面為你介紹其相關信息。

相關資源鏈接

• 論文：https://arxiv.org/pdf/2410.00847
• 模型：URM-LLaMa-3.1-8B，該模型基於 Skywork-Reward-Llama-3.1-8B 進行微調。

模型架構

URM 是圖中的獎勵模型（RM）之一。

對齊結果

使用不確定性估計來改善大語言模型對齊的結果。不確定性較低的獎勵更可靠，能帶來更好的對齊效果。

✨ 主要特性

URM-LLaMa-3.1-8B 是一個不確定感知獎勵模型，由基礎模型和不確定感知且特定屬性的值頭組成。該獎勵模型涉及兩階段訓練：

1. 屬性迴歸
2. 門控層學習

📚 詳細文檔

屬性迴歸

• 數據集：HelpSteer2

在訓練期間，不確定感知值頭的輸出不是多屬性分數，而是正態分佈的參數，分數從該分佈中採樣得到。然後，我們對輸出與標籤進行迴歸以訓練值頭。為了實現梯度反向傳播，使用了重參數化技術。

門控層學習

• 數據集：Skywork-Reward-Preference-80K-v0.1

受 ArmoRM 的啟發，我們學習一個門控層來組合多屬性分數，而不是使用 SteerLM-RM 中的固定權重。門控層的學習目標是通過 BT 損失使選擇的響應優先於拒絕的響應。我們僅使用來自 HelpSteer2 的五個屬性：有用性、正確性、連貫性、複雜性和冗長性。在此過程中，值頭和基礎模型保持凍結。

💻 使用示例

基礎用法

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

📄 許可證

文檔中未提及許可證相關信息。

🔧 技術細節

數據集

屬性	詳情
數據集1	nvidia/HelpSteer2，用於屬性迴歸階段
數據集2	Skywork/Skywork-Reward-Preference-80K-v0.1，用於門控層學習階段

訓練階段

• 屬性迴歸：訓練時值頭輸出正態分佈參數，通過重參數化技術實現梯度反向傳播進行迴歸訓練。
• 門控層學習：受 ArmoRM 啟發學習門控層組合多屬性分數，使用 BT 損失優化，僅考慮 HelpSteer2 的五個屬性，訓練時值頭和基礎模型凍結。

📖 引用

請引用以下論文：

@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}
}