adu2021 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 batch
- Branching: Generate k different step-t candidates from shared prefix
- Evaluation: LLM judges each step t independently on three dimensions
- Reward Aggregation: Combine three reward signals (reasoning + query + answer)
- Policy Update: Gradient step on isolated step with reduced variance
- Output: 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': question},
{'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 steps
# Example: "Who was the first person to reach the South Pole?"
# Step 0: User question
# Step 1: Search "first person South Pole" → retrieve docs about Amundsen, Scott
# Step 2: Synthesize: "Roald Amundsen reached it in 1911"
# Step 3: Formulate follow-up: "Roald Amundsen nationality" (verify credibility)
# Step 4: Retrieve clarification docs
# Step 5: Generate final answer with sources
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
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.
When to Use:
- 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
When NOT to Use:
- 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
Common Pitfalls:
- 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.