把评分换到优化时踩过的坑

在实际项目中,我对 reward model 有这几个具体需求:

  1. 能打分:给定 (prompt, response) 对,输出一个分数(0-10,或者 0-1)
  2. 能排序:对于同一个 prompt 的多个 response,能正确排序(好的 > 坏的)
  3. 训练成本可接受:不能像 GPT-4 那样花几百万,最好能用小模型搞定
  4. 推理够快:实际训练 PPO 的时候要频繁调用 reward model,推理太慢会卡死整个流程

简单说:不追求 state-of-the-art,但求能用、够用、成本低

最近在做 LLM 对齐相关的工作,绕不开 reward model(奖励模型)。以前只在论文里见过,实际跑起来才发现:看着简单,踩坑无数。下面是从"想搞懂 reward model 是什么"到"能跑通一个简单 PPO 流程"的完整记录,弯路都写在里头了。

背景:为什么需要奖励模型

事情是这样的。我们在做一个内容生成系统,需要模型输出高质量、安全的回复。但是 baseline 模型(比如 Llama 2 7B)有时候会:

  • 生成一些不相关的内容(问题问东他答西)
  • 输出格式不对(要求 JSON 结果给了纯文本)
  • 说话风格不一致(有时候正经有时候突然来个表情包)

一开始尝试用 prompt engineering 来解决,加了各种 instruction、example、constraint。效果有,但不稳定——有时候工作,有时候又不行,而且每次都要重新调 prompt,成本很高。

后来接触到 RLHF(Reinforcement Learning from Human Feedback),发现核心就是:训练一个能"打分"的模型,告诉主模型哪些回答好、哪些回答不好

这个"打分"的模型,就是 reward model。

实现:从零构建奖励模型

3.1 数据准备

第一个坑就是数据。一开始以为随便搞点问答对就行,结果完全不是这样。

reward model 需要的是 preference data(偏好数据),也就是同一个问题,有多个回答,而且有人工标注哪个回答更好。

数据格式大概长这样:

{
  "prompt": "如何提高编程能力?",
  "responses": [
    {"text": "多写代码,多实践。", "rank": 2},
    {"text": "推荐《代码大全》《重构》等经典书籍,每天坚持编码至少 2 小时,参与开源项目...", "rank": 1},
    {"text": "随便学学就行。", "rank": 3}
  ]
}

一开始自己搞了点数据,质量不行。后来找了几个开源数据集:

  • HH-RLHF:Anthropic 的数据集,质量不错,量也大
  • Stanford Human Preferences:比较多样,覆盖多种场景
  • OpenAI WebGPT Comparisons:偏重长文本生成

但有个问题:这些数据集都是英文的,我们要用中文。怎么办?

偷懒方案:用 GPT-4 把英文数据翻译成中文。虽然可能损失点语义细节,但实践下来效果还行。

数据预处理代码(简化版):

def load_preference_data(data_path):
    data = []
    with open(data_path, 'r') as f:
        for line in f:
            item = json.loads(line)
            # 提取 prompt 和多个 response
            prompt = item['prompt']
            responses = item['responses']
            # 构造 preference pairs(选对的,选错的)
            for i in range(len(responses)):
                for j in range(i+1, len(responses)):
                    if responses[i]['rank'] < responses[j]['rank']:
                        # i 比 j 好
                        data.append({
                            'prompt': prompt,
                            'chosen': responses[i]['text'],
                            'rejected': responses[j]['text']
                        })
    return data

3.2 模型选择

第二个坑是模型选型。一开始想直接用 GPT-4 打分,但:

  • 成本太高(每 1k tokens 要花不少钱)
  • 推理太慢(PPO 的时候要频繁调用)

后来想了想:reward model 本质就是个二分类器,不需要太大

方案:

  1. 用一个小模型(比如 BERT-base,或者 Llama 2 7B 的某个变种)
  2. 输入是 [prompt, response],输出是单个标量(分数)
  3. 训练目标是:chosen 的分数 > rejected 的分数

我最后选了 BertForSequenceClassification,原因:

  • 速度快(比 Llama 快多了)
  • 对短文本效果好(我们的场景不需要处理超长文本)
  • 训练成本低(显存占用小)

模型架构(示意):

from transformers import BertForSequenceClassification

model = BertForSequenceClassification.from_pretrained(
    'bert-base-uncased',
    num_labels=1,  # 输出单个分数
    problem_type='regression'  # 回归问题(不是分类)
)

3.3 训练过程

训练 reward model 的核心是 ranking loss,目标是让模型学会排序,而不是精确打分。

常用 loss 函数:

def ranking_loss(chosen_score, rejected_score, margin=1.0):
    """
    chosen_score: chosen response 的分数(batch_size, 1)
    rejected_score: rejected response 的分数(batch_size, 1)
    margin: 希望 chosen 比 rejected 高多少
    """
    # loss = max(0, margin - (chosen - rejected))
    loss = torch.relu(margin - (chosen_score - rejected_score))
    return loss.mean()

训练循环(简化):

optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)

