STaR（Self-Taught Reasoner）方法：讓語言模型自學推理能力

在大型語言模型（LLM）的推理能力優化中，STaR（Self-Taught Reasoner） 是一種引人注目的技術，屬于“修改提議分布（Modifying Proposal Distribution）”類別。與傳統的基于結果驗證（verifier）方法不同，STaR通過訓練模型生成更好的推理步驟（input-focused），直接調整采樣分布，使其傾向于選擇“推理相關”的token。本文將詳細介紹STaR的原理、工作流程，并提供一個可運行的Python代碼實現，幫助你理解和實踐這一方法。

參考：https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-reasoning-llms

---

1. STaR的原理

背景

傳統的LLM生成方法通常依賴貪婪解碼（選擇最高概率token）或隨機采樣，但這些方法可能無法生成邏輯嚴謹的推理步驟。STaR通過讓模型自生成推理數據并進行監督微調（Supervised Fine-Tuning），優化其推理能力，調整token的提議分布，使其更傾向于推理過程。

核心思想

• 自生成推理數據：模型首先生成推理步驟和答案。
• 驗證與修正：
  • 如果答案正確，直接將推理步驟和答案加入訓練數據集。
  • 如果答案錯誤，提供正確答案作為“提示”，讓模型重新推理并生成正確過程。
• 監督微調：用生成的數據集訓練模型，強化其推理行為。

目標

• 輸入聚焦：通過修改提議分布，使模型更擅長生成推理相關token，而非簡單輸出結果。
• 自增強：利用模型自身生成的數據，無需大量人工標注。

---

2. STaR的工作流程

STaR的核心是一個循環過程，包含以下步驟：

1. 生成推理步驟和答案：

   • 模型根據問題生成推理路徑和最終答案。

2. 驗證答案：

   • 正確（2a）：推理和答案正確，進入步驟3b。
   • 錯誤（2b）：答案錯誤，進入步驟4b。

3. 正確答案處理（3b）：

   • 將問題、推理步驟、答案組成三元組，加入訓練數據集。

4. 錯誤答案修正（4b）：

   • 提供正確答案作為提示，要求模型重新生成推理步驟。
   • 將修正后的推理加入訓練數據集。

5. 監督微調（5）：

   • 使用生成的三元組數據集，對模型進行監督微調，優化推理能力。

關鍵特點

• 合成數據：STaR通過自生成數據創建訓練樣本，類似于數據蒸餾。
• 迭代改進：多次循環生成和微調，逐步提升模型性能。

---

3. 代碼實現

以下是一個簡化的STaR實現，基于PyTorch。我們模擬一個數學推理任務（如“2 + 3 = ?”），展示其工作流程。

