TrajSelector: Aproveitando Representações Latentes para Seleção Eficiente e Eficaz do Melhor de N em Modelos de Raciocínio em Grande Escala

TrajSelector: Harnessing Latent Representations for Efficient and Effective Best-of-N in Large Reasoning Model

October 18, 2025

Autores: Bin Yu, Xinming Wang, Shijie Lian, Haotian Li, Changti Wu, Ruina Hu, Bailing Wang, Yuliang Wei, Kai Chen

cs.AI

Resumo

Modelos de linguagem de grande escala (LLMs) têm demonstrado progressos notáveis em tarefas de raciocínio complexo, em grande parte possibilitados por paradigmas de escalonamento em tempo de teste (TTS) que alocam recursos computacionais adicionais durante a inferência. Entre esses, o TTS externo (particularmente o paradigma de seleção Best-of-N) proporciona melhorias escaláveis de desempenho ao selecionar entre múltiplas trajetórias de raciocínio geradas de forma independente. No entanto, essa abordagem enfrenta limitações-chave: (i) a alta sobrecarga computacional de implantar modelos de recompensa de processo, (ii) a subutilização das representações latentes intrínsecas do LLM. Introduzimos o TrajSelector, um framework Best-of-N eficiente e eficaz que explora os estados ocultos no LLM amostrador para pontuação em nível de processo. Um verificador leve (com apenas 0,6 bilhões de parâmetros) avalia a qualidade da trajetória passo a passo e, em seguida, agrega essas pontuações para identificar a trajetória de raciocínio ótima. Nosso framework emprega uma receita de treinamento totalmente orientada por dados e de ponta a ponta, eliminando a dependência de anotações massivas em nível de etapa. Resultados experimentais em cinco benchmarks demonstram que o TrajSelector proporciona ganhos consistentes de desempenho. Em configurações Best-of-32, ele supera a votação majoritária em 4,61% de precisão e supera os modelos de recompensa de processo existentes em 4,31% a 12,21%, tudo isso mantendo custos de inferência mais baixos.

English

Large language models (LLMs) have shown remarkable progress in complex reasoning tasks, largely enabled by test-time scaling (TTS) paradigms that allocate additional compute during inference. Among these, external TTS (particularly the Best-of-N selection paradigm) yields scalable performance improvements by selecting from multiple independently generated reasoning trajectories. However, this approach faces key limitations: (i) the high computational overhead of deploying process reward models, (ii) the underutilization of the LLM's intrinsic latent representations. We introduce TrajSelector, an efficient and effective Best-of-N framework that exploit the hidden states in the sampler LLM for process-level scoring. A lightweight verifier (with only 0.6B parameters) evaluates the quality of step-wise trajectory, and then aggregates these scores to identify the optimal reasoning trajectory. Our framework employs a fully data-driven, end-to-end training recipe that eliminates reliance on massive step-level annotations. Experiential results across five benchmarks demonstrate that TrajSelector delivers consistent performance gains. In Best-of-32 settings, it surpasses majority voting by 4.61% accuracy and outperforms existing process reward models by 4.31% to 12.21%, all while maintaining lower inference costs.