Super Apriel: One Checkpoint, Many Speeds

Resumen Técnico: Super Apriel

1. El Problema

El costo de inferencia de los modelos de lenguaje grandes (LLM) se ha convertido en un cuello de botella crítico, especialmente para contextos largos y alto rendimiento de servicio (throughput).

• Limitación de la Atención Completa (FA): En la decodificación autoregresiva, la Atención Completa (Full Attention) requiere leer y actualizar una caché de claves y valores (KV) que crece linealmente con la longitud de la secuencia. Esto limita la memoria disponible para el batching y hace que el rendimiento esté limitado por el movimiento de datos (memoria) en lugar de por el cómputo.
• Rigidez de las Arquitecturas Híbridas Existentes: Las soluciones actuales que reemplazan la FA con mecanismos eficientes (como ventanas deslizantes o recurrencias lineales) suelen tener una colocación fija (placement) determinada antes del entrenamiento o durante la conversión. Esto significa que un modelo solo puede operar en un único punto de la curva "velocidad-calidad".
• Necesidad de Flexibilidad: Diferentes cargas de trabajo requieren diferentes equilibrios:
  • Tráfico de alta latencia/bajo contexto necesita máxima calidad.
  • Tráfico de alto volumen/contexto largo necesita máxima velocidad.
  • Cambiar entre estos puntos requiere actualmente múltiples checkpoints, fases de entrenamiento separadas y validaciones repetidas.

2. Metodología

Los autores proponen Super Apriel, un "supernet" (red superpuesta) de 15 mil millones de parámetros que permite cambiar la arquitectura en tiempo de ejecución sin recargar pesos.

• Arquitectura Supernet:

  • Derivado de Apriel 1.6 (un transformador multimodal de 48 capas).
  • Cada capa del decodificador contiene cuatro tipos de "mezcladores" (mixers) entrenados simultáneamente:
    1. FA (Full Attention): Atención estándar.
    2. SWA (Sliding Window Attention): Ventana deslizante (ventana fija de 4096 tokens).
    3. KDA (Kimi Delta Attention): Atención delta con puertas por canal.
    4. GDN (Gated DeltaNet): Regla delta con puertas escalares.
  • Una colocación (placement) selecciona un mezclador por capa. Al cambiar la colocación, se obtiene una arquitectura diferente que comparte la mayoría de los parámetros (FFNs, embeddings, normalización), cambiando solo los bloques de mezcladores.

• Proceso de Entrenamiento:

  1. Destilación Estocástica (S1): Se entrena el supernet distilando desde un profesor congelado (Apriel 1.6). En cada paso, se selecciona aleatoriamente un mezclador por capa. Esto permite que todos los mezcladores aprendan a operar bajo la misma representación compartida.
  2. Optimización de Colocación: Se utiliza un surrogado basado en expansión de clústeres (técnica de física estadística) para modelar la relación entre la asignación de mezcladores y la calidad/velocidad. Esto permite encontrar el óptimo de Pareto (mejor calidad para un costo dado) de manera exacta mediante programación dinámica, explorando un espacio combinatorio masivo ($4^{48}$).
  3. Ajuste Fino Supervisado (SFT): Se realiza un SFT dirigido (targeted) en un subconjunto de colocaciones óptimas identificadas previamente, utilizando datos de instrucción de alta calidad.

• Inferencia:

  • Implementado en vLLM y Fast-LLM.
  • Permite cambiar de "preset" (configuración de colocación) entre solicitudes sin descargar el modelo de disco, solo cargando los bloques de mezcladores necesarios en la GPU.
  • Habilita decodificación especulativa nativa: la colocación de FA actúa como verificador y cualquier colocación eficiente (ej. todo GDN) actúa como modelo borrador, sin necesidad de un modelo externo.

3. Contribuciones Clave

1. Supernet de Atención: El primer checkpoint único que ofrece cuatro opciones de mezcladores entrenados en cada capa, permitiendo múltiples presets de velocidad desde una sola implementación.
2. Optimización de Colocación Exacta: Uso de un modelo sustituto (surrogate) con expansión de clústeres para mapear el paisaje de velocidad-calidad y encontrar los mejores compromisos (frente de Pareto) antes del despliegue.
3. Análisis de Dinámicas del Paisaje: Estudio sobre la estabilidad de las clasificaciones de colocaciones durante el entrenamiento. Descubrieron que, aunque las clasificaciones generales se estabilizan temprano, las colocaciones en el "frente de Pareto" (las más eficientes) son más volátiles a gran escala (15B), advirtiendo contra la extrapolación directa de modelos pequeños (0.5B).
4. Liberación de Código Abierto: Pesos del supernet, código de entrenamiento (Fast-LLM), código de servicio (vLLM) y una herramienta de optimización de colocaciones.

4. Resultados

• Rendimiento vs. Velocidad:
  • El preset todo-FA iguala el rendimiento del profesor Apriel 1.6 en todos los benchmarks.
  • Los presets híbridos recomendados logran un aumento de velocidad de decodificación de 2.9x a 10.7x (a 32k tokens) manteniendo entre el 96% y el 77% de la calidad original.
  • Escalado con Contexto: Las ventajas de velocidad de las colocaciones eficientes se compounding con la longitud del contexto. De 16k a 32k tokens, las colocaciones eficientes ganan un 80-155% de velocidad relativa adicional, superando ampliamente a los híbridos externos (que solo ganan 5-46%).
• Comparativa: Super Apriel supera a modelos híbridos externos de mayor tamaño (como Nemotron-3-Nano 30B o Qwen-3.5 27B) en eficiencia relativa, ofreciendo múltiples puntos de operación desde un solo modelo de 15B.
• Ajuste Fino (SFT): El SFT dirigido mejora significativamente las capacidades de razonamiento matemático en las colocaciones eficientes, recuperando habilidades que la destilación sola no podía preservar completamente.

5. Significado e Impacto

• Flexibilidad Operativa: Elimina la necesidad de entrenar y desplegar múltiples modelos para diferentes necesidades de latencia o calidad. Un solo artefacto puede servir desde tareas de alta precisión hasta inferencia masiva de bajo costo.
• Control en Tiempo de Ejecución: Convierte el compromiso entre velocidad y calidad en un control de runtime (ajustable por carga de trabajo) en lugar de una decisión de diseño estática.
• Eficiencia en Contextos Largos: Demuestra que las arquitecturas híbridas dinámicas son superiores para contextos largos, donde los mecanismos eficientes (recurrencia lineal) evitan los cuellos de botella de memoria de la atención completa.
• Nueva Metodología de Entrenamiento: Introduce un enfoque donde la búsqueda de arquitectura (dónde poner FA vs. eficiente) se desacopla del entrenamiento principal mediante el uso de supernets y optimización basada en sustitutos, permitiendo una exploración exhaustiva del espacio de diseño.

En resumen, Super Apriel representa un avance significativo hacia la infraestructura de IA adaptable, permitiendo a los operadores de modelos ajustar dinámicamente el equilibrio entre rendimiento y costo sin sacrificar la calidad del modelo base ni requerir múltiples versiones del modelo.