Resumen Técnico: Compresión de Pesos Compensada por Movimiento (MCWC)

Declaración del Problema

El despliegue de redes neuronales está cada vez más limitado no por la capacidad de cómputo, sino por el almacenamiento, el ancho de banda y la sobrecarga de inicio en frío requeridos para cargar los puntos de control en los aceleradores. Las pipelines de compresión existentes, dominadas por la cuantización y el recorte posteriores al entrenamiento, suelen tratar cada capa como un objeto independiente. Este enfoque pasa por alto la redundancia sustancial entre capas inducida por simetrías que preservan la función (por ejemplo, la capacidad de reindexar unidades ocultas o cabezas de atención sin cambiar la función de entrada-salida del modelo). En consecuencia, estos métodos no logran aprovechar la fuerte similitud y predictibilidad entre capas inherentes a arquitecturas modernas como los Transformers.

Metodología

MCWC propone un códec solo de pesos que trata la profundidad del modelo como una señal secuencial, análoga a los fotogramas de video. La premisa central es que, al eliminar el "movimiento" inducido por simetrías de permutación, la evolución restante de los pesos a través de las capas se vuelve altamente predecible, lo que permite una codificación residual eficiente. La pipeline consta de cuatro etapas principales:

1. Compensación de Movimiento mediante Alineación Funcional

El codificador identifica permutaciones que preservan la función para bloques de pesos (por ejemplo, unidades ocultas de FFN, cabezas de atención) que maximizan la correspondencia entre capas.

• Mecanismo: Para cada capa $\ell$, el método selecciona una permutación $\Pi_\ell$ que alinea los bloques con sus contrapartes en la capa anterior $\ell-1$. Esto se logra resolviendo un problema de asignación lineal para maximizar la similitud (una combinación de similitud del coseno basada en pesos y similitud basada en estadísticas de activación) entre los bloques candidatos y la referencia previamente alineada.
• Resultado: Esto transforma la secuencia de pesos en un "sistema de coordenadas alineado" donde las capas adyacentes son geométricamente comparables, reduciendo significativamente la energía de las diferencias entre ellas.

2. Codificación Predictiva Secuencial por Capas

Una vez alineados, los pesos se tratan como una secuencia temporal.

• Predicción: Un predictor ligero y secuencial por capas ($g_\theta$) estima los bloques alineados de la capa actual basándose en el contexto previamente decodificado (bloques alineados de la capa $\ell-1$), condicionado por incrustaciones de índice de capa y tipo de bloque.
• Residuos: El códec codifica únicamente los residuos de predicción cuantizados ($R = \text{ObjetivoAlineado} - \text{Predicción}$).
• Fotogramas Clave: Para evitar la acumulación de errores (deriva) en secuencias largas, se codifican periódicamente "fotogramas clave" como valores absolutos en lugar de residuos. Estos actúan como puntos de sincronización, permitiendo al decodificador reiniciar su contexto.

3. Cuantización Aprendida y Modelado de Entropía

• Cuantización: Los residuos se mapean a índices de código discretos utilizando un cuantizador escalar aprendido con tamaños de paso y medias por canal.
• Codificación de Entropía: Un modelo de entropía aprendido asigna probabilidades a los símbolos cuantizados, habilitando la codificación aritmética. El modelo se condiciona sobre el índice de capa, el tipo de bloque y las estadísticas locales derivadas de la salida del predictor.
• Entrenamiento: Todo el sistema (predictor, cuantizador, modelo de entropía) se entrena de extremo a extremo bajo un objetivo de tasa-distorsión ($L = L_{task} + \lambda E[-\log p_\psi(c)]$), donde $L_{task}$ es una pérdida de destilación que asegura que los pesos reconstruidos mantengan el rendimiento en la tarea.

4. Información Lateral y Decodificación

El flujo de bits incluye los residuos/fotogramas clave codificados por entropía, la información lateral de permutación (codificada mediante códigos de Lehmer codificados en delta para reducir la sobrecarga) y los parámetros del cuantizador. El decodificador invierte el proceso: decodifica por entropía, descuantiza, reconstruye los residuos mediante predicción y aplica permutaciones inversas para restaurar la estructura original del tensor para su despliegue.

