DiBA: Diagonal and Binary Matrix Approximation for Neural Network Weight Compression

Resumen Técnico: DiBA (Aproximación de Matriz Diagonal y Binaria)

Enunciado del Problema
Las matrices densas constituyen una unidad de almacenamiento y computación dominante en las redes neuronales modernas, apareciendo en capas lineales, incrustaciones de tokens, proyecciones de atención y convoluciones $1 \times 1$. Si bien se han desarrollado técnicas como el podado, la cuantización y la descomposición de bajo rango para comprimir estos pesos, enfrentan compensaciones específicas. La cuantización, aunque compatible con el hardware, introduce problemas de optimización intrínsecamente discretos difíciles de resolver, que a menudo requieren formulaciones complejas de enteros mixtos o estimadores directos. Las factorizaciones de bajo rango a menudo carecen de las restricciones estructurales específicas necesarias para un almacenamiento y aritmética eficientes a nivel de bits. El artículo aborda la necesidad de una representación matricial compacta que equilibre la reducción teórica del almacenamiento, la precisión de la aproximación y la viabilidad de la adaptación a tareas posteriores sin reoptimizar componentes discretos.

Metodología
El artículo propone DiBA (Aproximación de Matriz Diagonal y Binaria), una factorización matricial estructurada que aproxima una matriz objetivo densa $A \in \mathbb{R}^{m \times n}$ como:
$$\hat{A} = D_1 B_1 D_2 B_2 D_3$$
donde:

• $D_1, D_2, D_3$ son matrices diagonales de valores reales de tamaños $m \times m$, $k \times k$ y $n \times n$, respectivamente.
• $B_1 \in \{0, 1\}^{m \times k}$ y $B_2 \in \{0, 1\}^{k \times n}$ son matrices binarias.
• $k$ es una dimensión intermedia que controla la compensación entre almacenamiento y precisión.

Propiedades Teóricas:

• Almacenamiento: Los factores binarios se almacenan a 1 bit por entrada, mientras que los factores diagonales utilizan precisión de punto flotante estándar ($Q$ bits). La relación de almacenamiento teórica respecto a una matriz densa de $Q$ bits es $\rho_{\text{DiBA}} = \frac{k(m+n) + Q(m+k+n)}{Qmn}$.
• Computación: Los productos matriz-vector se descomponen en tres operaciones de escalado elemento a elemento (diagonales) y dos operaciones de mezcla binaria (selección y suma), reduciendo las multiplicaciones de punto flotante de $mn$a$m + k + n$.
• Estructura: La aproximación es una superposición de $k$ patrones binarios de rango uno escalados, corregidos por escalas diagonales de filas y columnas.

Optimización: DiBA-Greedy
Para ajustar los factores DiBA a una matriz de pesos preentrenada, los autores introducen DiBA-Greedy, un solver alternante:

1. Actualizaciones Diagonales: Con los factores binarios fijos, las matrices diagonales se actualizan mediante soluciones de mínimos cuadrados de forma cerrada.
2. Actualizaciones Binarias: Con los factores diagonales fijos, las matrices binarias se actualizan evaluando el cambio exacto en el error cuadrático de Frobenius al invertir un solo bit. Una inversión se acepta solo si la mejora supera una tolerancia $\tau$.
   Este enfoque evita la no convexidad de la optimización simultánea separando las actualizaciones continuas y discretas, aunque es un solver local codicioso y no garantiza un óptimo global.

Adaptación Posterior: DiBARD
Reconociendo que minimizar el error de reconstrucción no garantiza el rendimiento en tareas posteriores, el artículo propone DiBARD (DiBA con ajuste fino solo de factores diagonales). En este protocolo:

• Las matrices binarias $B_1$ y $B_2$ se congelan después del ajuste inicial DiBA-Greedy.
• Solo las entradas diagonales de $D_1, D_2, D_3$ se reajustan en datos de entrenamiento posteriores utilizando optimización estándar basada en gradientes.
• Esto transforma la fase de adaptación en un problema de optimización continua pequeño, evitando la necesidad de búsqueda discreta durante el ajuste fino.

Contribuciones Clave

1. Formulación: La definición de DiBA como una factorización diagonal y binaria para la compresión de pesos densos, incluyendo su contabilidad teórica de almacenamiento.
2. Algoritmo: El desarrollo de DiBA-Greedy, que combina actualizaciones de mínimos cuadrados de forma cerrada con actualizaciones codiciosas exactas de un bit.
3. Evaluación Empírica (Nivel de Matriz): Una evaluación en 40 matrices de pesos densas extraídas de modelos preentrenados públicos (NLP y Visión). Los resultados muestran mejoras consistentes en la Relación Señal-Ruido (SNR) a medida que aumenta la relación de almacenamiento teórica, con matrices de incrustación que muestran una SNR particularmente alta.
4. Protocolo de Adaptación: La introducción de DiBARD y su validación en dos estudios de reemplazo de componentes:
   • DistilBERT/WikiText: Reemplazo de la proyección de incrustación/salida atada. DiBARD mejoró la precisión de tokens enmascarados de 0.4447 (DiBA pre-ajuste) a 0.5210, superando a Int4-scale-RT (0.5110) y acercándose al modelo original (0.4906) a pesar de una relación de almacenamiento teórica más baja.
   • AST/Comandos de Voz: Reemplazo de 48 matrices de proyección de atención en un Transformador de Espectrograma de Audio. DiBARD recuperó la precisión de prueba de 0.7684 a 0.9781 (el original fue 0.9837), recuperando el 97.4% de la brecha de rendimiento inducida por la compresión inicial.

Significado y Afirmaciones
El artículo afirma que DiBA ofrece una representación de pesos compacta que combina efectivamente la mezcla binaria fija con un reajuste diagonal ligero. El significado principal radica en demostrar que:

• Los pesos densos pueden aproximarse intercalando escalado diagonal y mezcla binaria compartida sin cuantización elemento a elemento.
• El rendimiento posterior puede recuperarse después de una compresión agresiva reajustando solo los factores diagonales, evitando así la dificultad computacional de reoptimizar variables binarias discretas dentro de una pérdida de red no lineal.

Limitaciones y Modestia
Los autores señalan explícitamente varias limitaciones:

• Almacenamiento Teórico: Las relaciones de almacenamiento son teóricas, asumiendo matrices binarias empaquetadas en bits; no se proporcionan aceleraciones medidas de memoria o kernels empaquetados en bits.
• Alcance: Los experimentos 2-1 y 2-2 reportan relaciones de compresión relativas solo a los componentes reemplazados, no al modelo completo.
• Solver: DiBA-Greedy es un solver local codicioso, no un optimizador global.
• Evaluación: Las evaluaciones posteriores utilizan divisiones de datos fijas y semillas limitadas.
  El artículo no afirma aceleraciones de inferencia medidas ni reducciones de tiempo de ejecución de extremo a extremo, centrándose en cambio en la calidad de la aproximación, el almacenamiento teórico y la viabilidad de la recuperación solo con diagonales.