Depth Registers Unlock W4A4 on SwiGLU: A Reader/Generator Decomposition

1. El Problema

El artículo aborda el desafío crítico de la cuantización post-entrenamiento W4A4 (4 bits para pesos y 4 bits para activaciones) en modelos de lenguaje basados en la arquitectura SwiGLU (utilizada en modelos como LLaMA, Mistral y Qwen).

• Fallo catastrófico: En un modelo controlado de 300M parámetros entrenado con 5 mil millones de tokens, la cuantización W4A4 "ingenua" (redondeo al más cercano) colapsa la perplejidad de validación de 23.6 (FP16) a 1727, lo que representa una pérdida total de capacidad del modelo.
• Laguna en el conocimiento: Aunque métodos existentes (como SmoothQuant, QuaRot, AWQ) han demostrado recuperar parte de esta capacidad mediante escalado o rotación, no existe una comprensión clara de qué sitios de activación de entrada son los responsables dominantes del error ni si estos errores son accesibles mediante intervenciones en el "eje residual" (la corriente de residuo principal del modelo).

2. Metodología y Marco Teórico

El autor propone una descomposición operatoria de los cinco sitios de activación de entrada dentro de un bloque SwiGLU, clasificándolos en dos categorías basadas en la fuente de sus entradas:

A. Taxonomía: Lectores vs. Generadores

1. Lectores (Readers): Son las líneas que reciben la corriente residual normalizada (RMSNorm).
   • Componentes: qkv (atención), w1 y w3 (MLP).
   • Propiedad: Su magnitud está estrictamente acotada por cualquier intervención que controle la magnitud del residuo.
2. Generadores (Generators): Son las líneas que reciben cantidades generadas internamente dentro del bloque.
   • Componentes: o_proj (salida de atención) y w2 (producto bilineal de SwiGLU).
   • Propiedad:
     • o_proj: Su entrada es una mezcla lineal ponderada por softmax. El control residual es débil y depende de los pesos de atención.
     • w2: Su entrada es el producto elemento a elemento silu(w1) ⊙ w3. Teóricamente, incluso si los factores están acotados, el producto puede generar una "cola pesada" (outliers) debido a la dependencia positiva de eventos raros. El control residual no puede acotar este producto más allá de un límite trivial.

B. Intervención Propuesta: Depth Registers + Hinge Loss (DR+sink)

Para probar esta hipótesis, el autor introduce un método de entrenamiento que actúa exclusivamente en el eje residual sin usar calibración post-hoc ni rotaciones:

• Depth Registers (DR): Se divide el eje de canales residuales en una región semántica y una región de "registros" (k canales).
• Pérdida Hinge (Sink Loss): Se aplica una penalización de tipo hinge que fuerza la magnitud máxima de los canales de registro a estar por debajo de un umbral fijo ($\tau$).
• Objetivo: Los canales de registro actúan como un "sumidero" para absorber las magnitudes extremas (outliers), manteniendo los canales semánticos calibrados y libres de picos, sin alterar la capacidad del modelo en FP16.

3. Contribuciones Clave

1. Descomposición Operatoria: Identifica que el daño de W4A4 no es uniforme; los "lectores" son responsables de la mayor parte del error y son accesibles desde el eje residual, mientras que el "generador" w2 (producto bilineal) genera un residuo de error que el control residual no puede eliminar.
2. Método DR+sink: Un método de entrenamiento que reduce la perplejidad W4A4 de 1727 a 119 (una mejora de ~14x) manteniendo la capacidad en FP16 y cero-shot.
3. Análisis de "Skip-Ablation": Define una métrica de daño en exceso de NLL ($\Delta NLL$) para cuantificar cuánto contribuye cada sitio. Demuestra que al mantener los "lectores" en FP16, el error restante se estabiliza en un "residuo" plano de ~0.28 nats, dominado casi exclusivamente por w2.
4. Validación de la Hipótesis de Rotación: Demuestra empíricamente que la rotación ortogonal (usada en QuaRot) no puede suprimir la cola del producto bilineal de SwiGLU, confirmando la predicción teórica de que la rotación no es una solución universal para este tipo de no linealidades.

4. Resultados Experimentales

Los experimentos se realizaron en un modelo de 300M parámetros (20 capas, 1024 dimensiones) entrenado con 5B tokens de FineWeb-Edu.

• Rendimiento de Perplejidad (PPL):

  • FP16 (Referencia): 23.6
  • W4A4 Ingenuo: 1727
  • W4A4 + DR+sink: 119 (~14x mejora sobre el ingenuo).
  • W4A4 + DR+sink + SmoothQuant: 39.9 (Mejora adicional combinada).
  • W4A4 + Per-Linear QuaRot: ~25.4 (Casi idéntico a DR+sink, validando que el problema principal son los "lectores").
  • W4A4 + Full QuaRot (incluye rotación en w2): ~26.0 (No cierra la brecha restante, confirmando que la rotación no arregla el problema de w2).

• Análisis de Kurtosis (Colas de distribución):

  • Con DR+sink, la kurtosis de los "lectores" (ej. qkv) colapsa de ~35.9 a ~2.85 (casi Gaussiano).
  • La kurtosis de los "generadores" (especialmente w2) permanece alta (~1900), indicando que el método no afecta su distribución.

• Descomposición del Daño ($\Delta NLL$):

  • Al mantener solo los "lectores" en FP16, el modelo alcanza un piso de ~31 PPL en todos los checkpoints.
  • El daño restante (~0.28 nats) es dominado por w2 ($\Delta remove(w2) \approx 0.69-0.88$ nats), mientras que o_proj contribuye muy poco una vez que los lectores están limpios.

5. Significado y Conclusiones

• Diagnóstico Preciso: El trabajo demuestra que el error en la cuantización W4A4 de modelos SwiGLU tiene dos componentes distintos: uno dominado por los "lectores" (accesible mediante control de magnitud residual) y otro dominado por el producto bilineal de w2 (inaccesible mediante control residual o rotación ortogonal).
• Limitación de la Rotación: Se prueba empíricamente que las rotaciones ortogonales (como las de QuaRot) no pueden eliminar las colas pesadas generadas por productos bilineales, desafiando la idea de que la rotación es una solución general para todos los outliers.
• Implicaciones Futuras: Dado que el residuo de error de w2 es irreducible con métodos puramente residuales o de rotación, el artículo sugiere que las soluciones prácticas para W4A4 en SwiGLU podrían requerir:
  • Precisión mixta específica para la capa w2.
  • Cuantización consciente del entrenamiento (QAT).
  • Técnicas de recorte de productos (product clipping).

En resumen, el artículo no solo ofrece un método de entrenamiento (DR+sink) que mejora drásticamente la cuantización W4A4, sino que proporciona una explicación teórica y empírica de por qué ciertos métodos fallan y dónde reside el límite fundamental de la compresión en arquitecturas modernas de transformadores.