Contribuciones Clave

1. Compensación de Movimiento en el Espacio de Pesos: La introducción de la selección de permutaciones por bloques para aumentar la correspondencia entre capas manteniendo la función representada, análogo a la estimación de movimiento en la compresión de video.
2. Códec Predictivo para Profundidad: Una arquitectura novedosa que trata la profundidad como un eje temporal, utilizando fotogramas clave explícitos y codificación residual para convertir las capas alineadas en "fotogramas" predecibles.
3. Entrenamiento de Tasa-Distorsión de Extremo a Extremo: Entrenamiento conjunto de modelos de cuantización y entropía para asegurar que el flujo de residuos discretos sea preciso después de la decodificación y de baja entropía bajo el modelo de probabilidad aprendido.
4. Validación Empírica: Demostración de fronteras de Pareto de tasa-precisión mejoradas tanto en modelos de lenguaje Transformer (Pythia) como en clasificadores de visión (ViT), con estudios de ablación que confirman que la alineación, la predicción y el modelado de entropía son todos necesarios para obtener las ganancias completas.

Resultados

MCWC fue evaluado en Pythia-1.4B (modelado de lenguaje) y ViT-B/16 (clasificación de visión) frente a bases sólidas que incluyen cuantización uniforme INT8/INT4, SmoothQuant, GPTQ, AWQ y cuantización escalar aprendida.

• Modelado de Lenguaje: En tasas de compresión agresivas (por ejemplo, 2.8 bits/parámetro), MCWC logra una perplejidad de 16.40, superando a Recorte+Cuantización (18.10) y a la cuantización escalar aprendida (17.30). En el régimen de tasa media (4.2 bits/parámetro), logra una perplejidad de 15.95, superando a INT4 uniforme (16.85) y manteniéndose competitivo con métodos conscientes del Hessiano como GPTQ (16.40) y AWQ (16.25).
• Clasificación de Visión: A 2.8 bits/parámetro, MCWC retiene un 80.15% de precisión top-1 en ImageNet-1k, superando a Recorte+Cuantización (77.60%) y a la cuantización escalar aprendida (78.70%). A 4.2 bits/parámetro, alcanza una precisión del 80.92%, superando a INT4 uniforme (79.35%).
• Eficiencia de Decodificación: Los tiempos de decodificación permanecen competitivos con las bases de cuantización estándar (por ejemplo, ~2.5s para Pythia-1.4B en CPU), ya que la reconstrucción predictiva y la permutación inversa son operaciones de tiempo lineal, siendo el costo dominante la decodificación de entropía y la materialización del tensor.
• Diagnósticos: Los estudios de ablación confirman que la alineación funcional aumenta la similitud del coseno entre capas adyacentes (de ~0.28 a ~0.67 para Pythia-1.4B) y reduce la energía residual normalizada, validando que la alineación expone un sistema de coordenadas más predecible.

Significado y Afirmaciones

El artículo posiciona a MCWC no como un reemplazo de los kernels de inferencia de baja precisión (por ejemplo, multiplicación de matrices INT8/INT4), sino como un códec especializado para almacenamiento y transporte de puntos de control. Su valor principal radica en reducir la huella de almacenamiento, el ancho de banda de red y la latencia de inicio en frío asociadas con la distribución y la materialización repetida de pesos de modelos grandes.

Los autores afirman que, al aprovechar la redundancia entre capas mediante compensación de movimiento, MCWC aborda un cuello de botella ortogonal a la cuantización tradicional. El método es más efectivo en regímenes de despliegue donde los puntos de control se mueven, almacenan o replican muchas veces (por ejemplo, inicios en frío sin servidor, registros de modelos, distribución en el borde), permitiendo que el costo de codificación fuera de línea único se amortice sobre los ahorros repetidos de despliegue. El trabajo sugiere que tratar los pesos neuronales como trayectorias estructuradas a través de la profundidad, en lugar de tensores aislados, ofrece una nueva vía para la compresión eficiente de modelos.