菜单
Improve credit assignment in retrieval-augmented reasoning by truncating trajectories at single decision points. Generate k samples sharing a common prefix, differing only at the next step to isolate variation and reduce gradient variance by T-fold on T-step tasks.
默认使用会先检查来源的 Prompt；你也可以切换为直接命令，或下载本地副本。
决定是否安装前，请先阅读 SKILL.md，以及 SkillsMP 当前展示的配套文件。
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
直接命令不会经过审查 Prompt；运行前请先检查来源。
npx skills add https://github.com/ADu2021/skillXiv --skill truncated-step-level-retrieval-reasoning命令会保持在同一行。复制前请横向滚动并检查完整内容。
想先保存到本地？可下载 SkillsMP 当前能够提供的文件。
正在显示 SKILL.md
基于 SOC 职业分类
| name | truncated-step-level-retrieval-reasoning |
| title | Truncated Step-Level Sampling with Process Rewards for Retrieval-Augmented Reasoning |
| version | 0.0.2 |
| engine | skillxiv-v0.0.2-claude-opus-4.6 |
| license | MIT |
| url | https://arxiv.org/abs/2602.23440 |
| keywords | ["Retrieval-Augmented Generation","Variance Reduction","Step-Level Rewards","Multi-Hop Reasoning","Credit Assignment"] |
| description | Improve credit assignment in retrieval-augmented reasoning by truncating trajectories at single decision points. Generate k samples sharing a common prefix, differing only at the next step to isolate variation and reduce gradient variance by T-fold on T-step tasks. |
Standard reinforcement learning for retrieval-augmented generation (RAG) tasks suffers from poor credit assignment: high-variance gradient estimates make it hard to distinguish which search query or reasoning step improved final outcomes. Truncated step-level sampling fixes this by generating trajectories that share a prefix and diverge only at a single decision point. This isolation dramatically reduces variance while enabling per-step reward supervision through dense LLM-as-Judge scoring.
The core innovation treats each reasoning step as an isolated decision problem. By fixing all prior steps, the method achieves T-fold variance reduction on T-step tasks while maintaining gradient accuracy through reward decomposition.
Truncated sampling operates through two coordinated mechanisms:
Prefix-Sharing Trajectories: Generate k candidate continuations from identical prefixes, isolating variation to a single step. This eliminates confounding from earlier decisions.
Dense Step-Level Rewards: Score each step independently via LLM judge on three dimensions:
These two components combine to provide stable, interpretable learning signals.
Step 1: Structure multi-hop reasoning task with explicit steps
Define the reasoning process to enable step-level isolation.
def decompose_rag_task(question):
"""
Decompose RAG into explicit steps: search → retrieve → reason → search → answer.
Each step is atomic and independently scorable.
"""
steps = [
{: , : , : question},
{: , : , : },
{: , : , : },
{: , : , : },
{: , : , : },
{: , : , : }
]
steps
Step 2: Generate prefix and sample k step-t candidates
Execute deterministically up to step t-1, then generate k diverse continuations for step t.
def generate_truncated_batch(model, question, num_candidates=4, target_step=2):
"""
Generate batch of trajectories sharing prefix, differing only at target_step.
Example with target_step=2:
- Execute steps 0, 1 once (deterministic or single sample)
- Generate k candidates for step 2
- Execute steps 3+ identically for all k candidates
"""
steps = decompose_rag_task(question)
# Phase 1: Execute steps before target_step deterministically
prefix_state = question
for step_idx in range(target_step):
step = steps[step_idx]
if step_idx == 0:
prefix_state = generate_initial_query(model, question)
elif step_idx == 1:
documents = retrieve_documents(prefix_state)
prefix_state = documents
# ... continue deterministically
# Phase 2: Generate k candidates for step target_step
candidates = []
for _ in range(num_candidates):
candidate_action = model.generate(
f"Given context: {prefix_state}, generate next reasoning step:",
max_tokens=150,
temperature=0.7,
top_p=0.9
)
candidates.append(candidate_action)
# Phase 3: Continuation after target_step (reuse for all candidates)
trajectories = []
for candidate in candidates:
# Forward pass through remaining steps
trajectory = {
'prefix': prefix_state,
'step_t_action': candidate,
'target_step': target_step,
'full_trajectory': None
}
# Execute remaining steps using this candidate
current_state = candidate
for step_idx in range(target_step + 1, len(steps)):
current_state = execute_step(model, steps[step_idx], current_state)
trajectory['full_trajectory'] = current_state
trajectories.append(trajectory)
return trajectories
Step 3: Compute per-step LLM-as-Judge rewards
Score each step independently on reasoning quality, query quality, and answer correctness.
def compute_three_part_reward(trajectory, question, reference_answer,
judge_model):
"""
Decompose reward into three scoreable components:
1. Reasoning quality: Is the thinking sound?
2. Query quality: Is the search formulation likely to retrieve relevant docs?
3. Answer correctness: Does final answer match reference?
"""
target_step = trajectory['target_step']
step_action = trajectory['step_t_action']
full_trajectory = trajectory['full_trajectory']
# Component 1: Reasoning quality (LLM judge)
reasoning_prompt = f"""
Evaluate the reasoning quality of this step:
Question: {question}
Prefix context: {trajectory['prefix'][:100]}
Reasoning step: {step_action}
Rate on scale 0-1: Is the logic sound? Does it make progress toward answer?
"""
reasoning_score = judge_model.score(
reasoning_prompt,
return_numeric=True # 0-1 score
)
# Component 2: Query quality (relevance of search formulation)
if 'query' in step_action.lower() or 'search' in step_action.lower():
query_prompt = f"""
Evaluate query formulation quality:
Original question: {question}
Formulated query: {step_action}
Rate 0-1: Will this retrieve relevant documents?
"""
query_score = judge_model.score(query_prompt, return_numeric=True)
else:
query_score = 0.5 # Neutral if not a query step
# Component 3: Answer correctness (binary: correct/incorrect)
final_answer = full_trajectory.split('\n')[-1] # Last line of trajectory
is_correct = final_answer.lower() in reference_answer.lower()
answer_score = 1.0 if is_correct else 0.0
# Combine components (equal weighting)
total_reward = (reasoning_score + query_score + answer_score) / 3.0
return {
'reasoning_score': reasoning_score,
'query_score': query_score,
'answer_score': answer_score,
'total_reward': total_reward
}
# Evaluate batch
batch = generate_truncated_batch(model, question, num_candidates=4, target_step=2)
rewards = [compute_three_part_reward(traj, question, reference_answer, judge)
for traj in batch]
Step 4: Compute policy gradients with variance reduction
Calculate gradients using isolated step rewards and prove T-fold variance reduction.
def compute_pg_loss_truncated(model, batch, rewards, target_step):
"""
Policy gradient loss with variance reduction from truncation.
Key insight: By fixing prefix and only varying step t,
var(gradient) is reduced by ~T-fold relative to full-trajectory sampling.
"""
total_loss = 0.0
for trajectory, reward in zip(batch, rewards):
# Get log probability of the step-t action under current policy
step_action = trajectory['step_t_action']
prefix_context = trajectory['prefix']
logits = model.forward(
f"Context: {prefix_context}\nAction:",
output_logits=True
)
# Log probability of sampled action
action_logprob = compute_logprob(logits, step_action, model.tokenizer)
# Policy gradient: maximize log-prob of high-reward actions
# Loss = -logprob * reward (negative for gradient descent)
step_loss = -action_logprob * reward['total_reward']
total_loss += step_loss
return total_loss / len(batch)
def analyze_variance_reduction(trajectories_full, trajectories_truncated):
"""
Empirically verify T-fold variance reduction.
trajectories_full: full-trajectory samples from baseline
trajectories_truncated: truncated samples from this method
"""
# Compute gradient variance
grads_full = [compute_gradient(traj) for traj in trajectories_full]
grads_truncated = [compute_gradient(traj) for traj in trajectories_truncated]
var_full = np.var([g.norm().item() for g in grads_full])
var_truncated = np.var([g.norm().item() for g in grads_truncated])
# Expected T-fold reduction
T = 5 # Assume 5-step tasks
expected_ratio = T
observed_ratio = var_full / var_truncated
print(f"Variance reduction: {observed_ratio:.1f}x (expected ~{expected_ratio}x)")
return observed_ratio
Step 5: RL training loop with step-level updates
Integrate truncated sampling into standard policy gradient training.
def train_rag_with_truncated_sampling(model, judge_model, train_tasks,
num_steps=10000, batch_size=32):
"""
Train retrieval-augmented reasoning with truncated step-level sampling.
"""
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for step in range(num_steps):
# Sample tasks from training set
task_batch = random.sample(train_tasks, batch_size)
total_loss = 0.0
for question, reference_answer in task_batch:
# Randomly select which step to optimize
target_step = random.randint(1, 4)
# Generate truncated batch (k candidates for step target_step)
trajectories = generate_truncated_batch(
model,
question,
num_candidates=4,
target_step=target_step
)
# Compute rewards
rewards = [
compute_three_part_reward(traj, question, reference_answer, judge_model)
for traj in trajectories
]
# Compute loss
loss = compute_pg_loss_truncated(model, trajectories, rewards, target_step)
total_loss += loss
# Gradient step
optimizer.zero_grad()
(total_loss / len(task_batch)).backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
if (step + 1) % 100 == 0:
print(f"Step {step + 1}: Loss = {(total_loss / len(task_batch)).item():.4f}")
return model
Step 6: Evaluation on multi-hop QA benchmarks
Benchmark on tasks requiring multiple retrieval steps (TREC, HotpotQA, etc.).
def evaluate_on_multihop_qa(model, eval_tasks):
"""
Evaluate on multi-hop QA: questions requiring 2-4 retrieval steps.
"""
correct = 0
total = 0
for question, reference_answer, num_hops in eval_tasks:
# Execute full trajectory
trajectory = execute_full_rag(model, question, max_steps=num_hops)
# Extract final answer
final_answer = trajectory.split('\n')[-1]
# Check correctness
if final_answer.lower() in reference_answer.lower():
correct += 1
total += 1
accuracy = correct / total
print(f"Multi-hop QA accuracy: {accuracy * 100:.1f}%")
return accuracy
Hyperparameter Selection:
When to Use:
When NOT to Use:
Common Pitfalls: