PACE comprime o estado de LLMs lineares para 4 bits e economiza um terço da memória

Uma equipe de pesquisadores propôs o PACE, um quantizador pós-treinamento que armazena o estado recorrente de modelos com atenção linear em quatro bits e não exige ajuste fino. Em quatro modelos híbridos, de 4B a 35B parâmetros, a variante com cabeças mistas de 4 e 8 bits fica no mesmo nível do DAMP de 8 bits em três dos modelos maiores, ocupando cerca de dois terços da memória dele. A versão com todas as cabeças em 4 bits pontua mais que o Q-Mamba com o mesmo volume de memória.

Por que o estado vira gargalo

Em modelos com atenção linear, o contexto não fica em um KV-cache que cresce com a sequência, mas em um estado recorrente de tamanho fixo. À primeira vista, isso parece conveniente. Mas o sistema de inferência precisa manter uma cópia separada desse estado para cada sequência que decodifica ao mesmo tempo. Quanto mais cópias cabem na memória do acelerador, mais requisições são atendidas em paralelo e, portanto, mais barato fica cada token.

Daí o interesse em quantizar o próprio estado. Os métodos existentes arredondam o estado depois de cada token gerado. O erro de arredondamento se acumula passo a passo e, em quatro bits, a precisão cai de forma perceptível. Os autores de trabalhos anteriores chegavam à mesma conclusão: ou se aumenta a largura em bits, ou se retreina o modelo para a quantização. O PACE mostra que existe um terceiro caminho.

As duas observações por trás do PACE

A primeira diz respeito à frequência do arredondamento. O PACE arredonda o estado apenas a cada poucos tokens e mantém os tokens intermediários em um pequeno buffer. Algo parecido já acontece nos kernels chunked na etapa de prefill: eles atualizam o estado uma vez por chunk, e não a cada token. Menos arredondamentos significam menos erro acumulado.

A segunda observação diz respeito à estrutura do estado nas camadas gated-delta. Os valores grandes não se espalham de maneira uniforme, e sim se concentram em colunas. O motivo é que alguns canais da atualização gravada a cada token continuam grandes ao longo de muitos tokens. Por isso, o PACE usa uma escala de quantização compartilhada por coluna inteira. Assim, uma única entrada grande deixa de prejudicar a precisão de todos os outros valores, que seriam obrigados a dividir a escala com ela.

O terceiro ponto é a precisão mista. Algumas cabeças retêm o erro de arredondamento por mais tempo que outras. O PACE mantém justamente essas em oito bits e converte as demais para quatro.

• Método: pós-treinamento, sem ajuste fino do modelo
• Formato-alvo: quatro bits por elemento do estado, oito para as cabeças de erro mais "duradouro"
• Arredondamento: a cada poucos tokens, com os tokens intermediários guardados em um buffer
• Escala de quantização: compartilhada por coluna do estado
• Validação: quatro modelos híbridos de 4B a 35B parâmetros

O que as medições mostraram

O PACE misto, com cabeças de 4 e 8 bits, ficou no nível do DAMP de 8 bits em três dos quatro modelos, os maiores. A memória usada foi de cerca de dois terços da do DAMP. A variante totalmente em 4 bits supera o Q-Mamba na comparação com o mesmo volume de memória. Os autores também observam que arredondar a cada poucos tokens preserva a aceleração da decodificação proporcionada pelo buffer: a economia de memória não é paga com perda de velocidade.

Esse último ponto é importante na prática. Uma quantização de estado que adiciona operações de arredondamento a cada passo pode consumir o ganho obtido com o batch mais denso. Aqui, o arredondamento sai do laço crítico, e esse risco diminui.

Onde isso pode ser útil na infraestrutura

Arquiteturas híbridas, que alternam camadas lineares com atenção convencional, aparecem cada vez mais em modelos abertos e comerciais. Para quem as opera, o PACE oferece uma alavanca clara: reduzir em cerca de um terço a memória do estado em relação ao esquema de 8 bits, sem retreinamento nem ajuste fino. O espaço liberado pode ser usado para um batch maior ou para uma fila mais longa de sessões simultâneas.

Há limitações. O resultado foi confirmado em quatro modelos híbridos e, no menor deles, o PACE misto não alcançou o nível do DAMP. Equipes que planejam adotar o método devem verificar como ele se comporta no próprio modelo e na própria carga antes de reorganizar a produção em torno de um estado de quatro bits. Ainda assim, a ideia em si, arredondar com menos frequência e ajustar a escala à estrutura real dos dados, parece transferível e barata de implementar.