Resumen Técnico: BCJR-QAT

Enunciado del Problema

Lograr 2 bits por peso (bpw) o menos para Modelos de Lenguaje Grande (LLM) en hardware de consumo es un cuello de botella crítico. Si bien los métodos de Cuantización Post-Entrenamiento (PTQ) como QTIP han establecido una frontera sólida para la cuantización de 2 bits utilizando Cuantización Codificada en Trellis (TCQ), están limitados por el "techo PTQ". Empujar más allá de esto requiere Entrenamiento Consciente de la Cuantización (QAT). Sin embargo, las técnicas estándar de QAT fallan en códigos de trellis porque el algoritmo de decodificación de Viterbi implica una operación argmax no diferenciable sobre un número exponencial de caminos. Aplicar un Estimador Directo (STE) al trellis destruye el acoplamiento entre decisiones de peso adyacentes, que es la ventaja principal de la codificación en trellis sobre la cuantización escalar.

Metodología: BCJR-QAT

El artículo introduce BCJR-QAT, una relajación diferenciable de la cuantización codificada en trellis que reemplaza el argmax duro de Viterbi con una expectativa suave derivada del algoritmo BCJR (Bahl-Cocke-Jelinek-Raviv) hacia adelante-hacia atrás.

Mecanismo Central

1. Analogía de la Mecánica Estadística: El trellis se trata como un grafo de factores equivalente a una cadena de espines tipo Ising unidimensional. El argmax de Viterbi se considera el límite de temperatura cero ($T \to 0$) de una distribución de Boltzmann sobre los caminos.
2. Palabra de Código Suave: A una temperatura finita $T > 0$, el método calcula las probabilidades posteriores marginales exactas de cada estado utilizando el algoritmo BCJR. La salida cuantizada es la expectativa de Boltzmann de la palabra de código emitida:
   $$\hat{w}_t(T) = \langle c(s_t) \rangle_{p_T} = \sum_s c(s) p_T(s_t = s | w)$$
   Esta palabra de código suave es exactamente diferenciable con respecto a la preimagen del peso y los parámetros del trellis.
3. Recocido: Durante el entrenamiento, la temperatura $T$ se recuece desde un valor inicial $T_0$ hasta un valor cercano a cero $T_{end}$. A medida que $T \to 0$, la palabra de código suave converge a la salida dura de Viterbi, "cristalizando" efectivamente la optimización continua en una solución discreta.

Innovaciones Técnicas Clave

• Diferenciabilidad: La recursión hacia adelante-hacia atrás de BCJR (complejidad $O(L \cdot S^2)$) proporciona gradientes exactos a través del trellis, preservando el acoplamiento inter-estado que los métodos STE escalares pierden.
• Kernel Triton: Los autores implementaron un kernel fusionado de Triton que combina operaciones de recolección, log-sum-exp y suma en un solo lanzamiento. Esto reduce el paso hacia adelante/hacia atrás desde cientos de lanzamientos de kernel a uno, logrando una aceleración de 6.57× sobre una implementación de referencia de autograd de PyTorch mientras mantiene la paridad numérica FP32.
• Corrección de Programación (Saltar-Alta-T): El artículo diagnostica un modo de fallo crítico en las programaciones estándar de recocido simulado (comenzando en $T_0=1.0$). A altas temperaturas, el gradiente de la palabra de código suave está "difuminado" e ininformado, causando que el optimizador se desvíe hacia una cuenca de Voronoi peor antes de que la fase de enfriamiento pueda recuperarse. La solución propuesta es saltar la fase de alta-$T$, inicializando a una temperatura más baja (por ejemplo, $T_0=0.3$) donde los gradientes ya son informativos.

Resultados Experimentales

El artículo evalúa BCJR-QAT en dos arquitecturas: OLMoE-1B-7B (MoE) y Llama-3.2-1B (Densa).

1. OLMoE-1B-7B (Resultado Negativo: Brecha de Proxy)

• Configuración: QAT codiciosa por capa optimizando para MSE de reconstrucción (un objetivo proxy).
• Resultado: BCJR-QAT redujo exitosamente el MSE de reconstrucción por capa en un 17–20% en comparación con QTIP-PTQ. Sin embargo, esta mejora no se tradujo en el rendimiento de la tarea final. Los modelos retrocedieron en $\sim 1.3$ de perplejidad (PPL) en WikiText-2 en comparación con la línea base QTIP-PTQ.
• Conclusión: Esto confirma una "brecha de proxy" donde minimizar el error de reconstrucción local no garantiza un mejor rendimiento en la tarea final a 2 bpw.

2. Llama-3.2-1B (Resultado Positivo: KL de Extremo a Extremo)

• Configuración: QAT de una sola capa optimizando para destilación KL hacia adelante de extremo a extremo contra un maestro FP16.
• Resultado:
  • Exceso de Programación: El uso de una programación ingenua de alta-$T$ ($T_0=1.0$) resultó en un rendimiento ligeramente peor que PTQ (+0.005 a +0.022 PPL).
  • Solución Saltar-Alta-T: El uso de la programación corregida ($T_0=0.3$) eliminó el exceso. En la Capa 4, el modelo logró 10.1347 PPL, una mejora de −0.084 PPL sobre la línea base QTIP-PTQ (10.2189).
  • Superaditividad: Al combinar capas entrenadas (Capa 4 y Capa 8), la ganancia conjunta (−0.077 PPL) superó la suma de las ganancias individuales (−0.062 PPL), demostrando un "superávit de cooperación" donde las palabras de código entrenadas con BCJR se refuerzan mutuamente.

3. Análisis Teórico

• Teoría del Presupuesto de Deriva: Los autores derivan una condición para escapar de la cuenca de Voronoi de PTQ: la deriva acumulada de pesos ($\eta \cdot N \cdot g_{max}$) debe exceder el radio de la cuenca ($r_{Voronoi} \approx 10^{-3}$). Los experimentos confirman que las programaciones que no cumplen esta condición no pueden escapar de la solución PTQ, mientras que aquellas que la exceden pueden encontrar mejores soluciones si el objetivo y la programación son correctos.

Significado y Afirmaciones

El artículo afirma tres contribuciones principales:

1. Algorítmica: BCJR-QAT es el primer método que hace que la cuantización codificada en trellis sea nativamente diferenciable, permitiendo QAT en TCQ sin perder los beneficios del acoplamiento de caminos.
2. Ingeniería: El kernel fusionado de Triton hace que esto sea computacionalmente factible en hardware de consumo (una sola GPU de 12 GB), eliminando la barrera de implementación en tiempo de inferencia para la investigación.
3. Insight Empírico:
   • Identifica que el MSE de reconstrucción por capa es un proxy pobre para QAT de 2 bits, mientras que la destilación KL de extremo a extremo es efectiva.
   • Cuestiona la sabiduría convencional del recocido simulado para relajaciones discretas, mostrando que comenzar "caliente" ($T=1.0$) es perjudicial para BCJR-QAT porque el gradiente de alta temperatura está dominado por el ruido. La estrategia óptima es comenzar a una temperatura donde el gradiente ya sea informativo.

El artículo concluye que, aunque los resultados de una sola capa son prometedores, el potencial completo de BCJR-QAT probablemente radique en el entrenamiento conjunto de extremo a extremo a través de todas las capas, una tarea hecha asequible por el kernel liberado pero que requiere hardware a escala de nube (H100/H200) para modelos mayores a 1B de parámetros. Los autores declaran explícitamente que no realizaron este entrenamiento conjunto a gran escala en este trabajo.