trajectory-selection-reasoning

name	trajectory-selection-reasoning
title	TrajSelector: Harnessing Latent Representations for Efficient Best-of-N in LRMs
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2510.16449
keywords	["best-of-N selection","trajectory scoring","process rewards","latent representations","reasoning LLMs"]
description	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.

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.

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.
    """

     ():
        ().__init__()
        
        .mlp = nn.Sequential(
            nn.Linear(hidden_dim, ),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(, ),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(, output_dim)
        )
        .sigmoid = nn.Sigmoid()

     ():
        
        logits = .mlp(hidden_states)
        scores = .sigmoid(logits)  
         scores


 :
    

     ():
        .model = model
        .verifier = verifier

     () -> :
        
        
        step_scores = .verifier(hidden_states)  

        
        trajectory_score = step_scores.mean().item()
         trajectory_score

     () -> [, ]:
        
        scores = []
         traj, hidden  (trajectories, hidden_states_list):
            score = .score_trajectory(traj, hidden)
            scores.append(score)

        best_idx = torch.argmax(torch.tensor(scores)).item()
        best_score = scores[best_idx]
         best_idx, best_score


 ():
    
    optimizer = torch.optim.Adam(verifier.parameters(), lr=)
    criterion = nn.BCELoss()  

     better_traj, worse_traj  training_pairs:
        
         torch.no_grad():
            _, better_hidden = model.generate_with_hidden_states(better_traj)
            _, worse_hidden = model.generate_with_hidden_states(worse_traj)

        
        better_scores = verifier(better_hidden.to(device))  
        worse_scores = verifier(worse_hidden.to(device))   

        
        better_score_agg = better_scores.mean()
        worse_score_agg = worse_scores.mean()

        
        
        loss = criterion(better_score_agg, torch.tensor(, device=device))
        loss += criterion(worse_score_agg, torch.tensor(, device=device))

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

     loss.item()


 ():
    
    trajectories = []
    hidden_states_list = []

    
     _  (n_samples):
        tokens, hidden = model.generate_with_hidden_states(
            prompt,
            max_steps=max_steps
        )
        trajectories.append(tokens)
        hidden_states_list.append(hidden)

    
    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)

     answer, best_idx, best_score