613. AI Mathematical Olympiad - Progress Prize 1 | ai-mathematical-olympiad-prize
我们微调了两个 DeepSeek-Math-7B-RL 模型，一个用作策略模型（生成解决方案），一个用作奖励模型（为加权多数投票评分）。训练基础设施 adopted from (Sun et al., 2024)，bibtex 见文末。
我们的数据集是 AMC、AIME 和 Odyssey-Math 的组合。我们只选择那些答案为整数的问题，因为测试问题的答案将是整数。我们也删除了 AMC 中多选题的选项。
我们使用 few-shot 示例提示 GPT4 为我们数据集中的问题采样代码解决方案，并选择正确的解决方案作为策略模型的训练数据集。
数据集可以在我们的 GitHub AIMO-CMU-MATH 找到。
我们使用学习率 2e-5 对模型进行了 3 个 epoch 的微调。
我们使用 MATH、AIME、AMC 和 Odyssey-Math 中具有非负整数答案的问题作为训练奖励模型的问题集。
在 MATH 数据集上，如果我们用 DeepSeek-MATH-7B RL 采样解决方案，我们只会得到某些问题的正确解决方案，而且这些解决方案可能相似。另一方面，用 DeepSeek-MATH-7B Base 采样解决方案在某些问题上无法输出任何正确解决方案。
在 AMC、AIME 和 Odyssey 上，观察结果类似。当使用 DeepSeek-Math-7B-RL 采样解决方案时，大多数结果是错误的；当使用我们微调的策略模型时，准确率很高。
我们相信为了训练一个好的奖励模型，数据集应该满足：
因此，我们采用以下采样策略。
对于 MATH 数据集，我们插值 DeepSeek-Math-7B RL 和 DeepSeek-Math-7B Base 的模型参数，以获得具有不同能力水平的模型。分别将 DeepSeek-Math-7B RL 和 DeepSeek-Math-7B Base 的模型参数表示为 $M_{RL}$ 和 $M_{Base}$。插值模型参数 $M_{\alpha}$ 是 $M_{RL}$ 和 $M_{Base}$ 的线性组合。
$$M_{\alpha}=\alpha M_{RL} + (1-\alpha)M_{Base}$$
我们选择了 8 个不同的模型，$\alpha=0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0$。每个模型为每个问题生成两个解决方案。$\alpha$ 越大，模型采样的正确答案越多。
通过插值，我们可以得到不同的模型，产生多样的解决方案，并且正标签和负标签的比例比仅用一个模型采样更加平衡。
对于 AMC、AIME、Odyssey，我们使用为策略模型收集的训练数据集，并用 0、1、2 和 3 个 epoch 训练 DeepSeek-Math-RL 7B，以获得 4 个不同的模型。每个模型为每个问题生成 12 个解决方案。
通过这种方式，我们获得了不同的解决方案，并且每个问题更有可能同时拥有正确解决方案和错误解决方案。
一旦我们收集了所有带标签的解决方案，我们进行以下过滤：
我们发现上述方法并没有使模型在我们的验证数据集上表现更好。因此我们使用前面描述的方法。
对于策略模型，我们使用 vLLM 来加速采样，对于奖励模型使用 huggingface transformer。评分解决方案只是一个前向传递，不需要太多时间。每个模型使用一个 GPU (T4)。
我们为每个问题采样 42 个解决方案，策略模型被要求编写代码。Prompt 简单地是 problem + tool_instruction。
tool_instruction = '\nPlease integrate natural language reasoning with programs to solve the problem above, and put your final answer within \\boxed{}.'
如果代码执行结果不是整数，我们给予反馈并要求策略模型重试，代码如下：
def vllm_round_inference(problem, num_sequences, T, rounds):
prompt = problem + tool_instruction
texts = [prompt for _ in range(num_sequences)]
results = []
result_texts = []
for i in range(rounds):
sampling_params = SamplingParams(temperature=T, top_p=1.0, n=1, max_tokens=1024, stop="```output")
generated_texts = llm.generate(texts, sampling_params)
next_round_texts = []
completed = []
print(len(generated_texts))
for j, (text, generated_text) in enumerate(zip(texts, generated_texts)):
new_text = generated_text.outputs[0].text
completed = False
if '```python' in new_text:
try:
code_text = new_text.split('```python')[-1].split("```")[0]
code_result, CODE_STATUS = process_code(code_text, return_shell_output=True)
try:
result = int(float(code_result.strip()))
if result >= 0:
results.append(result % 1000)
result_texts.append((prompt + "```python" + new_text.split('```python')[-1]).split("```output")[0])
completed = True
except Exception as e:
code_result += "The code output is not an integer, final answer should be an integer."
except Exception as e:
code_result = "Code execution fail"
pass
if completed == False:
next_round_texts.append(text + new_text + "```output\n" + code_result + "\n```")
texts = next_round_texts
print("remaining text:", len(texts))
if texts == []:
break
return results, result_texts其中 num_sequences 设置为 42, T=0.9, rounds=2。
结果奖励模型 (ORM) 为每个解决方案分配一个分数。ORM 的输入是 problem + code 而不是 problem + solution。我们让 ORM 对策略模型生成的代码进行评分，因为最终候选答案仅仅是代码的执行结果。
def orm_score(reward_model, tokenizer, text):
input_id = torch.tensor([tokenizer.encode(text)]).to(reward_model.device)
with torch.no_grad():
logits = reward_model(input_id).logits
score = logits[0].mean(dim=-1).sigmoid().cpu().tolist()[-1]
return score为了计算一个答案的总权重，我们将其计算为几何平均数乘以 $N$，即如果 $N$ 个解决方案导致相同的答案 $a$，第 $i$ 个解决方案的得分为 $S_i$，答案 a 的权重为：
$$N(\prod_{i=1}^n S_i)^{\frac{1}{N}}.$$
我们选择几何平均数乘以 $N$ 而不是权重之和，因为它可以防止过度优化。例如，在我们的实验中，问题“存在一个唯一的递增几何序列，包含五个两位正整数。它们的和是多少？”的答案是 211。ORM 给一个带有 211 的解决方案评分 0.75，但给答案 50 的解决方案评分 0.85, 0.03, 0.15。0.85 可能是奖励模型的假阳性，几何平均数防止我们选择这样的答案。
我们还注意到策略模型倾向于生成答案为 0 的错误解决方案，所以我们对答案 0 添加了一些惩罚。以下是加权多数投票的代码。
def weighted_geometric_mean_times_num(answers, scores):
import math
max_weight = -1
number_weight_dict = {}
max_weight_answer = None
for answer, score in zip(answers, scores):
if answer == 0:
score /= 10
if answer not in number_weight_dict:
number_weight_dict[answer] = [score]
else:
number_weight_dict[answer].append(score)
for answer in number_weight_dict.keys():
score_list = number_weight_dict[answer]
weight = len(score_list) * math.prod(score_list) ** (1 / len(score_list))
if weight > max_weight:
max_weight = weight
max_weight_answer = answer
return max_weight_answer我们将 Notebook 的高性能归因于以下两点：
在 10 个问题的 AIMO 训练集上，使用微调策略模型，准确率从 1/10 提高到 2/10，奖励模型进一步将其从 2/10 提高到 4/10。
微调 代码库 采用自以下论文的代码：
@article{sun2024easy,
title={Easy-to-Hard Generalization: Scalable Alignment Beyond Human Supervision},
author={Sun, Zhiqing and Yu, Longhui and Shen, Yikang and Liu, Weiyang and Yang, Yiming and Welleck, Sean and Gan, Chuang},
journal={arXiv preprint arXiv:2403.09472},
year={2024}
}