Resumo Técnico: Compressão de Tensores Sem Perdas como Síntese de Programas (Brevis)

Declaração do Problema

A rápida proliferação de modelos de linguagem de grande escala (LLMs) e outras arquiteturas de aprendizado profundo levou a um crescimento exponencial no número e no tamanho dos checkpoints de modelos. Armazenar, transferir e implantar esses ativos tornou-se cada vez mais dispendioso. Embora compressores de uso geral (ex: gzip, zstd) sejam amplamente utilizados, eles tratam os pesos dos tensores como sequências de bytes genéricas, ignorando as propriedades estruturais específicas dos tensores (ex: tipos de dados, formas, layouts de ponto flutuante e relações entre elementos). Por outro outro lado, os compressores específicos para tensores existentes frequentemente dependem de pipelines fixos e específicos de formato ou transformações predefinidas que podem falhar em capturar diversos padrões específicos de tensores, tais como valores repetidos, subsequências recorrentes e relações simples entre elementos.

O desafio central é alcançar uma compressão sem perdas (bit-a-bit exata) que se adapte à estrutura única de cada tensor sem depender de esquemas rígidos e predefinidos.

Metodologia: Brevis

Os autores apresentam o Brevis, um sistema que formula a compressão de tensores sem perdas como um problema de síntese de programas. Em vez de selecionar um codec fixo, o Brevis sintetiza um programa compacto e autossuficiente para cada tensor que, quando executado, reconstrói o tensor original bit a bit.

1. Linguagem de Domínio Específico (DSL) Tipada

O Brevis utiliza uma DSL tipada com operadores reversíveis projetados para capturar estruturas recorrentes de tensores:

• Lit: Um fallback universal que armazena fluxos de palavras arbitrários usando várias codificações físicas (palavras brutas, empacotamento de bits, Huffman, rANS) para selecionar a representação de menor tamanho.
• Const: Armazena uma única palavra e sua contagem de repetição.
• Concat: Une programas para regiões adjacentes.
• Repeat: Expande um programa filho não vazio múltiplas vezes para capturar subsequências repetidas.
• Map: Aplica bijeções que preservam a largura (ex: XOR, adição modular, ZigZag, codificação Gray) para expor relações que não são visíveis como strings simples.
• Scan: Armazena uma palavra inicial seguida por atualizações (XOR ou adição modular), expondo correlações adjacentes.
• Merge: Combina filhos de igual comprimento cujas larguras somam a largura do pai, útil para dividir campos de ponto flutuante (sinal, expoente, fração) ou planos de bits.

2. Síntese Direcionada ao Alvo com Busca A* Limitada

O processo de síntese busca por um programa $P$ tal que $Exec(P) = Bits(X)$, minimizando o tamanho serializado $L(P)$.

• Busca Reversa (Backward Search): Em vez de enumerar programas, o Brevis busca de trás para frente a partir do tensor alvo. Cada "espaço vazio" (hole) no programa parcial é pareado com o fluxo exato que ele deve gerar. Um operador só é aplicado se sua decomposição do fluxo alvo for válida (ou seja, a composição das saídas dos filhos reconstrói o alvo exatamente).
• Prioridade de Produção Específica do Checkpoint: Para guiar a busca de forma eficiente, o Brevis aprende uma distribuição de probabilidade sobre as produções da DSL a partir de uma pequena amostra determinística de tensores dentro de um checkpoint. Este prior captura as tendências estruturais específicas do modelo (ex: uso frequente de Repeat ou Merge para camadas específicas).
• Algoritmo A Limitado: A busca utiliza um algoritmo A limitado onde os estados são priorizados por $g(s) + h(s)$.
  • $g(s)$ é o custo acumulado baseado no prior aprendido.
  • $h(s)$ é um heurístico admissível para o custo de conclusão restante.
  • Um limite inferior de bytes ($LB_L$) descarta estados que não podem melhorar a solução atual de melhor desempenho.
  • A busca é restringida por um orçamento de expansão (ex: 1 expansão por tensor na avaliação principal) e limites de nós/profundidade.

