Resumen Técnico: Más allá del emparejamiento de salidas en la destilación de NVFP4 para LLM

1. Planteamiento del problema

A medida que los Modelos de Lenguaje de Gran Escala (LLMs) se despliegan en entornos de producción con restricciones de latencia y memoria, surge una demanda creciente de formatos de inferencia de baja precisión, específicamente NVFP4 (un formato de punto flotante de 4 bits con escalado de bloque de grano fino). Si bien la Destilación con Conciencia de Cuantización (QAD) ha emergido como un método robusto para recuperar la precisión perdida durante la cuantización post-entrenamiento (PTQ) mediante la alineación de la distribución de salida de un estudiante cuantizado con un profesor de precisión completa vía divergencia KL, este artículo identifica un fallo crítico en el paradigma prevaleciente.

Los autores argumentan que el alineamiento únicamente a nivel de salida es insuficiente para una recuperación fiel de las capacidades. Su análisis revela que la QAD estándar, que optimiza únicamente la divergencia KL, puede enmascarar una degradación representacional interna significativa. Específicamente, mientras que el modelo estudiante puede lograr emparejar exitosamente los logits de salida del profesor, sus geometrías de activación intermedias experimentan una reorganización no deseada (deriva o drift). Este "aprendizaje de atajos" (shortcut learning) es particularmente severo en modelos que han pasado por un post-entrenamiento de Aprendizaje por Refuerzo (RL), donde las jerarquías semánticas internas necesarias para tareas complejas de razonamiento y codificación se ven comprometidas a pesar de una alta precisión en la salida.

2. Metodología: CKA-QAD

Para abordar la desconexión entre el emparejamiento de la salida y la geometría interna, los autores proponen CKA-QAD, un método que aumenta la QAD estándar con un regularizador de Alineación de Kernel Centrado (CKA).

Mecanismo Central

En lugar de depender exclusivamente del destilado a nivel de logits, CKA-QAD introduce un regularizador ligero que alinea las matrices de Gram por capa de las activaciones del estudiante y del profesor.

• Métrica CKA: El método utiliza CKA lineal, que mide la similitud entre matrices de activación centradas ($\bar{H}_t$ y $\bar{H}_s$). Se elige CKA por su invariancia a las transformaciones ortogonales y al escalado isotrópico, lo que lo hace robusto ante las distribuciones de activación no estacionarias y los cambios de escala inherentes a la cuantización NVFP4.
• Función de Objetivo: La pérdida total combina la pérdida de divergencia KL Top-K estándar ($L_{TopK-KL}$) con un término de alineación CKA dinámico ($L_{CKA}$).
  $$L_{total} = L_{TopK-KL} + \text{sg}\left(\frac{L_{TopK-KL}}{L_{CKA} + \epsilon}\right) L_{CKA}$$
  El uso de un operador de parada de gradiente (stop-gradient o sg) y el escalado dinámico permite equilibrar el término CKA frente a la pérdida de logits sin necesidad de un ajuste manual del hiperparámetro global $\lambda$.

Detalles de Implementación

• Selección de Capas: El CKA se calcula en las salidas de los bloques decodificadores emparejados (post-atención/MLP/actualizaciones de Mamba). Los autores excluyen explícitamente las activaciones de pre-normalización, los pesos de atención y los estados KV para evitar artefactos de las escalas de la corriente residual y las dependencias autoregresivas.
• Optimizaciones de Ingeniería: Para gestionar la sobrecarga computacional de calcular el CKA en muchas capas:
  1. Destilación Top-K: La pérdida KL se calcula solo sobre los 8192 logits superiores del profesor.
  2. CKA en el Espacio de Características: El CKA se calcula mediante matrices de covarianza de características ($\bar{H}^\top \bar{H}$) en lugar de matrices de Gram explícitas de token-a-token, reduciendo la complejidad de memoria.
  3. Captura mediante Hooks: Las activaciones se capturan mediante hooks de propagación hacia adelante (forward hooks), con las activaciones del profesor desvinculadas del grafo de computación.

3. Contribuciones Clave