import torch
import torch.nn as nn
import torch.nn.functional as F
from copy import deepcopy# 超參數
vocab_size = 10  # 詞匯表大小（0-9數字）
embed_size = 16
num_heads = 2
hidden_size = 32
num_layers = 2
max_steps = 3   # 最大推理步驟# 生成模型
class SimpleReasoner(nn.Module):def __init__(self):super(SimpleReasoner, self).__init__()self.embedding = nn.Embedding(vocab_size, embed_size)self.transformer = nn.TransformerDecoderLayer(embed_size, num_heads, hidden_size)self.output_layer = nn.Linear(embed_size, vocab_size)def forward(self, x):x = self.embedding(x)x = self.transformer(x, x)return self.output_layer(x)def generate(self, prompt, max_len=3, temperature=1.0):seq = prompt.copy()inputs = torch.tensor([seq], dtype=torch.long).to(device)for _ in range(max_len - len(seq)):logits = self.forward(inputs)[:, -1, :]probs = F.softmax(logits / temperature, dim=-1)next_token = torch.multinomial(probs, 1).item()seq.append(next_token)inputs = torch.tensor([seq], dtype=torch.long).to(device)return seqdef train_step(self, data, optimizer):self.train()optimizer.zero_grad()inputs = torch.tensor([d[0] + d[1][:-1] for d in data], dtype=torch.long).to(device)targets = torch.tensor([d[1] for d in data], dtype=torch.long).to(device)logits = self.forward(inputs)loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))loss.backward()optimizer.step()return loss.item()# STaR實現
class STaR:def __init__(self, model):self.model = modelself.device = next(model.parameters()).devicedef generate_reasoning(self, prompt, correct_answer=None):if correct_answer is None:return self.model.generate(prompt, max_steps)# 提供正確答案作為提示hint_prompt = prompt + [correct_answer]return self.model.generate(hint_prompt, max_steps)def verify_answer(self, sequence, correct_answer):return sequence[-1] == correct_answerdef star_iteration(self, prompts, correct_answers, iterations=3):training_data = []for _ in range(iterations):new_model = deepcopy(self.model)  # 保存當前模型狀態optimizer = torch.optim.Adam(new_model.parameters(), lr=0.001)for prompt, correct_answer in zip(prompts, correct_answers):# 步驟1：生成推理步驟和答案sequence = self.generate_reasoning(prompt)# 步驟2：驗證答案if self.verify_answer(sequence, correct_answer):# 步驟3b：正確答案加入訓練數據training_data.append((prompt, sequence))else:# 步驟4b：錯誤答案，提供提示重新生成corrected_sequence = self.generate_reasoning(prompt, correct_answer)training_data.append((prompt, corrected_sequence))# 步驟5：監督微調if training_data:loss = new_model.train_step(training_data, optimizer)print(f"Iteration {_+1}, Loss: {loss}")self.model = new_model  # 更新模型return training_data# 初始化并運行
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleReasoner().to(device)
star = STaR(model)# 示例數據
prompts = [[2, 3]]  # "2 + 3"
correct_answers = [5]# 執行STaR
training_data = star.star_iteration(prompts, correct_answers, iterations=3)
print("Generated training data:", training_data)# 測試優化后的模型
test_prompt = [2, 3]
result = model.generate(test_prompt)
print(f"Test prompt: {test_prompt}, Generated result: {result}")

---

4. 代碼解析

生成模型（SimpleReasoner）

• generate：根據提示生成推理序列，模擬推理步驟和答案。
• train_step：使用監督微調優化模型，輸入為問題+推理步驟，目標為完整序列。

STaR實現

• generate_reasoning：
  • 無提示時：自由生成推理。
  • 有提示時：基于正確答案生成推理。
• verify_answer：檢查生成序列的最后一個token是否正確。
• star_iteration：
  • 步驟1：生成推理和答案。
  • 步驟2a/2b：驗證答案，正確則直接記錄，錯誤則用提示修正。
  • 步驟3b/4b：收集三元組（問題、推理、答案）。
  • 步驟5：用生成的數據微調模型。

運行邏輯

• 每次迭代生成數據，優化模型，逐步提高推理能力。
• 使用 deepcopy 保留模型狀態，確保迭代獨立。

---

5. 運行結果示例

運行代碼可能得到：

Iteration 1, Loss: 2.305
Iteration 2, Loss: 2.287
Iteration 3, Loss: 2.251
Generated training data: [([2, 3], [2, 3, 5]), ([2, 3], [2, 3, 5]), ([2, 3], [2, 3, 5])]
Test prompt: [2, 3], Generated result: [2, 3, 5]

• 未訓練模型初始生成隨機，STaR通過微調逐漸傾向于正確答案 [2, 3, 5]。
• 實際中需更多數據和迭代。

---

6. STaR的意義與改進

意義

• 自增強：無需大量人工數據，模型自生成訓練樣本。
• 推理優化：調整提議分布，強化推理token的選擇。
• 數據蒸餾：生成合成數據，可用于其他模型訓練。

改進方向

• 多樣化提示：增加問題類型（如數學、自然語言問答）。
• 獎勵函數：引入PRM評估推理步驟質量，而非僅驗證答案。
• 迭代控制：動態調整迭代次數或數據篩選標準。
• 預訓練模型：基于已有LLM（如GPT）實現，提升初始性能。

---

7. 總結

STaR通過自生成推理數據和監督微調，優化LLM的推理能力。其流程從生成到驗證再到修正，利用合成數據調整token分布，是“修改提議分布”的典型方法。代碼實現展示了從 [2, 3] 到 [2, 3, 5] 的優化過程，體現了其核心思想。運行這段代碼，你可以體驗STaR的自學過程。希望這篇博客對你理解和實踐STaR有所幫助！如需進一步優化，歡迎討論。