3. Formato de Arquivo e Reconstrução

O arquivo resultante preserva o cabeçalho e o layout originais do safetensors. Cada registro de tensor contém seu nome, dtype, forma e o programa sintetizado. A descompressão envolve o parsing do programa, a validação das regras de tipo e a execução direta, sem busca ou necessidade do prior aprendido, garantindo a reconstrução bit-a-bit exata.

Principais Contribuições

1. Formulação: O artigo formula a compressão de checkpoints sem perdas como síntese de programas, representando cada tensor como um programa de DSL compacto e autossuficiente.
2. Design da DSL: Uma linguagem de domínio específico tipada com operadores reversíveis que capturam diversas estruturas de tensores (valores, sequências, campos e planos) enquanto garantem a reconstrução bit-a-bit exata.
3. Algoritmo de Síntese: Um algoritmo de síntese A* limitado, guiado por um prior de produção específico do checkpoint aprendido a partir de uma pequena amostra, permitindo a descoberta eficiente de programas compactos dentro de um orçamento prático.
4. Avaliação: Avaliação abrangente em 10 checkpoints públicos (linguagem, áudio, imagem) totalizando 2,13 TB de dados.

Resultados

Os autores avaliaram o Brevis em 10 checkpoints públicos abrangendo modelos de geração de linguagem, áudio e imagem (2,13 TB de dados de origem).

• Redução de Armazenamento: O Brevis reduziu o tamanho total dos dados de 2,13 TB para 1,41 TB, alcançando uma redução de armazenamento de 33,93%.
• Comparação com Compressores de Uso Geral: O Brevis produziu arquivos menores que zstd, gzip, LZ4 e Snappy em todos os 10 checkpoints. A redução média sobre essas bases variou de 12,94% (vs. zstd) a 30,87% (vs. Snappy).
• Comparação com Compressores Específicos para Tensores:
  • Comparado ao ZipNN, o Brevis produziu arquivos menores em todos os 10 checkpoints, economizando 10,21 GB no agregado. A economia combinada é de 0,72% (com uma média de economia por checkpoint de 0,71%).
  • Comparado ao DFloat11 (no checkpoint Llama-3.1-8B, onde um resultado nato validado estava disponível), o Brevis foi 2,90% menor (316,27 MB economizados).
• Throughput (Vazão): Sob uma configuração de concorrência prática, o Brevis alcançou 3,60 GB/s de compressão e 6,61 GB/s de descompressão. Ele está na fronteira de Pareto, oferecendo melhor redução de armazenamento que compressores genéricos enquanto mantém um throughput competitivo.
• Significância Estatística: As melhorias sobre todas as bases comparadas permaneceram estatisticamente significativas após a correção de Holm ($p_H = 0,0098$).

Significância e Alegações

O artigo afirma que, ao tratar a compressão como síntese de programas, o Brevis supera as limitações de pipelines fixos. Ele demonstra que um programa autossuficiente pode servir como uma representação altamente eficiente e bit-a-bit exata de um tensor, adaptando-se a redundâncias estruturais específicas que codecs genéricos perdem.

Os autores enfatizam que sua abordagem:

• Preserva cada byte de origem (bit-a-bit exato), tornando-a adequada para arquivamento e transferência exata de checkpoints, ao contrário de métodos com perda (quantização/poda).
• Não requer decodificadores específicos de hardware ou famílias de codecs predefinidas.
• Alcança economias de armazenamento significativas através de diversos tipos de dados e formatos de modelos (BF16, FP32, FP8).

O artigo reconhece modestamente as limitações, observando que a avaliação cobre checkpoints públicos em vez de todas as cargas de trabalho de tensores, e que a implementação atual sintetiza tensores de forma independente (sem síntese entre diferentes tensores). Os autores posicionam o Brevis como um passo em direção a uma compressão sem perdas mais adaptável e consciente da estrutura para o crescente ecossistema de grandes modelos.