Resumen Técnico: Compresión de Tensores sin Pérdida como Síntesis de Programas (Brevis)

Planteamiento del Problema

La rápida proliferación de modelos de lenguaje de gran tamaño (LLMs) y otras arquitecturas de aprendizaje profundo ha provocado un crecimiento exponencial en el número y el tamaño de los checkpoints de los modelos. Almacenar, transferir y desplegar estos activos se ha vuelto cada vez más costoso. Si bien se utilizan ampliamente compresores de propósito general (por ejemplo, gzip, zstd), estos tratan los pesos de los tensores como secuencias de bytes genéricas, ignorando las propiedades estructurales específicas de los tensores (por ejemplo, tipos de datos, formas, distribuciones de punto flotante y relaciones entre elementos). Por el contrario, los compresores específicos para tensores existentes suelen depender de pipelines fijos y específicos de un formato o de transformaciones predefinidas que pueden no capturar la diversidad de patrones específicos de cada tensor, tales como valores repetidos, subsecuencias recurrentes y relaciones simples entre elementos.

El desafío central es lograr una compresión sin pérdida (bit-exacta) que se adapie a la estructura única de cada tensor sin depender de esquemas rígidos y predefinidos.

Metodología: Brevis

Los autores presentan Brevis, un sistema que formula la compresión de tensores sin pérdida como un problema de síntesis de programas. En lugar de seleccionar un codec fijo, Brevis sintetiza un programa compacto y autónomo para cada tensor que, al ejecutarse, reconstruye el tensor original bit a bit.

1. Lenguaje de Dominio Específico (DSL) Tipado

Brevis emplea un DSL tipado con operadores reversibles diseñados para capturar estructuras de tensores recurrentes:

• Lit: Un fallback universal que almacena flujos de palabras arbitrarias utilizando diversas codificaciones físicas (palabras puras, empaquetado de bits, Huffman, rANS) para seleccionar la representación más pequeña.
• Const: Almacena una sola palabra y su recuento de repeticiones.
• Concat: Une programas para regiones adyacentes.
• Repeat: Expande un programa hijo no vacío múltiples veces para capturar subsecuencias repetidas.
• Map: Aplica biyecciones que preservan el ancho (por ejemplo, XOR, suma modular, ZigZag, codificación Gray) para exponer relaciones que no son visibles como simples cadenas.
• Scan: Almacena una palabra inicial seguida de actualizaciones (XOR o suma modular), exponiendo correlaciones adyacentes.
• Merge: Combina hijos de igual longitud cuyas anchuras suman el ancho del padre, útil para dividir campos de punto flotante (signo, exponente, fracción) o planos de bits.

2. Síntesis Dirigida por Objetivos con Búsqueda A* Acotada

El proceso de síntesis busca un programa $P$ tal que $Exec(P) = Bits(X)$, minimizando el tamaño serializado $L(P)$.

• Búsqueda hacia atrás (Backward Search): En lugar de enumerar programas, Brevis busca hacia atrás desde el tensor objetivo. Cada "hueco" en el programa parcial se empareja con el flujo exacto que debe generar. Solo se aplica un operador si su descomposición del flujo objetivo es válida (es decir, si la composición de las salidas de los hijos reconstruye el objetivo exactamente).
• Prioridad de Producción Basada en el Checkpoint: Para guiar la búsqueda de manera eficiente, Brevis aprende una distribución de probabilidad sobre las producciones del DSL a partir de una pequeña muestra determinista de tensores dentro de un checkpoint. Este prior captura las tendencias estructurales específicas del modelo (por ejemplo, el uso frecuente de Repeat o Merge para capas específicas).
• Algoritmo A Acotado: La búsqueda utiliza un algoritmo A acotado donde los estados se priorizan mediante $g(s) + h(s)$.
  • $g(s)$ es el costo acumulado basado en el prior aprendido.
  • $h(s)$ es una heurística admisible para el costo de finalización restante.
  • Un límite inferior de bytes ($LB_L$) poda los estados que no pueden mejorar la mejor solución actual.
  • La búsqueda está restringida por un presupuesto de expansión (por ejemplo, 1 expansión por tensor en la evaluación principal) y límites de nodos/profundidad.

