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