🚀 [URM-LLaMa-3.1-8B 不确定感知奖励模型]
[URM-LLaMa-3.1-8B 是一个不确定感知奖励模型,由基础模型和不确定感知且特定属性的值头组成。该模型通过两阶段训练,能有效提升大语言模型的对齐效果。]
🚀 快速开始
本项目主要围绕 URM-LLaMa-3.1-8B 不确定感知奖励模型展开,下面为你介绍其相关信息。
相关资源链接
模型架构
URM 是图中的奖励模型(RM)之一。
对齐结果
使用不确定性估计来改善大语言模型对齐的结果。不确定性较低的奖励更可靠,能带来更好的对齐效果。
✨ 主要特性
URM-LLaMa-3.1-8B 是一个不确定感知奖励模型,由基础模型和不确定感知且特定属性的值头组成。该奖励模型涉及两阶段训练:
- 属性回归
- 门控层学习
📚 详细文档
属性回归
在训练期间,不确定感知值头的输出不是多属性分数,而是正态分布的参数,分数从该分布中采样得到。然后,我们对输出与标签进行回归以训练值头。为了实现梯度反向传播,使用了重参数化技术。
门控层学习
受 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}]
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)
📄 许可证
文档中未提及许可证相关信息。
🔧 技术细节
数据集
训练阶段
- 属性回归:训练时值头输出正态分布参数,通过重参数化技术实现梯度反向传播进行回归训练。
- 门控层学习:受 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}
}