Decodificação Especulativa Guiada por Recompensa para Raciocínio Eficiente em 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

Resumo

Apresentamos a Decodificação Especulativa Guiada por Recompensa (RSD), um novo framework destinado a melhorar a eficiência da inferência em grandes modelos de linguagem (LLMs). A RSD combina de forma sinérgica um modelo preliminar leve com um modelo alvo mais poderoso, incorporando um viés controlado para priorizar saídas de alta recompensa, em contraste com os métodos de decodificação especulativa existentes que impõem imparcialidade estrita. A RSD utiliza um modelo de recompensa de processo para avaliar etapas de decodificação intermediárias e decidir dinamicamente se deve invocar o modelo alvo, otimizando o equilíbrio entre custo computacional e qualidade de saída. Demonstramos teoricamente que uma estratégia de mistura baseada em limiar alcança um equilíbrio ótimo entre utilização de recursos e desempenho. Avaliações extensivas em benchmarks desafiadores de raciocínio, incluindo tarefas de nível olímpico, mostram que a RSD proporciona ganhos significativos de eficiência em comparação com a decodificação apenas com o modelo alvo (até 4,4 vezes menos FLOPs), enquanto alcança uma precisão significativamente melhor do que o método de decodificação paralela em média (até +3,5). Esses resultados destacam a RSD como uma abordagem robusta e econômica para implantar LLMs em cenários intensivos em 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.