Escalar de manera óptima el cálculo del tiempo de prueba de LLM puede ser más efectivo que escalar los parámetros del modelo.

Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters

August 6, 2024

Autores: Charlie Snell, Jaehoon Lee, Kelvin Xu, Aviral Kumar

cs.AI

Resumen

Permitir a los LLM mejorar sus salidas mediante el uso de más cálculos en tiempo de prueba es un paso crítico hacia la construcción de agentes generalmente auto-mejoradores que puedan operar en lenguaje natural de forma abierta. En este documento, estudiamos la escalabilidad de los cálculos en tiempo de inferencia en LLMs, centrándonos en responder a la pregunta: si a un LLM se le permite utilizar una cantidad fija pero significativa de cálculos en tiempo de inferencia, ¿cuánto puede mejorar su rendimiento en una tarea desafiante? Responder a esta pregunta tiene implicaciones no solo en el rendimiento alcanzable de los LLMs, sino también en el futuro del preentrenamiento de LLM y cómo se debe equilibrar el cálculo en tiempo de inferencia y el preentrenamiento. A pesar de su importancia, hay poca investigación que haya intentado comprender los comportamientos de escalabilidad de varios métodos de inferencia en tiempo de prueba. Además, el trabajo actual en su mayoría proporciona resultados negativos para varias de estas estrategias. En este trabajo, analizamos dos mecanismos principales para escalar el cálculo en tiempo de prueba: (1) buscar en modelos de recompensa verificadores densos basados en procesos; y (2) actualizar la distribución del modelo sobre una respuesta de forma adaptativa, dada la tarea en tiempo de prueba. Descubrimos que, en ambos casos, la efectividad de diferentes enfoques para escalar el cálculo en tiempo de prueba varía críticamente según la dificultad de la tarea. Esta observación motiva la aplicación de una estrategia de escalado "óptima en cálculo", que actúa para asignar de manera más efectiva el cálculo en tiempo de prueba de forma adaptativa por tarea. Utilizando esta estrategia óptima en cálculo, podemos mejorar la eficiencia del escalado del cálculo en tiempo de prueba en más de 4 veces en comparación con un modelo de referencia de mejor de N. Además, en una evaluación de FLOPs equiparada, descubrimos que en problemas donde un modelo base más pequeño alcanza tasas de éxito algo significativas, el cálculo en tiempo de prueba puede utilizarse para superar a un modelo 14 veces más grande.

English

Enabling LLMs to improve their outputs by using more test-time computation is a critical step towards building generally self-improving agents that can operate on open-ended natural language. In this paper, we study the scaling of inference-time computation in LLMs, with a focus on answering the question: if an LLM is allowed to use a fixed but non-trivial amount of inference-time compute, how much can it improve its performance on a challenging prompt? Answering this question has implications not only on the achievable performance of LLMs, but also on the future of LLM pretraining and how one should tradeoff inference-time and pre-training compute. Despite its importance, little research attempted to understand the scaling behaviors of various test-time inference methods. Moreover, current work largely provides negative results for a number of these strategies. In this work, we analyze two primary mechanisms to scale test-time computation: (1) searching against dense, process-based verifier reward models; and (2) updating the model's distribution over a response adaptively, given the prompt at test time. We find that in both cases, the effectiveness of different approaches to scaling test-time compute critically varies depending on the difficulty of the prompt. This observation motivates applying a "compute-optimal" scaling strategy, which acts to most effectively allocate test-time compute adaptively per prompt. Using this compute-optimal strategy, we can improve the efficiency of test-time compute scaling by more than 4x compared to a best-of-N baseline. Additionally, in a FLOPs-matched evaluation, we find that on problems where a smaller base model attains somewhat non-trivial success rates, test-time compute can be used to outperform a 14x larger model.