# Trajectory Selection Reasoning > Select best reasoning trajectories from multiple samples using step-level scoring from a 0.6B lightweight verifier that exploits hidden states, outperforming external reward models by 4-12% without massive annotations. - Skill: `adu2021/trajectory-selection-reasoning` (Agent Skill) - Install (CLI): `npx skillmds@latest add adu2021/trajectory-selection-reasoning` - Raw SKILL.md: https://api.skillmd.com/api/skills/adu2021/trajectory-selection-reasoning/raw - Safety review: pending - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: Coding & Dev Tools - License: MIT - Author: adu2021 (https://skillmd.com/u/adu2021) - Updated: 2026-09-17 - Page: https://skillmd.com/skills/adu2021/trajectory-selection-reasoning --- # Technique: Trajectory Selection via Latent Scoring — Efficient Reasoning Verification When sampling multiple reasoning trajectories (chain-of-thought paths), selecting the best one is critical but expensive. Traditional approaches either use majority voting (simple but weak) or external process reward models (effective but costly). TrajSelector exploits the language model's own hidden states to score reasoning steps efficiently. The key insight is that LLMs encode quality signals in their internal representations: good reasoning branches tend to produce specific activation patterns. A tiny 0.6B verifier trained to recognize these patterns can score steps more efficiently than external reward models, requiring no massive step-level annotations. ## Core Concept TrajSelector operates on three principles: - **Hidden State Extraction**: Access intermediate activations during generation - **Step-Level Verification**: Lightweight verifier scores each reasoning step - **Trajectory Aggregation**: Sum step scores to rank complete trajectories - **Data-Driven Training**: Train verifier end-to-end on trajectory pairs, no manual step annotations The result is better than majority voting (+4.6% accuracy) and nearly matches external process reward models (+4.3-12.2%) while being 100× faster. ## Architecture Overview - **Language Model Backbone**: Base 7B-70B LLM generates reasoning trajectories - **Hidden State Extractor**: Capture internal activations at each step - **Lightweight Verifier**: 0.6B model trained to distinguish good/bad steps - **Aggregator**: Sum per-step scores to get trajectory-level ranking - **Training Pipeline**: Preference learning on (better, worse) trajectory pairs ## Implementation Steps The core algorithm extracts hidden states and trains a lightweight verifier. This example shows how to implement step scoring and trajectory selection. ```python import torch import torch.nn as nn from typing import List, Tuple class LightweightStepVerifier(nn.Module): """ 0.6B parameter model that scores reasoning steps. Takes hidden states and outputs quality scores. """ def __init__(self, hidden_dim=4096, output_dim=1): super().__init__() # Lightweight: only 2 layers with moderate width self.mlp = nn.Sequential( nn.Linear(hidden_dim, 512), nn.ReLU(), nn.Dropout(0.1), nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.1), nn.Linear(256, output_dim) ) self.sigmoid = nn.Sigmoid() def forward(self, hidden_states): """ Args: hidden_states: (batch, seq_len, hidden_dim) or (seq_len, hidden_dim) Returns: scores: (batch, seq_len, 1) or (seq_len, 1) - [0, 1] score per step """ logits = self.mlp(hidden_states) scores = self.sigmoid(logits) # Normalize to [0, 1] return scores class TrajectorySelector: """ Select best trajectory from multiple samples using step-level scoring. """ def __init__(self, model, verifier: LightweightStepVerifier): self.model = model self.verifier = verifier def score_trajectory( self, trajectory_tokens: List[int], hidden_states: torch.Tensor ) -> float: """ Compute trajectory score as sum of step scores. Args: trajectory_tokens: list of token IDs hidden_states: (seq_len, hidden_dim) hidden states from model Returns: trajectory_score: float, sum of step scores """ # Get step-level scores step_scores = self.verifier(hidden_states) # (seq_len, 1) # Aggregate: mean score across steps trajectory_score = step_scores.mean().item() return trajectory_score def select_best_trajectory( self, trajectories: List[List[int]], hidden_states_list: List[torch.Tensor] ) -> Tuple[int, float]: """ Score all trajectories and return best. Args: trajectories: list of token sequences (L trajectories) hidden_states_list: list of (seq_len, hidden_dim) tensors Returns: best_idx: index of best trajectory best_score: score of best trajectory """ scores = [] for traj, hidden in zip(trajectories, hidden_states_list): score = self.score_trajectory(traj, hidden) scores.append(score) best_idx = torch.argmax(torch.tensor(scores)).item() best_score = scores[best_idx] return best_idx, best_score def train_step_verifier( verifier: LightweightStepVerifier, model, training_pairs: List[Tuple[List[int], List[int]]], device='cuda' ): """ Train verifier using preference learning on trajectory pairs. Each pair is (better_trajectory, worse_trajectory). """ optimizer = torch.optim.Adam(verifier.parameters(), lr=1e-4) criterion = nn.BCELoss() # Binary: better trajectory = 1, worse = 0 for better_traj, worse_traj in training_pairs: # Generate hidden states for both trajectories with torch.no_grad(): _, better_hidden = model.generate_with_hidden_states(better_traj) _, worse_hidden = model.generate_with_hidden_states(worse_traj) # Verify: better trajectory should score higher better_scores = verifier(better_hidden.to(device)) # (seq, 1) worse_scores = verifier(worse_hidden.to(device)) # (seq, 1) # Mean score across steps better_score_agg = better_scores.mean() worse_score_agg = worse_scores.mean() # Loss: encourage better > worse # Simple ranking loss: penalize when worse scores higher loss = criterion(better_score_agg, torch.tensor(1.0, device=device)) loss += criterion(worse_score_agg, torch.tensor(0.0, device=device)) optimizer.zero_grad() loss.backward() optimizer.step() return loss.item() def best_of_n_with_trajectory_selector( prompt: str, model, verifier: LightweightStepVerifier, n_samples: int = 32, max_steps: int = 100 ): """ Sample N trajectories and select best using TrajSelector. """ trajectories = [] hidden_states_list = [] # Generate N samples for _ in range(n_samples): tokens, hidden = model.generate_with_hidden_states( prompt, max_steps=max_steps ) trajectories.append(tokens) hidden_states_list.append(hidden) # Select best trajectory selector = TrajectorySelector(model, verifier) best_idx, best_score = selector.select_best_trajectory( trajectories, hidden_states_list ) best_trajectory = trajectories[best_idx] answer = model.tokenizer.decode(best_trajectory) return answer, best_idx, best_score ``` The verifier learns to distinguish quality without explicit step-level labels. Preference learning on trajectory pairs is sufficient: it learns that certain activation patterns correlate with correct reasoning. ## Practical Guidance | Scenario | N Samples | Verifier Size | Speedup vs Reward Model | |----------|-----------|---------------|----------------------| | Basic math | 8-16 | 0.6B | 50× faster | | Complex reasoning | 16-32 | 0.6B-1.2B | 30-40× faster | | Code generation | 32-64 | 1.2B | 20-30× faster | **When to Use:** - You generate multiple reasoning samples and need to select best - External reward model is too slow or unavailable - You can collect or have preference-labeled trajectory pairs - Inference latency matters (lightweight verifier is fast) **When NOT to Use:** - Single-shot generation (no samples to choose from) - No training data with trajectory preferences - Accuracy doesn't improve enough to justify verifier cost - Hidden states not accessible (non-standard model implementation) **Common Pitfalls:** - Verifier overfits to training distribution → use diverse trajectory sources - Hidden state dimension mismatch → ensure model output layer dims correct - Not normalizing hidden states → add layer norm before verifier input - Training on identical trajectories → use hard negatives (similar wrong vs right) - Aggregating over full sequence length including padding → mask padding tokens ## Reference [TrajSelector: Harnessing Latent Representations for Efficient Best-of-N in LRMs](https://arxiv.org/abs/2510.16449)