3. Formato de Archivo y Reconstrucción

El archivo resultante preserva el encabezado safetensors y el diseño original. Cada registro de tensor contiene su nombre, dtype, forma y el programa sintetizado. La descompresión consiste en parsear el programa, validar las reglas de tipo y ejecutarlo directamente sin búsqueda ni el uso del prior aprendido, lo que garantiza una reconstrucción bit-exacta.

Contribuciones Clave

1. Formulación: El artículo formula la compresión de checkpoints sin pérdida como síntesis de programas, representando cada tensor como un programa de DSL compacto y autónomo.
2. Diseño del DSL: Un lenguaje de dominio específico tipado con operadores reversibles que capturan diversas estructuras de tensores (valores, secuencias, campos y planos) garantizando la reconstrucción bit-exacta.
3. Algoritmo de Síntesis: Un algoritmo de síntesis A* acotado guiado por un prior de producción específico del checkpoint aprendido de una pequeña muestra, permitiendo el descubrimiento eficiente de programas compactos dentro de un presupuesto práctico.
4. Evaluación: Evaluación exhaustiva en 10 checkpoints públicos (lenguaje, audio, imagen) que totalizan 2.13 TB de datos.

Resultados

Los autores evaluaron Brevis en 10 checkpoints públicos que abarcan modelos de generación de lenguaje, audio e imagen (2.13 TB de datos de origen).

• Reducción de Almacenamiento: Brevis redujo el tamaño total de los datos de 2.13 TB a 1.41 TB, logrando una reducción de almacenamiento del 33.93%.
• Comparación con Compresores de Propósito General: Brevis produjo archivos más pequeños que zstd, gzip, LZ4 y Snappy en los 10 checkpoints. La reducción promedio sobre estas bases osciló entre el 12.94% (vs. zstd) y el 30.87% (vs. Snappy).
• Comparación con Compresores Específicos de Tensores:
  • Comparado con ZipNN, Brevis produjo archivos más pequeños en todos los 10 checkpoints, ahorrando 10.21 GB en agregado. El ahorro total es del 0.72% (con un ahorro medio por checkpoint de 0.71%).
  • Comparado con DFloat11 (en el checkpoint Llama-3.1-8B donde había un resultado natio validado disponible), Brevis fue un 2.90% más pequeño (316.27 MB ahorrados).
• Rendimiento (Throughput): Bajo una configuración de concurrencia práctica, Brevis alcanzó 3.60 GB/s de compresión y 6.61 GB/s de descompresión. Se sitúa en la frontera de Pareto, ofreciendo una mejor reducción de almacenamiento que los compresores generales manteniendo un rendimiento competitivo.
• Significancia Estadística: Las mejoras sobre todas las bases emparejadas se mantuvieron estadísticamente significativas tras la corrección de Holm ($p_H = 0.0098$).

Significancia y Reivindicaciones

El artículo afirma que, al tratar la compresión como síntesis de programas, Brevis supera las limitaciones de los pipelines fijos. Demuestra que un programa autónomo puede servir como una representación altamente eficiente y bit-exacta de un tensor, adaptándose a las redundancias estructurales específicas que los codecs genéricos pasan por alto.

Los autores enfatizan que su enfoque:

• Preserva cada byte de origen (bit-exacto), lo que lo hace adecuado para el archivo y la transferencia exacta de checkpoints, a diferencia de los métodos con pérdida (cuantización/poda).
• No requiere decodificadores específicos de hardware o familias de codecs predefinidas.
• Logra ahorros de almacenamiento significativos en diversos tipos de modelos y formatos de datos (BF16, FP32, FP8).

El documento reconoce modestamente sus limitaciones, señalando que la evaluación cubre checkpoints públicos en lugar de todas las cargas de trabajo de tensores, y que la implementación actual sintetiza los tensores de forma independiente (sin síntesis entre tensores). Los autores posicionan a Brevis como un paso hacia una compresión sin pérdida más adaptativa y consciente de la estructura para el creciente ecosistema de los grandes modelos.