Skip to main content ホーム クリエイター adu2021 skillxiv truncated-step-level-retrieval-reasoning
truncated-step-level-retrieval-reasoning 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.
インストールへ移動 Skills Marketplace コミュニティが作成したAIスキルを発見・探索
Codex または Claude でインストール この Prompt をコピーして Codex、Claude、または他のアシスタントに貼り付けると、Skill ページを確認してインストールできます。
直接コマンドでは確認用 Prompt が省略されます。実行前にソースを確認してください。
npx skills add https://github.com/ADu2021/skillXiv --skill truncated-step-level-retrieval-reasoningコマンドは1行のまま表示されます。コピー前に横へスクロールして全体を確認してください。
ローカルで確認しますか？SkillsMP が現在取得できるファイルをダウンロードできます。
Zipをダウンロード ダウンロード中... 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.
Truncated Step-Level Sampling for Retrieval-Augmented Reasoning
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.
Core Concept
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:
Reasoning quality (is the thinking sound?)
Query quality (is the search formulation effective?)
Answer correctness (does the final answer match?)
These two components combine to provide stable, interpretable learning signals.
Architecture Overview
Input : Multi-hop retrieval task (question requiring multiple document lookups)Prefix Generation : Execute steps 1 to t-1 deterministically or sample once, reuse across batchBranching : Generate k different step-t candidates from shared prefixEvaluation : LLM judges each step t independently on three dimensionsReward Aggregation : Combine three reward signals (reasoning + query + answer)Policy Update : Gradient step on isolated step with reduced varianceOutput : Improved retrieval and reasoning policy
Implementation Steps
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 = [
{'step'
0
'action'
'formulate_initial_query'
'state'
'step'
1
'action'
'retrieve_documents'
'state'
None
'step'
2
'action'
'synthesize_reasoning'
'state'
None
'step'
3
'action'
'formulate_followup_query'
'state'
None
'step'
4
'action'
'retrieve_evidence'
'state'
None
'step'
5
'action'
'generate_final_answer'
'state'
None
return
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)
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
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)
trajectories = []
for candidate in candidates:
trajectory = {
'prefix' : prefix_state,
'step_t_action' : candidate,
'target_step' : target_step,
'full_trajectory' : None
}
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' ]
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
)
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
final_answer = full_trajectory.split('\n' )[-1 ]
is_correct = final_answer.lower() in reference_answer.lower()
answer_score = 1.0 if is_correct else 0.0
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
}
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):
step_action = trajectory['step_t_action' ]
prefix_context = trajectory['prefix' ]
logits = model.forward(
f"Context: {prefix_context} \nAction:" ,
output_logits=True
)
action_logprob = compute_logprob(logits, step_action, model.tokenizer)
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
"""
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])
T = 5
expected_ratio = T
observed_ratio = var_full / var_truncated
print (f"Variance reduction: {observed_ratio:.1 f} 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):
task_batch = random.sample(train_tasks, batch_size)
total_loss = 0.0
for question, reference_answer in task_batch:
target_step = random.randint(1 , 4 )
trajectories = generate_truncated_batch(
model,
question,
num_candidates=4 ,
target_step=target_step
)
rewards = [
compute_three_part_reward(traj, question, reference_answer, judge_model)
for traj in trajectories
]
loss = compute_pg_loss_truncated(model, trajectories, rewards, target_step)
total_loss += loss
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():.4 f} " )
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:
trajectory = execute_full_rag(model, question, max_steps=num_hops)
final_answer = trajectory.split('\n' )[-1 ]
if final_answer.lower() in reference_answer.lower():
correct += 1
total += 1
accuracy = correct / total
print (f"Multi-hop QA accuracy: {accuracy * 100 :.1 f} %" )
return accuracy
Practical Guidance Hyperparameter Selection:
Number of candidates k : 2-4. More candidates = better coverage; diminishing returns beyond 4.Target step selection : Sample uniformly from [1, T] for balanced training, or focus on early steps for faster learning.Reward weights : Equal [1/3, 1/3, 1/3] standard; can emphasize answer_score for final tuning.Judge model temperature : 0.5-0.7 for stable scoring; lower = more consistent, higher = more nuanced.
Multi-hop retrieval tasks (2-5 steps) where single-step methods fail
Scenarios where every reasoning step affects downstream success
Settings with access to LLM judge for step-level scoring
Tasks where variance reduction significantly improves convergence
Single-hop retrieval (no multi-step credit assignment benefit)
Tasks with very long trajectories (T > 10); variance reduction diminishes
Real-time systems where multiple forward passes are expensive
Scenarios without reliable reference answers for reward feedback
Inconsistent judge scoring : If judge is non-deterministic, rewards are noisy. Use deterministic judge or averaging over multiple runs.Step isolation breaks dependencies : Some steps may be impossible without specific prefix context. Validate that sampled candidates are valid.Reward signal too sparse : If most samples have zero answer_score, provide auxiliary rewards (query quality, reasoning soundness) to maintain signal.Target step bias : If always sampling early steps, later steps never get optimized. Ensure uniform or curriculum-based step selection.
Reference 