基于大型語言模型改進 STaR 方法：以 LLaMA 3 或 Qwen 2.5 為例

在之前的STaR（Self-Taught Reasoner）實現中，我們使用了一個簡化的模型來展示其工作原理。然而，為了在實際任務中獲得更好的推理能力，可以基于Hugging Face（HF）上的預訓練大型語言模型（LLM）如 LLaMA 3 或 Qwen 2.5 進行改進。本文將以中文博客的形式，結合改進方向（多樣化提示、獎勵函數、迭代控制、預訓練模型），詳細說明如何基于這些HF模型優化STaR，并提供改進后的代碼實現。

---

1. 改進背景與目標

原始實現局限

• 模型能力：SimpleReasoner 未經過預訓練，生成隨機且缺乏推理能力。
• 提示單一：僅支持簡單數學任務。
• 獎勵簡單：僅驗證答案，未評估推理質量。
• 靜態迭代：固定次數，缺乏靈活性。

改進目標

• 預訓練模型：利用LLaMA 3或Qwen 2.5的強大語言理解能力。
• 多樣化提示：支持數學和自然語言問答。
• 獎勵函數：引入過程獎勵模型（PRM）評估推理步驟。
• 迭代控制：動態調整迭代次數和數據篩選。

---

2. 改進方案

1. 基于預訓練模型：LLaMA 3 或 Qwen 2.5

• 選擇原因：
  • LLaMA 3：高效、適合微調，廣泛用于研究。
  • Qwen 2.5：開源，支持多語言，推理能力強。
• 實現：使用Hugging Face的 transformers 庫加載預訓練模型，替換 SimpleReasoner。

2. 多樣化提示

• 數學任務：如“2 + 3 = ?”。
• 自然語言問答：如“中國的首都是哪里？”。
• 方法：擴展輸入格式，支持文本和符號混合。

3. 獎勵函數：引入PRM

• 目的：評估推理步驟的邏輯性和完整性，而非僅答案。
• 實現：使用一個小型預訓練模型（如BERT）作為PRM，評分推理質量。

4. 迭代控制

• 動態調整：根據數據質量或損失收斂動態停止迭代。
• 數據篩選：僅保留高質量推理樣本。

---

3. 改進后的代碼實現

以下基于 Qwen 2.5（也可替換為LLaMA 3）的STaR實現，展示改進后的完整流程。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSequenceClassification
from copy import deepcopy
import random# 超參數
max_steps = 50  # 最大生成長度
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 初始化生成模型（Qwen 2.5）
model_name = "Qwen/Qwen2.5-7B-Instruct"  # 可替換為 "meta-llama/Llama-3-8B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
generator = AutoModelForCausalLM.from_pretrained(model_name).to(device)# 初始化PRM（使用BERT評估推理質量）
prm_name = "bert-base-uncased"
prm_tokenizer = AutoTokenizer.from_pretrained(prm_name)
prm_model = AutoModelForSequenceClassification.from_pretrained(prm_name, num_labels=1).to(device)# STaR實現
class STaR:def __init__(self, generator, tokenizer, prm_model, prm_tokenizer):self.generator = generatorself.tokenizer = tokenizerself.prm_model = prm_modelself.prm_tokenizer = prm_tokenizerdef generate_reasoning(self, prompt, correct_answer=None, temperature=0.7):"""生成推理步驟和答案"""if correct_answer is None:input_text = f"問題: {prompt}\n推理步驟和答案:"else:input_text = f"問題: {prompt}\n正確答案: {correct_answer}\n請提供推理步驟:"inputs = self.tokenizer(input_text, return_tensors="pt").to(device)outputs = self.generator.generate(**inputs, max_length=max_steps, temperature=temperature,do_sample=True, pad_token_id=self.tokenizer.eos_token_id)return self.tokenizer.decode(outputs[0], skip_special_tokens=True)def verify_answer(self, response, correct_answer):"""驗證答案是否正確"""answer_part = response.split("答案:")[-1].strip()return str(correct_answer) in answer_partdef evaluate_reasoning(self, response):"""使用PRM評估推理質量"""inputs = self.prm_tokenizer(response, return_tensors="pt", truncation=True, max_length=512).to(device)with torch.no_grad():score = self.prm_model(**inputs).logits.item()return score  # 返回正值表示推理質量def star_iteration(self, prompts, correct_answers, max_iterations=5, min_loss=0.1):training_data = []model = deepcopy(self.generator)optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)for iteration in range(max_iterations):new_data = []total_loss = 0.0for prompt, correct_answer in zip(prompts, correct_answers):# 步驟1：生成推理和答案response = self.generate_reasoning(prompt)# 步驟2：驗證答案if self.verify_answer(response, correct_answer):# 步驟3b：正確答案，檢查推理質量score = self.evaluate_reasoning(response)if score > 0.5:  # 篩選高質量推理new_data.append((prompt, response))else:# 步驟4b：錯誤答案，提供提示重新生成corrected_response = self.generate_reasoning(prompt, correct_answer)score = self.evaluate_reasoning(corrected_response)if score > 0.5:new_data.append((prompt, corrected_response))# 步驟5：監督微調if new_data:model.train()optimizer.zero_grad()inputs = self.tokenizer([d[0] + "\n" + d[1] for d in new_data], return_tensors="pt", padding=True, truncation=True, max_length=max_steps).to(device)labels = inputs["input_ids"].clone()outputs = model(**inputs, labels=labels)loss = outputs.lossloss.backward()optimizer.step()total_loss += loss.item()training_data.extend(new_data)print(f"Iteration {iteration+1}, Loss: {total_loss / len(new_data) if new_data else 0}")if total_loss / len(new_data) < min_loss and new_data:breakself.generator = modelreturn training_data# 示例數據
prompts = ["2 + 3等于多少？","中國的首都是哪里？"
]
correct_answers = ["5", "北京"]# 初始化STaR
star = STaR(generator, tokenizer, prm_model, prm_tokenizer)# 執行STaR
training_data = star.star_iteration(prompts, correct_answers)
print("Generated training data:", training_data)# 測試優化后的模型
for prompt in prompts:result = star.generate_reasoning(prompt)print(f"Prompt: {prompt}, Generated result: {result}")

---

4. 代碼解析

1. 預訓練模型：Qwen 2.5

• 加載：使用 AutoModelForCausalLM 加載Qwen 2.5，替換簡化的 SimpleReasoner。
• 生成：generate_reasoning 使用 model.generate 支持多樣化提示，生成文本而非token序列。
• 優勢：Qwen 2.5 已具備語言理解能力，初始生成更接近推理。

2. 多樣化提示

• 輸入格式：
  • 數學："2 + 3等于多少？\n推理步驟和答案:"。
  • 問答："中國的首都是哪里？\n推理步驟和答案:"。
• 輸出：支持自然語言，生成完整句子，如“推理：2加3等于5，答案：5”。

3. 獎勵函數：PRM

• 實現：使用BERT作為PRM，評分推理文本的邏輯性。
• 篩選：score > 0.5 保留高質量推理，避免噪聲數據。
• 改進：可訓練BERT區分正確推理（如“2+3=5”）和錯誤推理（如“2*3=5”）。

4. 迭代控制

• 動態停止：若損失低于 min_loss（如0.1），提前終止。
• 數據篩選：結合PRM分數，確保訓練數據質量。

---

5. 運行結果示例

運行代碼可能得到：

Iteration 1, Loss: 0.85
Iteration 2, Loss: 0.62
Iteration 3, Loss: 0.09
Generated training data: [('2 + 3等于多少？', '問題: 2 + 3等于多少？\n推理步驟和答案: 首先，2加上3，等于5。\n答案: 5'),('中國的首都是哪里？', '問題: 中國的首都是哪里？\n推理步驟和答案: 中國是一個國家，其首都是北京。\n答案: 北京')
]
Prompt: 2 + 3等于多少？, Generated result: 問題: 2 + 3等于多少？\n推理步驟和答案: 首先，2加上3，等于5。\n答案: 5
Prompt: 中國的首都是哪里？, Generated result: 問題: 中國的首都是哪里？\n推理步驟和答案: 中國是一個國家，其首都是北京。\n答案: 北京

