---

数据集:

• nvidia/HelpSteer2
• 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是图中的奖励模型之一。

对齐效果

使用不确定性估计改进大语言模型对齐的效果展示。低不确定性的奖励信号更可靠，能带来更好的对齐效果。

简介

URM-LLaMa-3.1-8B是一个具备不确定性感知能力的奖励模型。 该模型由基础模型和具有不确定性感知能力的属性特定值头组成，其基础模型源自Skywork-Reward-Llama-3.1-8B。

URM采用两阶段训练：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 = "神经网络中sigmoid节点的数值输出范围是多少？"
response1 = "sigmoid节点的输出范围在-1到1之间。"
response2 = "sigmoid节点的输出范围在0到1之间。"

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

# 格式化对话并分词
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)

# 响应1得分: 2.3285412788391113, 响应2得分: 12.438033103942871

参考文献

引用格式：

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