adu2021 Title: Align Reasoning Trajectories for Factual Correctness
Large reasoning models sometimes identify correct facts during thinking but fail to incorporate them into final answers. MR-Align detects this "reasoning-answer hit gap" by modeling reasoning as transitions between 15 distinct meta-cognitive states (framing, decomposition, verification, etc.). The framework reweights preference optimization based on state transition patterns, amplifying trajectories that consistently lead to factuality.
The key is treating reasoning as a navigable state space, not monolithic text.
Core Concept
State-Transition-Aware Preference Optimization:
- Meta-Reasoning States: 15 cognitive patterns grouped into 4 strategy categories
- Transition Matrices: Track probability of moving between states in factual vs. defective outputs
- Implicit Rewards: Reweight preference signals based on local-to-global transition probabilities
- Segment-Level Supervision: Annotate reasoning with fine-grained cognitive labels
- KTO Adaptation: Enhanced Kahneman-Tversky Optimization with transition-aware weighting
Architecture Overview
- Reasoning Taxonomy: 15 meta-reasoning states (framing, decomposition, chaining, verification, etc.)
- Segment Annotation: LLM-based coarse-to-fine labeling of reasoning text
- Transition Analysis: EM algorithm to estimate state transition matrices from paired examples
- Reward Computation: Token-level signals reweighted by transition probabilities
- Preference Optimization: KTO with segment-aware reward shaping
Implementation Steps
1. Define Meta-Reasoning Taxonomy
Create a taxonomy of cognitive patterns in reasoning.
class MetaReasoningTaxonomy:
# 4 macro-strategies, 15 total cognitive states
STRATEGIES = {
'Meta-Cognitive Regulation': [
'framing', # Problem understanding/context setting
'backtracking', # Reconsidering previous steps
'verification' # Checking intermediate results
],
'Problem-Solving Operations': [
'decomposition', # Breaking down complex problems
'chaining', # Linking multiple steps
'transformation' # Converting problem form
],
'Knowledge Operations': [
'retrieval', # Fetching relevant knowledge
'synthesis', # Combining information
'analogy' # Drawing parallels
],
'Explanatory & Communication': [
'explanation', # Clarifying reasoning
'summarization', # Condensing information
'refinement' # Improving clarity
]
}
@staticmethod
def tag_reasoning_segment(text, llm):
# Use LLM to annotate text with cognitive state
prompt = f"""Label this reasoning segment with ONE cognitive state from:
{', '.join(MetaReasoningTaxonomy.STRATEGIES.keys())}
Segment: {text}
State:"""
state = llm.generate(prompt).strip()
return state
2. Implement Segment-Level Annotation
Parse reasoning into segments and annotate with cognitive states.
def annotate_reasoning_trajectory(reasoning_text, llm, chunk_size=100):
# Split into segments (roughly sentences)
sentences = reasoning_text.split('.')
segments = [s.strip() + '.' for s in sentences if s.strip()]
annotations = []
for segment in segments:
if len(segment.split()) > 10: # Only meaningful segments
state = MetaReasoningTaxonomy.tag_reasoning_segment(segment, llm)
annotations.append({
'text': segment,
'state': state
})
return annotations
def create_positive_negative_pairs(correct_example, incorrect_example, llm):
# Annotate both correct and incorrect reasoning
correct_anno = annotate_reasoning_trajectory(correct_example['reasoning'], llm)
incorrect_anno = annotate_reasoning_trajectory(incorrect_example['reasoning'], llm)
return {
'positive': correct_anno,
'negative': incorrect_anno,
'correct_answer': correct_example['answer'],
'incorrect_answer': incorrect_example['answer']
}
3. Estimate State Transition Matrices
Learn transition probabilities using EM algorithm.
class TransitionMatrixEstimator:
def __init__(self, states):
self.states = states
self.num_states = len(states)
# Initialize uniform transitions
self.P_positive = np.ones((self.num_states, self.num_states)) / self.num_states
self.P_negative = np.ones((self.num_states, self.num_states)) / self.num_states
def state_to_idx(self, state):
return self.states.index(state)
def e_step(self, positive_examples, negative_examples):
# Compute expected transitions given current P estimates
positive_transitions = np.zeros((self.num_states, self.num_states))
negative_transitions = np.zeros((self.num_states, self.num_states))
for example in positive_examples:
for i in range(len(example) - 1):
curr_state = self.state_to_idx(example[i]['state'])
next_state = self.state_to_idx(example[i + 1]['state'])
positive_transitions[curr_state, next_state] += 1
for example in negative_examples:
for i in range(len(example) - 1):
curr_state = self.state_to_idx(example[i]['state'])
next_state = self.state_to_idx(example[i + 1]['state'])
negative_transitions[curr_state, next_state] += 1
return positive_transitions, negative_transitions
def m_step(self, positive_transitions, negative_transitions):
# Update transition matrices from counts
self.P_positive = positive_transitions / positive_transitions.sum(axis=1, keepdims=True)
self.P_negative = negative_transitions / negative_transitions.sum(axis=1, keepdims=True)
def fit(self, positive_examples, negative_examples, num_iterations=5):
for _ in range(num_iterations):
p_trans, n_trans = self.e_step(positive_examples, negative_examples)
self.m_step(p_trans, n_trans)
4. Compute Transition-Aware Rewards
Reweight preference optimization based on transition deviation.
def compute_transition_aware_reward(reasoning_trajectory, P_positive, P_negative):
# For each token, compute how much its transition deviates from global patterns
state_sequence = [seg['state'] for seg in reasoning_trajectory]
rewards = []
for i in range(len(state_sequence) - 1):
curr_idx = state_to_idx(state_sequence[i])
next_idx = state_to_idx(state_sequence[i + 1])
# Probability of this transition in positive/negative examples
prob_positive = P_positive[curr_idx, next_idx]
prob_negative = P_negative[curr_idx, next_idx]
# Reward: how much this transition favors positive over negative
# Higher prob in positive, lower in negative = higher reward
transition_reward = prob_positive / (prob_negative + 1e-8)
# This reward applies to all tokens in the segment
segment_tokens = len(reasoning_trajectory[i]['text'].split())
for _ in range(segment_tokens):
rewards.append(transition_reward)
return rewards
def apply_transition_aware_preference_optimization(
positive_reasoning, negative_reasoning,
P_positive, P_negative,
model, tokenizer
):
# Compute transition-aware rewards
pos_rewards = compute_transition_aware_reward(positive_reasoning, P_positive, P_negative)
neg_rewards = compute_transition_aware_reward(negative_reasoning, P_positive, P_negative)
# Tokenize reasoning
pos_tokens = tokenizer(positive_reasoning_text)
neg_tokens = tokenizer(negative_reasoning_text)
# KTO loss with transition-aware weighting
pos_loss = -pos_rewards.mean() * cross_entropy_loss(model(pos_tokens), pos_tokens)
neg_loss = neg_rewards.mean() * cross_entropy_loss(model(neg_tokens), neg_tokens)
return pos_loss + neg_loss
Practical Guidance
When to Use:
- Reasoning models where reasoning quality doesn't directly correlate with final answer
- Factuality-critical applications
- Fine-tuning models post-training to improve trustworthiness
Hyperparameters:
- num_em_iterations: 5-10 (usually converges quickly)
- segment_min_tokens: 10 (avoid overly granular segmentation)
- transition_reward_exponent: 2.0 (how much to amplify good transitions)
When NOT to Use:
- Models where reasoning text isn't available
- Supervised fine-tuning from scratch (use for alignment only)
- Tasks where reasoning steps don't decompose into meta-cognitive patterns
Pitfalls:
- Sparse transition matrices: With limited examples, EM may estimate unreliable transitions; use Laplace smoothing
- Mislabeled states: Automatic annotation errors propagate to transition matrices; validate manually
- Negative examples quality: If "negative" examples aren't truly inferior, reward learning fails
Integration Strategy: Apply after initial training convergence. Use on 5-10% of data with careful transition matrix validation.