for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for batch in dataloader:
        # 输入格式:[CLS] prompt [SEP] response [SEP]
        chosen_input = tokenize(batch['prompt'], batch['chosen'])
        rejected_input = tokenize(batch['prompt'], batch['rejected'])

        # 前向传播
        chosen_score = model(**chosen_input).logits
        rejected_score = model(**rejected_input).logits

        # 计算 loss
        loss = ranking_loss(chosen_score, rejected_score)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    print(f"Epoch {epoch}, Loss: {total_loss / len(dataloader)}")

3.4 评估

训练完怎么知道好不好用?

方法 1:看 accuracy(但这个不太准确)

def evaluate_ranking_accuracy(model, test_data):
    correct = 0
    total = 0
    for item in test_data:
        chosen_score = model.score(item['prompt'], item['chosen'])
        rejected_score = model.score(item['prompt'], item['rejected'])
        if chosen_score > rejected_score:
            correct += 1
        total += 1
    return correct / total

方法 2:看分数分布

import matplotlib.pyplot as plt

def plot_score_distribution(scores):
    plt.hist(scores, bins=50)
    plt.xlabel('Reward Score')
    plt.ylabel('Count')
    plt.title('Distribution of Reward Scores')
    plt.show()

方法 3:人工抽查(这个最靠谱,但成本高)

挑几个样本,看模型给的分数是否合理:

PromptResponse AResponse BA 分数B 分数合理?
如何学习 Python?“多写代码就行。”“推荐《Python 编程从入门到实践》,先学基础语法…”3.28.7
解释什么是机器学习?“机器学习是让机器从数据中学习规律的技术…”“不知道。”7.51.2

踩坑记录

坑 1:数据太少,模型过拟合

一开始只用了几千条数据,结果训练完在训练集上 accuracy 99%,但在测试集上只有 60%。

解决

  • 增加数据量(至少上万条)
  • 加正则化(dropout、weight decay)
  • 用早停(early stopping)

坑 2:loss 没怎么降

训练了 10 个 epoch,loss 还是差不多,基本没变化。

排查发现:

  • 学习率太大(一开始用的 1e-3,后来降到 2e-5)
  • 数据预处理有问题(有些 prompt 太长被截断了)
  • 模型初始化不好(换了预训练权重)

解决

  • 调低学习率,用 warmup
  • 检查 tokenizer,确保长文本不会被截断
  • 用更好的预训练权重

坑 3:分数分布奇怪

训练完发现所有分数都集中在 4-6 之间,没法区分好坏。

原因

  • 没用 margin(ranking loss 里的 margin 参数)
  • 模型太保守(预测都往中间值靠)

解决

  • 加 margin(比如 1.0 或 2.0),强迫模型拉开差距
  • 后处理:把分数标准化到 [0, 1] 区间
def normalize_score(score, min_score, max_score):
    return (score - min_score) / (max_score - min_score)

坑 4:推理太慢

PPO 训练的时候,每步都要调用 reward model 打分,结果推理成了瓶颈。

解决

  • 模型量化(从 fp32 降到 int8)
  • 用批推理(batch inference)
  • 换更小的模型(比如从 BERT-base 换到 DistilBERT)
# 量化
from transformers import BertForSequenceClassification, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0
)

model = BertForSequenceClassification.from_pretrained(
    'bert-base-uncased',
    quantization_config=quantization_config
)

结果:实际效果怎么样

训练完的 reward model 在我们的场景下表现如下。

用得上的地方:好的回答平均 7.8 分,差的平均 3.2 分,区分度够用;批量处理 1000 条约 200ms,PPO 训练不会卡死;整个训练成本不到 $50。

还差点意思的地方:超过 512 token 的内容经常被截断,打分不准;偶尔会翻车,给明显差的回答高分;新领域的 prompt 没见过,打分飘。

PPO 效果

用这个 reward model 训练 PPO 后,主模型的生成质量确实有提升:

指标训练前训练后
相关性(人工评分)6.2/107.5/10
格式正确率68%85%
风格一致性52%78%

PPO 训练前后三项指标的提升幅度不同,并列柱状图能看清 reward model 主要改善了哪些维度。

PPO 训练前后在相关性、格式正确率和风格一致性上的对比

不算惊艳,但格式与风格一致性改善明显,说明 reward model 已足够驱动一轮可用的对齐训练。

写在后面

搞 reward model 一圈下来,几条教训:

数据质量比模型大小重要——preference data 烂,BERT 也救不了。够用就行,别在调参上耗太久。训练过程中要定期人工抽查,模型会学到奇怪的东西。成本、速度、显存这些硬约束,比 SOTA 更实际。

下一步打算试 Llama 2 7B 做 reward model,多收几个领域的偏好数据,也看看能不能用 GPT-4 自动构造 preference pair。

reward model 这个领域还在变,新的训练方法和数据构造策略出得很快。如果你也在搞这个,欢迎交流踩坑经验。

版权声明: 本文首发于 指尖魔法屋-把评分换到优化时踩过的坑https://blog.thinkmoon.cn/post/396-ai-reward-model-rating-optimization-practice/) 转载或引用必须申明原指尖魔法屋来源及源地址!