Resumen Técnico: Computar Donde Cuenta: Modelos de Lenguaje Auto-Optimizables

Enunciado del Problema

La investigación actual en la eficiencia de la inferencia de Modelos de Lenguaje Grandes (LLM) se centra principalmente en reducir el costo computacional de los pasos individuales de decodificación mediante técnicas estáticas como la cuantización, el podado y la atención dispersa. Estos métodos suelen aplicar un presupuesto de computación uniforme a cada token generado. Sin embargo, el artículo argumenta que la dificultad del token varía significativamente; algunos pasos son localmente predecibles, mientras que otros requieren contexto de largo alcance o cálculos intermedios precisos. Un presupuesto fijo por token conduce a sobrecomputación en pasos fáciles y subcomputación en pasos difíciles. Además, las decisiones de eficiencia no son independientes; una compresión agresiva en el paso $t$ puede degradar el estado de la memoria caché de claves y valores (KV), provocando una degradación de calidad retardada ("contaminación de KV") en pasos futuros. Los métodos existentes de "salida temprana" o adaptables a la profundidad abordan el salto de capas, pero no modulan la ejecución de un paso de profundidad fija en cuanto a contexto, precisión y dispersión.

Metodología: Modelos de Lenguaje Auto-Optimizables (SOL)

Los autores proponen Modelos de Lenguaje Auto-Optimizables (SOL), un marco que aprende a asignar dinámicamente presupuestos de computación por token dentro de un único LLM congelado.

Arquitectura Central

• Modelo Base Congelado: Los parámetros subyacentes del LLM ($\theta$) permanecen sin cambios.
• Red de Política Ligera: Un pequeño transformador autoregresivo (el controlador) lee el estado oculto del LLM y características simples de progreso/presupuesto. Selecciona una acción de eficiencia discreta en cada paso de decodificación.
• Espacio de Acciones Unificado: Cada acción controla simultáneamente tres mecanismos de eficiencia ortogonales:
  1. Dispersión de atención a nivel de token (mediante el método Quest).
  2. Podado estructurado de activaciones MLP (mediante el enmascaramiento de canales estilo TEAL).
  3. Ancho de bits de cuantización de activaciones (mediante cuantización falsa estilo ZeroQuant).
• Horizonte de Control: La decodificación se trata como un problema de decisión secuencial sobre episodios cortos de longitud $T$ (por defecto $T=16$). Para mitigar los efectos acumulativos de la contaminación de KV, el sistema realiza actualizaciones periódicas de la caché KV (recalculando estados densos para los últimos $T$ tokens) después de cada episodio.

Procedimiento de Entrenamiento

La política se entrena utilizando Optimización de Política Relativa por Grupos (GRPO) con episodios forzados por el profesor:

1. Programas Contrafactuales: Para una entrada fija y una trayectoria de token fija (forzamiento por el profesor), el sistema muestrea $K$ programas de computación diferentes (trayectorias) que varían únicamente en sus acciones de eficiencia.
2. Función de Recompensa: La recompensa equilibra la calidad del modelo de lenguaje (log-verosimilitud del token) frente a una penalización suave por desviarse de un objetivo de presupuesto solicitado. El objetivo de presupuesto se proporciona como una característica de entrada a la política, permitiendo que una sola política opere en un rango de regímenes de eficiencia.
3. Asignación de Crédito: Se utiliza un retorno descontado hacia adelante para asignar crédito por efectos retardados dentro del horizonte del episodio.

Inferencia

En el momento de la inferencia, la política selecciona acciones de manera codiciosa basándose en el estado oculto actual y el objetivo de presupuesto solicitado. El controlador incurre en una sobrecarga mínima y constante proporcional al horizonte de control.

Contribuciones Clave

1. Asignación de Presupuesto por Token Aprendible: Introducción de una política ligera que asigna computación dinámicamente en cada paso mientras mantiene el LLM base congelado.
2. Control de Eficiencia Unificado: La política opera en un espacio de acciones unificado que controla simultáneamente la dispersión de atención, el podado de MLP y la cuantización, en lugar de ajustar un único mecanismo.
3. Eficiencia Controlable mediante GRPO: Una receta de entrenamiento práctica que utiliza trayectorias contrafactuales forzadas por el profesor y optimización relativa a grupos para aprender políticas que cumplen de manera fiable los objetivos de presupuesto solicitados mientras maximizan la calidad.
4. Rendimiento Óptimo de Pareto: El método descubre un frente de Pareto calidad-eficiencia superior en comparación con la asignación estática y la búsqueda aleatoria.

Resultados Experimentales

Los autores evaluaron SOL en diversas familias de modelos (Llama-3.2-1B, 3B y Llama-3.1-8B-Instruct) y regímenes de eficiencia.

• Compromiso Calidad-Eficiencia: En todos los experimentos, SOL logra consistentemente una menor perplejidad que las estrategias base fijas (estáticas) con tasas de retención neta equivalentes. Esta mejora persiste a medida que el espacio de acciones crece de 8 a 1.560 opciones discretas.
• Adherencia al Presupuesto: La política rastrea con alta fidelidad los objetivos de presupuesto solicitados (tasa de retención de tokens, tasa de retención de podado y relación de cuantización), cayendo cerca de la línea de identidad en las gráficas de solicitado versus realizado.
• Precisión en Tareas Posteriores: En tareas posteriores (MMLU, ARC, PiQA, Winogrande, GSM8K), SOL mejora la precisión hasta en un 7.3% sobre estrategias de asignación de presupuesto uniforme con presupuestos de computación equivalentes.
• Comparación con Búsqueda: SOL rastrea el envoltorio inferior del paisaje de programas aleatorios y supera la búsqueda aleatoria "mejor de 500" en la mayoría de los puntos de operación probados.
• Sensibilidad al Horizonte: El método permanece efectivo a través de diferentes horizontes de control ($T=4, 16, 64$), con un rendimiento que generalmente mejora a medida que aumenta el horizonte, demostrando robustez frente a estrategias de gestión de la contaminación de KV.

Significado y Afirmaciones

El artículo afirma que SOL abre un eje ortogonal de optimización de eficiencia. A diferencia de trabajos anteriores que seleccionan un único punto de operación estático (por ejemplo, "cuantización de 8 bits" o "50% de dispersión"), SOL aprende cuánta computación gastar por token basándose en la dificultad de la generación y las restricciones de servicio.

Los autores enfatizan que este enfoque permite un control en el momento de la implementación, donde una sola política puede adaptarse a requisitos de presupuesto variables sin reentrenamiento. Al tratar la asignación de computación como un problema de decisión secuencial en lugar de una heurística por paso, SOL aborda las dependencias temporales de las decisiones de eficiencia (contaminación de KV) que los métodos estáticos ignoran. El trabajo sugiere que aprender a "computar donde cuenta" ofrece una vía complementaria a las técnicas de compresión tradicionales para mejorar la eficiencia de la inferencia de LLM.