Decodificación Especulativa Guiada por Recompensa para un Razonamiento Eficiente en LLM

Reward-Guided Speculative Decoding for Efficient LLM Reasoning

January 31, 2025

Autores: Baohao Liao, Yuhui Xu, Hanze Dong, Junnan Li, Christof Monz, Silvio Savarese, Doyen Sahoo, Caiming Xiong

cs.AI

Resumen

Presentamos Decodificación Especulativa Guiada por Recompensa (RSD), un marco novedoso destinado a mejorar la eficiencia de la inferencia en modelos de lenguaje grandes (LLMs). RSD combina de manera sinérgica un modelo de borrador ligero con un modelo objetivo más potente, incorporando un sesgo controlado para priorizar salidas de alta recompensa, a diferencia de los métodos de decodificación especulativa existentes que imponen una imparcialidad estricta. RSD emplea un modelo de recompensa de proceso para evaluar pasos de decodificación intermedios y decidir dinámicamente si invocar al modelo objetivo, optimizando el equilibrio entre el costo computacional y la calidad de la salida. Demostramos teóricamente que una estrategia de mezcla basada en umbrales logra un equilibrio óptimo entre la utilización de recursos y el rendimiento. Evaluaciones extensas en desafiantes bancos de pruebas de razonamiento, incluidas tareas de nivel olímpico, muestran que RSD proporciona ganancias significativas de eficiencia en comparación con la decodificación solo con el modelo objetivo (hasta 4.4 veces menos FLOPs), al tiempo que logra una precisión significativamente mejor que el método de decodificación paralela en promedio (hasta +3.5). Estos resultados destacan a RSD como un enfoque sólido y rentable para implementar LLMs en escenarios intensivos en recursos.

English

We introduce Reward-Guided Speculative Decoding (RSD), a novel framework aimed at improving the efficiency of inference in large language models (LLMs). RSD synergistically combines a lightweight draft model with a more powerful target model, incorporating a controlled bias to prioritize high-reward outputs, in contrast to existing speculative decoding methods that enforce strict unbiasedness. RSD employs a process reward model to evaluate intermediate decoding steps and dynamically decide whether to invoke the target model, optimizing the trade-off between computational cost and output quality. We theoretically demonstrate that a threshold-based mixture strategy achieves an optimal balance between resource utilization and performance. Extensive evaluations on challenging reasoning benchmarks, including Olympiad-level tasks, show that RSD delivers significant efficiency gains against decoding with the target model only (up to 4.4x fewer FLOPs), while achieving significant better accuracy than parallel decoding method on average (up to +3.5). These results highlight RSD as a robust and cost-effective approach for deploying LLMs in resource-intensive scenarios.