Reasoning in LLMs
1. The Reasoning Challenge
Reasoning is the ability to derive new conclusions from known premises through logical steps. For LLMs, reasoning manifests as generating intermediate steps before arriving at a final answer.
1.1 Formal Definition
Given a problem requiring reasoning, the model generates a sequence of reasoning steps:
where are intermediate reasoning steps and is the final answer. The probability of the complete reasoning chain is:
where .
2. Chain-of-Thought (CoT) Prompting
2.1 Standard CoT
Chain-of-thought prompting (Wei et al., 2022) elicits step-by-step reasoning by including examples with intermediate steps in the prompt:
The model is trained to produce reasoning traces like:
Q: Roger has 5 tennis balls. He buys 2 more cans of 3 tennis balls each. How many does he have now?
A: Roger started with 5 balls. 2 cans of 3 tennis balls each is 2 * 3 = 6 tennis balls. 5 + 6 = 11. The answer is 11.
2.2 Zero-Shot CoT
Simply adding "Let's think step by step" to the prompt triggers reasoning:
This surprisingly effective technique activates the model's latent reasoning capabilities.
2.3 Self-Consistency
Self-consistency (Wang et al., 2023) samples multiple reasoning paths and takes the majority vote:
where is the -th sampled reasoning path and is its answer.
The probability of the correct answer increases with the number of samples:
where is the probability of the most likely reasoning path.
3. Process Reward Models (PRM)
3.1 Outcome vs Process Supervision
Outcome Reward Model (ORM): Scores only the final answer:
Process Reward Model (PRM): Scores each reasoning step:
3.2 PRM Training
Given a dataset of reasoning traces with step-level annotations:
where indicates whether step is correct.
The PRM is trained with binary cross-entropy at each step:
3.3 PRM vs ORM Comparison
| Aspect | ORM | PRM |
|---|---|---|
| Annotation | Final answer only | Each step |
| Training signal | Sparse | Dense |
| Credit assignment | Difficult | Clear |
| Sample efficiency | Lower | Higher |
| Annotation cost | Lower | Higher |
4. Monte Carlo Tree Search (MCTS) for Reasoning
4.1 MCTS Overview
MCTS (Coulom, 2006) builds a search tree by balancing exploration and exploitation:
4.2 MCTS Algorithm for Reasoning
function MCTS_Reasoning(q, LLM, PRM, max_iterations N):
root = Node(q)
for i = 1 to N:
node = Select(root) // UCT selection
child = Expand(node, LLM) // Generate next reasoning step
reward = Evaluate(child, PRM) // PRM score
Backpropagate(child, reward) // Update statistics
return BestChild(root) // Return most visited child
4.3 UCT (Upper Confidence Bound for Trees)
where:
- : estimated value (average reward)
- : visit count of parent node
- : visit count of action
- : exploration constant (typically )
4.4 Reasoning as Search
In reasoning, the search space is the space of all possible reasoning chains:
5. DeepSeek-R1 Style Training
5.1 Overview
DeepSeek-R1 (2024) demonstrated that reasoning capabilities can emerge through reinforcement learning without explicit chain-of-thought supervision.
5.2 Training Pipeline
Phase 1: Cold Start Start with a base model and fine-tune on a small set of reasoning demonstrations:
Phase 2: Reasoning RL Apply RL with a rule-based reward (correctness of final answer):
where is 1 if the answer matches ground truth, 0 otherwise.
Phase 3: Rejection Sampling Generate many reasoning traces, select those with correct answers, and fine-tune:
Phase 4: Alignment RL Final RL stage with additional helpfulness and safety rewards.
5.3 Reward Shaping
The composite reward combines correctness and format:
where:
- : binary correctness
- : reward for proper
<think>...</think>formatting - : penalty for excessive length
6. Test-Time Compute Scaling
6.1 The Insight
Traditional scaling laws focus on training compute:
Test-time compute scaling shows that inference compute can also improve performance:
6.2 Methods for Test-Time Scaling
Best-of-N sampling: Generate reasoning traces, select the best:
Beam search: Maintain top- reasoning candidates at each step:
MCTS: Search over reasoning space (Section 4).
Iterative refinement: Generate, critique, revise:
6.3 Scaling Curve
The relationship between test-time compute and performance follows a power law:
where determines the scaling exponent and is the baseline accuracy.
6.4 Compute-Optimal Allocation
Given a total budget , allocate between training and testing:
The optimal allocation depends on the task:
- Easy tasks: More training, less test compute
- Hard tasks: More test compute (search, sampling)
7. Reasoning Distillation
7.1 Process Distillation
Distill reasoning from a stronger teacher to a weaker student:
7.2 Step-Level Distillation
For each reasoning step, train the student to match the teacher:
8. Implementation
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
class ReasoningEngine:
def __init__(self, model_name, prm_name=None):
self.model = AutoModelForCausalLM.from_pretrained(model_name)
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.prm = load_prm(prm_name) if prm_name else None
def generate_with_cot(self, question, n_samples=1, temperature=1.0):
prompt = f"Question: {question}\nLet's think step by step.\n"
samples = []
for _ in range(n_samples):
inputs = self.tokenizer(prompt, return_tensors="pt")
outputs = self.model.generate(
**inputs,
max_new_tokens=2048,
temperature=temperature,
do_sample=True,
top_p=0.95
)
response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
samples.append(response)
return samples
def self_consistency(self, question, n_samples=32):
samples = self.generate_with_cot(question, n_samples)
answers = [self.extract_answer(s) for s in samples]
# Majority vote
from collections import Counter
counts = Counter(answers)
best_answer = counts.most_common(1)[0][0]
confidence = counts[best_answer] / len(answers)
return best_answer, confidence, samples
def mcts_reasoning(self, question, max_iterations=100):
root = ReasoningNode(question)
for _ in range(max_iterations):
# Select
node = self.select(root)
# Expand
child = self.expand(node)
# Evaluate
if self.prm:
reward = self.prm.score(question, child.trace)
else:
reward = self.rollout(child)
# Backprop
self.backpropagate(child, reward)
return self.best_child(root)
def select(self, node):
while not node.is_leaf():
node = max(node.children, key=lambda c: self.uct(c))
return node
def uct(self, node, c=1.414):
if node.visits == 0:
return float('inf')
return node.value + c * (math.log(node.parent.visits) / node.visits) ** 0.5
def expand(self, node):
# Generate next reasoning step
prompt = node.get_prompt()
next_step = self.generate_step(prompt)
child = ReasoningNode(
parent=node,
step=next_step,
trace=node.trace + [next_step]
)
node.children.append(child)
return child
9. Evaluation
9.1 Accuracy Metrics
Final Answer Accuracy:
Step Accuracy:
9.2 Reasoning Quality
Faithfulness: Does the reasoning lead to the correct answer?
Coherence: Are the reasoning steps logically connected?
9.3 Efficiency Metrics
Samples to solve: Number of samples needed to get correct answer:
10. Open Challenges
- Scalability: MCTS is expensive for long reasoning chains
- Generalization: Reasoning skills may not transfer across domains
- Verification: How to verify correctness without ground truth
- Compositionality: Combining reasoning skills
- Faithfulness: Ensuring reasoning reflects actual computation
References
- Wei et al. (2022). "Chain-of-Thought Prompting Elicits Reasoning in Large Language Models." NeurIPS.
- Wang et al. (2023). "Self-Consistency Improves Chain of Thought Reasoning in Language Models." ICLR.
- Lightman et al. (2023). "Let's Verify Step by Step." ICLR.
- DeepSeek-AI (2024). "DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning." arXiv.
- Snell et al. (2024). "Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters." arXiv.
- Yao et al. (2023). "Tree of Thoughts: Deliberate Problem Solving with Large Language Models." arXiv.