• 結果：Qwen 2.5初始生成已較合理，微調后更傾向推理。

---

6. 基于LLM的改進優勢

預訓練能力

• Qwen 2.5 或 LLaMA 3 自帶語言理解和生成能力，初始推理質量高于隨機模型。
• STaR在此基礎上進一步強化推理分布。

多樣化支持

• 處理文本輸入，支持數學和問答任務，擴展性強。

PRM增強

• BERT作為PRM評估推理邏輯，確保生成數據不僅是正確答案，還包含合理步驟。

動態優化

• 損失收斂后停止，節省計算資源。

---

7. 進一步優化建議

• 更大模型：使用LLaMA 3-70B或Qwen 2.5-72B，提升推理深度。
• 混合獎勵：結合PRM和答案正確性（ORM），綜合評分。
• 數據蒸餾：將STaR生成的數據用于其他模型（如小規模LLM）的訓練。

---

8. 總結

基于Qwen 2.5的STaR改進，利用預訓練LLM的強大能力，支持多樣化提示，通過PRM優化推理質量，并動態控制迭代。代碼展示了從數學到問答的推理生成，體現了“修改提議分布”的核心思想。運行此代碼，你可以體驗基于HF模型的STaR優化過程。希望這篇博客對你有所幫助！如需調整或擴展，歡迎討論。

解析 STaR 中 star_iteration 的逐迭代訓練設計

提出疑問：為什么訓練是每個iteration都要進行，而不是將所有數據處理好后再進行一次訓練？下面詳細解析這種逐迭代訓練的設計動機，分析其優劣勢，并探討替代方案。

---

1. 逐迭代訓練的背景

STaR的核心思想

STaR（Self-Taught Reasoner）是一種自監督方法，通過讓模型生成推理數據并進行監督微調（Supervised Fine-Tuning），優化其推理能力。其流程本質上是一個迭代改進的過程：

1. 模型基于當前參數生成推理和答案。
2. 驗證答案，收集正確或修正后的數據。
3. 用生成的數據微調模型。
4. 重復上述步驟。

代碼中的訓練位置

• 每次迭代內訓練：在每個 for iteration in range(max_iterations) 循環中，生成 new_data 后立即調用 model.train_step 進行微調。
• 累計數據：training_data.extend(new_data) 將每次迭代的數據加入總數據集，但訓練發生在每次迭代結束時。

---

2. 為什么每個Iteration都要訓練？

1. 動態優化模型分布

• 提議分布的修改：
  • STaR的目標是調整模型的token提議分布，使其傾向于生成推理相關的內容。
  • 每次迭代后，模型參數通過微調更新，下一次生成會基于更優的分布。
• 逐次改進：
  • 如果不訓練，模型在所有迭代中都使用初始參數，生成的推理質量可能持續較差。
  • 每次訓練后，模型更可能生成正確的推理步驟，逐步提升數據質量。

2. 自增強反饋循環

• 自生成數據：
  • STaR依賴模型自身生成訓練數據，每次迭代的 new_data 是當前模型能力的反映。
  • 訓練后，模型能力提升，下次生成的 new_data 更接近期望的推理模式。
• 反饋效應：
  • 類似強化學習，每次迭代強化模型的推理行為，形成正反饋。

3. 數據質量的逐步提高

• 初始數據可能較差：
  • 未訓練模型生成的推理可能隨機或錯誤（如 [2, 3, 1]）。
  • 第一次訓練后，模型學會部分正確模式（如 [2, 3, 5]），后續數據更優質。
• 避免積累噪聲：
  • 若等到最后訓練，可能積累大量低質量數據，影響微調效果。

4. 計算資源與時間優化

• 小批量訓練：
  • 每次迭代只處理當前生成的 new_data（如2個樣本），訓練負擔輕。
  • 若積累所有數據再訓練，可能需要更大批量或更多epoch，增加內存和時間成本。
• 提前終止：
  • if total_loss / len(new_data) < min_loss: 允許在損失收斂時停止，無需完成所有迭代。

代碼中的體現

• 訓練時機：

  if new_data:model.train()optimizer.zero_grad()# ... 微調代碼 ...optimizer.step()
  

  • 每次迭代立即訓練，確保模型實時更新。

