Scalable LLM Reasoning Acceleration with Low-rank Distillation

1. El Problema

Los Modelos de Lenguaje Grande (LLM) modernos, especialmente aquellos diseñados para el razonamiento matemático y la resolución de problemas complejos (modelos "de pensamiento" o thinking models), requieren generaciones largas y costosas computacionalmente. Para mitigar esto, se han desarrollado métodos de inferencia eficiente basados en esparsidad (como GRIFFIN y CATS) que seleccionan subconjuntos de neuronas en los bloques de alimentación frontal (Feedforward o FF) para reducir el costo computacional.

Sin embargo, existe una brecha crítica:

• Degradación en Razonamiento: Mientras que estos métodos mantienen un buen rendimiento en tareas de lenguaje estándar (resumen, comprensión lectora), destruyen severamente la capacidad de razonamiento matemático. Los errores de aproximación se acumulan a lo largo de las largas cadenas de pensamiento (Chain-of-Thought o CoT), desviando la trayectoria de generación y resultando en respuestas incorrectas.
• Ineficiencia en la Generación: La generación de CoTs es más costosa que la fase de prellenado (prefill), y los errores en un solo token pueden arruinar todo el proceso.
• Necesidad: Se requiere un método que reduzca la latencia y el número de parámetros activos sin sacrificar la precisión en tareas de razonamiento complejo.

2. Metodología: Caprese

Los autores proponen Caprese (CAPability REcovery with Scalable Efficiency), un método de destilación eficiente diseñado para recuperar las capacidades perdidas al aplicar algoritmos de inferencia eficiente.

Concepto Central

La idea fundamental es que, aunque la compresión de los bloques FF introduce errores, estos errores (residuos) tienen una estructura de bajo rango. En lugar de intentar corregir todo el modelo, Caprese aprende a aproximar estos residuos utilizando capas lineales de bajo rango adicionales.

Arquitectura y Proceso

1. Compresión Base: Se utiliza un algoritmo eficiente existente (como GRIFFIN o CATS) para comprimir el bloque FF original, reduciendo drásticamente el número de parámetros activos.
2. Adición de Capas de Bajo Rango: Se añaden capas lineales de bajo rango ($L$ y $R$) paralelas a la aproximación comprimida. Estas capas están diseñadas para modelar el residuo:
   $$\text{FF}(x) \approx \hat{\text{FF}}(x) + xLR$$
   Donde $\hat{\text{FF}}$ es la versión comprimida y $xLR$ es la corrección de bajo rango.
3. Entrenamiento por Destilación:
   • Distilación por Capas: Se entrena cada bloque FF individualmente para minimizar el error cuadrático medio (MSE) entre la salida del modelo original y la salida del modelo comprimido más la corrección.
   • Distilación de Extremo a Extremo (E2E): Se afinan todas las capas de corrección simultáneamente para minimizar el error en la representación final del modelo antes de la cabeza lineal.
   • Datos de Entrenamiento: El proceso es extremadamente eficiente, requiriendo solo 20,000 muestras sintéticas de matemáticas y aproximadamente un 1% de parámetros adicionales (con una dimensión interna $r=256$).
4. Integración: Las capas de corrección se pueden fusionar o concatenar con las matrices de proyección existentes, permitiendo una inferencia paralela sin sobrecarga significativa de latencia.

3. Contribuciones Clave

• Recuperación de Capacidades: Caprese logra recuperar la mayor parte, si no la totalidad, de la capacidad de razonamiento perdida por métodos de inferencia eficiente, a menudo superando el rendimiento del modelo original comprimido.
• Eficiencia de Recursos:
  • Reduce los parámetros activos en ~2 mil millones (para modelos de 8B-9B) gracias a la compresión subyacente.
  • Añade solo ~0.8% - 1% de parámetros adicionales (las capas de bajo rango).
  • Reduce la latencia de generación en más del 16% (tiempo hasta el siguiente token).
• Brevidad de Respuesta: Curiosamente, el método fomenta respuestas más breves (hasta un 8.5% menos de tokens) sin necesidad de regularización explícita, lo que ahorra aún más recursos.
• Versatilidad: Funciona tanto en modelos de instrucción (instruct) como en modelos de pensamiento (thinking), y es compatible con diferentes técnicas de compresión (GRIFFIN, CATS) y cuantización.

4. Resultados Experimentales

Los experimentos se realizaron en modelos como Llama 3.1/3.2, Gemma 2, DeepSeek-R1-Distill-Qwen y Qwen 2.5.

• Rendimiento en Razonamiento (Matemáticas):
  • En DeepSeek-R1-Distill-Qwen 7B, el método CATS redujo la precisión en AMC 2023 del 75.00% al 62.50%. Con Caprese (E2E), la precisión se recuperó al 78.13%, superando al modelo original.
  • En Gemma 2 2B, la precisión en GSM8K cayó del 51.02% al 34.42% con CATS, pero Caprese la restauró a niveles cercanos o superiores al original.
  • La técnica de re-selección (actualizar dinámicamente los criterios de esparsidad durante la generación) permite que Caprese con GRIFFIN iguale o supere el rendimiento del modelo completo.
• Escalado de Test-Time (Best-of-N):
  • Al escalar la cantidad de generaciones (muestrear múltiples respuestas y elegir la mejor), Caprese mantiene una cobertura (Pass@K) superior o comparable al modelo original, utilizando un 15.8% menos de parámetros activos.
• Tareas de Lenguaje:
  • En tareas de lenguaje estándar (resumen, QA), Caprese no degrada el rendimiento, preservando la utilidad general del modelo.
• Eficiencia de Latencia:
  • En Qwen 2.5 14B, se logró una reducción del 16% en la latencia total y del tiempo por token en comparación con el modelo completo, con una sobrecarga mínima comparada con el uso de solo GRIFFIN.

5. Significado e Impacto

El trabajo de Caprese es significativo porque resuelve la dicotomía actual entre eficiencia y capacidad de razonamiento.

• Viabilidad de Modelos Pequeños: Permite desplegar modelos más pequeños o comprimidos en entornos con recursos limitados sin sacrificar su capacidad para resolver problemas matemáticos complejos.
• Escalabilidad: Facilita el uso de técnicas de escalado en tiempo de prueba (test-time scaling) como CoTs largos o muestreo múltiple, que son prohibitivamente costosos con modelos completos.
• Enfoque Práctico: Al requerir tan pocos datos de entrenamiento y parámetros adicionales, ofrece una solución práctica y de bajo costo para la industria, permitiendo la adopción masiva de inferencia eficiente en aplicaciones críticas de razonamiento.

En resumen, Caprese demuestra que es posible tener inferencia de LLMs rápida, barata y precisa en tareas de razonamiento, cerrando la brecha de rendimiento que existía entre los métodos de compresión y la calidad del modelo original.