R2R: Navegación Eficiente de Rutas de Razonamiento Divergentes con Enrutamiento de Tokens de Modelos Pequeños-Grandes

R2R: Efficiently Navigating Divergent Reasoning Paths with Small-Large Model Token Routing

May 27, 2025

Autores: Tianyu Fu, Yi Ge, Yichen You, Enshu Liu, Zhihang Yuan, Guohao Dai, Shengen Yan, Huazhong Yang, Yu Wang

cs.AI

Resumen

Los Modelos de Lenguaje de Gran Escala (LLMs, por sus siglas en inglés) logran capacidades de razonamiento impresionantes a costa de un sobrecosto sustancial en la inferencia, lo que plantea desafíos significativos en su implementación. Aunque los Modelos de Lenguaje Pequeños (SLMs) destilados mejoran considerablemente la eficiencia, su rendimiento se ve afectado al no seguir las rutas de razonamiento de los LLMs. Afortunadamente, descubrimos que solo una pequeña fracción de tokens diverge genuinamente las rutas de razonamiento entre LLMs y SLMs. La mayoría de los tokens generados son idénticos o presentan diferencias neutrales, como variaciones menores en abreviaturas o expresiones. Aprovechando esta observación, presentamos **Roads to Rome (R2R)**, un método de enrutamiento neuronal de tokens que utiliza selectivamente los LLMs solo para estos tokens críticos que divergen las rutas, mientras delega la generación de la mayoría de los tokens al SLM. También desarrollamos una pipeline automática de generación de datos que identifica los tokens divergentes y genera etiquetas de enrutamiento a nivel de token para entrenar el enrutador ligero. Aplicamos R2R para combinar los modelos R1-1.5B y R1-32B de la familia DeepSeek, y evaluamos en benchmarks desafiantes de matemáticas, codificación y preguntas y respuestas. Con un tamaño promedio de parámetros activados de 5.6B, R2R supera la precisión promedio de R1-7B en 1.6x, superando incluso al modelo R1-14B. En comparación con R1-32B, ofrece una aceleración de 2.8x en tiempo de ejecución con un rendimiento comparable, avanzando la frontera de Pareto en la eficiencia de escalado en tiempo de prueba. Nuestro código está disponible en https://github.com/thu-nics/R2R.

English

Large Language Models (LLMs) achieve impressive reasoning capabilities at the cost of substantial inference overhead, posing substantial deployment challenges. Although distilled Small Language Models (SLMs) significantly enhance efficiency, their performance suffers as they fail to follow LLMs' reasoning paths. Luckily, we reveal that only a small fraction of tokens genuinely diverge reasoning paths between LLMs and SLMs. Most generated tokens are either identical or exhibit neutral differences, such as minor variations in abbreviations or expressions. Leveraging this insight, we introduce **Roads to Rome (R2R)**, a neural token routing method that selectively utilizes LLMs only for these critical, path-divergent tokens, while leaving the majority of token generation to the SLM. We also develop an automatic data generation pipeline that identifies divergent tokens and generates token-level routing labels to train the lightweight router. We apply R2R to combine R1-1.5B and R1-32B models from the DeepSeek family, and evaluate on challenging math, coding, and QA benchmarks. With an average activated parameter size of 5.6B, R2R surpasses the average accuracy of R1-7B by 1.6x, outperforming even the R1-14B model. Compared to R1-32B, it delivers a 2.8x wall-clock speedup with comparable performance, advancing the Pareto frontier of test-time scaling efficiency. Our code is available at https://github.com/thu-nics/R2R.