---

3. 模擬過程

任務

• prompts = ["2 + 3等于多少？"]，correct_answers = ["5"]。
• ( $\text{max\_iterations}=3$ )。

第一次迭代

• 生成：response = "問題: 2 + 3等于多少？\n推理和答案: 2 * 3 = 6\n答案: 6"。
• 驗證：錯誤。
• 修正：corrected_response = "問題: 2 + 3等于多少？\n正確答案: 5\n推理: 2 + 3 = 5"。
• 數據：new_data = [("2 + 3等于多少？", corrected_response)]。
• 訓練：微調模型，更新參數。

第二次迭代

• 生成：response = "問題: 2 + 3等于多少？\n推理和答案: 2 + 3 = 5\n答案: 5"（因訓練改進）。
• 驗證：正確，score > 0.5。
• 數據：new_data = [("2 + 3等于多少？", response)]。
• 訓練：進一步強化正確推理。

第三次迭代

• 生成：更穩定的正確推理。
• 數據：累計高質量樣本。
• 訓練：繼續優化。

對比假設

• 若最后訓練：
  • 第一次：[2, 3, 6]。
  • 第二次：[2, 3, 1]（仍隨機）。
  • 第三次：[2, 3, 4]。
  • 最后訓練可能因數據混雜，效果不佳。

---

4. 逐迭代訓練的優勢與劣勢

優勢

1. 實時反饋：每次迭代優化模型，提升后續生成質量。
2. 數據質量遞增：避免積累低質量數據。
3. 靈活終止：損失收斂時停止，節省資源。

劣勢

1. 計算開銷：頻繁訓練增加總計算時間。
2. 模型不穩定：小批量訓練可能導致參數波動。
3. 實現復雜性：需管理每次迭代的模型副本（如 deepcopy）。

---

5. 為何不等到所有數據處理好再訓練？

替代方案的問題

假設修改為收集所有數據后一次性訓練：

def star_iteration(self, prompts, correct_answers, max_iterations=5):training_data = []for _ in range(max_iterations):for prompt, correct_answer in zip(prompts, correct_answers):response = self.generate_reasoning(prompt)if self.verify_answer(response, correct_answer):if self.evaluate_reasoning(response) > 0.5:training_data.append((prompt, response))else:corrected_response = self.generate_reasoning(prompt, correct_answer)if self.evaluate_reasoning(corrected_response) > 0.5:training_data.append((prompt, corrected_response))# 一次性訓練if training_data:model = deepcopy(self.generator)optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)loss = model.train_step(training_data, optimizer)  # 假設支持多epochself.generator = modelreturn training_data

問題分析

1. 數據質量不一致：

   • 所有迭代使用初始模型，生成的 training_data 可能包含大量錯誤或低質量推理。
   • 無法利用中間訓練的改進。

2. 缺乏反饋：

   • 模型未在迭代中更新，每次生成無進步，可能浪費計算資源。

3. 訓練負擔：

   • 一次性處理大量數據需更多epoch或更高計算資源，可能超出現有硬件能力。

4. STaR目標偏離：

   • STaR強調自增強循環，逐迭代訓練是其核心機制，最后訓練削弱了這一特性。

---

6. 改進建議

折中方案

• 批次訓練：每隔幾輪迭代訓練一次，平衡反饋與效率：

  if new_data and iteration % 2 == 0:  # 每2輪訓練一次model.train_step(new_data, optimizer)
  

動態調整

• 自適應迭代：根據數據質量（如PRM分數）調整訓練頻率。
• 增量數據：僅訓練新增數據，避免重復計算。

---

7. 總結

STaR中逐迭代訓練的設計是為了：

• 動態優化：實時更新模型，提升每次生成的質量。
• 自增強：形成反饋循環，逐步強化推理能力。
• 效率：小批量訓練結合提前終止，適應資源限制。

相比之下，所有數據處理后再訓練可能導致數據質量低、缺乏反饋，違背STaR的自適應優化目標。代碼中的逐迭代訓練是其核心優勢的體現。

后記

2025年3月2日16點43分于上海，在grok3大模型輔助下完成。