1. Diagnóstico de Modos de Fallo de QAD: El artículo proporciona un diagnóstico a nivel de representación que muestra cómo la destilación basada solo en KL puede exacerbar la deriva representacional, particularmente en las capas FFN de profundidad media de los Transformers y en los bloques de estado de Mamba. Esta deriva se correlaciona con cuellos de botella de rendimiento en tareas de razonamiento y codificación.
2. CKA-QAD Guiado por NVFP4: La introducción de CKA-QAD, un método que preserva explícitamente la geometría representacional interna durante la destilación de baja cuantización. Restaura el alineamiento a niveles cercanos a BF16 con un overhead mínimo.
3. Validación Empírica: Evaluación exhaustiva en dos arquitecturas distintas:
   • Nemotron 3 Nano: Un modelo híbrido Mamba–Transformer MoE con post-entrenamiento de RL.
   • Qwen3-4B-Thinking-2507: Un modelo compacto de razonamiento basado en Transformer.

4. Resultados Experimentales

Los autores evaluaron CKA-QAD frente a profesores BF16, líneas base de PTQ de NVFP4 y QAD estándar basada solo en KL.

Alineación Representacional

• Nemotron 3 Nano: La QAD estándar (solo KL) resultó en una caída significativa del CKA promedio (0.958) en comparación con PTQ (0.993), con un colapso severo en la última capa (0.740). CKA-QAD restauró el CKA promedio a 0.994 y el CKA de la última capa a 0.985, igualando casi al profesor BF16.
• Qwen3-4B: Se observaron tendencias similares, con la QAD solo-KL reduciendo el CKA promedio a 0.98 y el de la última capa a 0.86. CKA-QAD los restauró a 0.99 y 0.97, respectivamente.

Precisión en Tareas de Destino (Downstream)

CKA-QAD demostró mejoras consistentes en benchmarks de razonamiento complejo y codificación:

• AIME25 (Razonamiento Matemático): Mejoró de 68.5% (solo KL) a 72.3% (CKA-QAD) en Qwen3-4B.
• GPQA-D (Conocimiento Experto): Mejoró de 59.5% a 61.1%.
• LiveCodeBench-v5 (Generación de Código): Mejoró de 57.9% a 59.8%.
• Nemotron 3 Nano: Mostró una recuperación similar, con AIME25 mejorando de 87.9% a 88.3% y LiveCodeBench de 68.9% a 70.3%.

Sobrecarga de Eficiencia

El método introduce costos de entrenamiento modestos:

• Tiempo de ejecución (Wall-clock time): +0.5% por paso.
• VRAM Pico: +7.0%.
• Throughput: -0.5%.
• Convergencia: Los pasos para alcanzar el objetivo de KL se redujeron ligeramente (-3.3%).

5. Significado y Reivindicaciones

El artículo posiciona el alineamiento representacional guiado por CKA como un complemento práctico para el emparejamiento de salida en la recuperación de LLMs cuantizados. Los autores afirman que:

• El alineamiento funcional es necesario pero insuficiente: Emparejar las distribuciones de salida no garantiza la preservación de la topología semántica interna necesaria para el razonamiento profundo.
• La preservación de la geometría ayuda a la generalización de baja cuantización: Al evitar que el estudiante aprenda mapeos de "atajo" que satisfacen la función de pérdida pero degradan la estructura interna, la regularización CKA ayuda a que la cuantización agresiva retenga las jerarquías semánticas finas codificadas durante el post-entrenamiento (especialmente el RL).
• Robustez ante Artefactos de Cuantización: A diferencia de las pérdidas de características punto a punto (ej. MSE), la invariancia de CKA a la escala y rotación lo hace único para NVFP4, donde el escalado por bloques induce distribuciones no estacionarias.

Los autores concluyen que preservar la geometría interna es crítico para tareas que requieren la propagación de características en múltiples pasos, como el razonamiento matemático y la síntesis de código, y que CKA-QAD ofrece un camino viable para el despliegue de LLMs de baja cuantización y alta fidelidad con un overhead de ingeniería mínimo.