Muestrear, No Buscar: Repensando la Alineación en Tiempo de Prueba para Modelos de Lenguaje

Sample, Don't Search: Rethinking Test-Time Alignment for Language Models

April 4, 2025

Autores: Gonçalo Faria, Noah A. Smith

cs.AI

Resumen

El aumento del cómputo en tiempo de prueba ha surgido como una dirección prometedora para mejorar el rendimiento de los modelos de lenguaje, particularmente en escenarios donde el ajuste fino del modelo es poco práctico o imposible debido a limitaciones computacionales o a la privacidad de los pesos del modelo. Sin embargo, los métodos existentes de búsqueda en tiempo de prueba que utilizan un modelo de recompensa (RM, por sus siglas en inglés) suelen degradarse en calidad a medida que escala el cómputo, debido a la sobreoptimización de lo que son, en esencia, proxies de recompensa imperfectos. Introducimos QAlign, un nuevo enfoque de alineación en tiempo de prueba. A medida que escalamos el cómputo en tiempo de prueba, QAlign converge a muestrear la distribución alineada óptima para cada indicación individual. Al adoptar avances recientes en cadenas de Markov Monte Carlo para la generación de texto, nuestro método permite obtener salidas mejor alineadas sin modificar el modelo subyacente ni siquiera requerir acceso a los logits. Demostramos la efectividad de QAlign en benchmarks de razonamiento matemático (GSM8K y GSM-Symbolic) utilizando un RM específico para la tarea, mostrando mejoras consistentes sobre métodos existentes de cómputo en tiempo de prueba como best-of-n y votación por mayoría. Además, cuando se aplica con RMs más realistas entrenados en el conjunto de datos de preferencias Tulu 3, QAlign supera a la optimización directa de preferencias (DPO, por sus siglas en inglés), best-of-n, votación por mayoría y votación por mayoría ponderada en una amplia gama de conjuntos de datos (GSM8K, MATH500, IFEval, MMLU-Redux y TruthfulQA). Como una solución práctica para alinear modelos de lenguaje en tiempo de prueba utilizando cómputo adicional sin degradación, nuestro enfoque expande los límites de la capacidad que se puede obtener de modelos de lenguaje listos para usar sin necesidad de entrenamiento adicional.

English

Increasing test-time computation has emerged as a promising direction for improving language model performance, particularly in scenarios where model finetuning is impractical or impossible due to computational constraints or private model weights. However, existing test-time search methods using a reward model (RM) often degrade in quality as compute scales, due to the over-optimization of what are inherently imperfect reward proxies. We introduce QAlign, a new test-time alignment approach. As we scale test-time compute, QAlign converges to sampling from the optimal aligned distribution for each individual prompt. By adopting recent advances in Markov chain Monte Carlo for text generation, our method enables better-aligned outputs without modifying the underlying model or even requiring logit access. We demonstrate the effectiveness of QAlign on mathematical reasoning benchmarks (GSM8K and GSM-Symbolic) using a task-specific RM, showing consistent improvements over existing test-time compute methods like best-of-n and majority voting. Furthermore, when applied with more realistic RMs trained on the Tulu 3 preference dataset, QAlign outperforms direct preference optimization (DPO), best-of-n, majority voting, and weighted majority voting on a diverse range of datasets (GSM8K, MATH500, IFEval, MMLU-Redux, and TruthfulQA). A practical solution to aligning language models at test time using additional computation without degradation, our approach expands the limits of the capability that can be obtained from off-the-shelf language models without further training.