O Google reduziu o consumo de memória de seus modelos de IA em seis vezes sem perder precisão – graças ao algoritmo TurboQuant.

O Google Research publicou um artigo sobre o TurboQuant, um algoritmo de quantização que reduz o cache de chave-valor de grandes modelos de linguagem para 3 bits sem perda de precisão ou necessidade de treinamento adicional. Em testes com aceleradores de IA Nvidia H100, o algoritmo TurboQuant de 4 bits (com um quarto bit para correção de erros) melhorou o desempenho do cálculo de logits de atenção em oito vezes em comparação com chaves de 32 bits não quantizadas; o tamanho do cache de chave-valor foi reduzido em seis vezes.

Fonte da imagem: and machines / unsplash.com

O cache KV armazena dados previamente calculados pelo mecanismo de atenção para evitar que o modelo precise recalculá-los a cada etapa de geração de tokens. À medida que a janela de contexto aumenta, esse cache cresce significativamente; os métodos de quantização tradicionais ajudam a reduzir o tamanho do cache, mas ao custo de vários bits extras por valor devido às constantes de quantização armazenadas junto com os dados comprimidos — uma espécie de sobrecarga semelhante a um dicionário em algoritmos tradicionais de arquivamento de dados ZIP e RAR. Mesmo com janelas de contexto grandes, essa sobrecarga ainda é significativa.

O TurboQuant elimina essa sobrecarga em um processo de duas etapas, eliminando completamente a necessidade de um dicionário. A primeira etapa envolve o método PolarQuant, que converte vetores de dados de coordenadas cartesianas para coordenadas polares. A cada vetor é atribuído um raio (comprimento) e um ângulo (direção). As distribuições angulares são previsíveis e concentradas, portanto, o PolarQuant elimina a etapa de normalização de cada bloco, que consome muitos recursos e é inevitável com quantizadores tradicionais. O resultado é uma compressão de alta qualidade sem a sobrecarga de armazenar constantes de atenção — dicionários.

Fonte da imagem: Conny Schneider / unsplash.com

O segundo estágio aplica uma camada de correção de erros de 1 bit usando o algoritmo Johnson-Lindenstrauss quantizado. O erro de quantização residual é projetado em um espaço de menor dimensão e cada valor é reduzido a um único bit de sinal, eliminando assim o viés sistemático nos cálculos da pontuação de atenção com sobrecarga insignificante.

O Google testou os dois componentes do algoritmo TurboQuant individualmente e em conjunto nos benchmarks de contexto longo LongBench, Needle In A Haystack, ZeroSCROLLS, RULER e L-Eval nos modelos de código aberto Gemma e Mistral. O TurboQuant demonstrou resultados perfeitos na classe de problemas “agulha no palheiro”, alcançando uma compressão de cache KV de pelo menos seis vezes. No conjunto de testes LongBench, que inclui resposta a perguntas, geração de código e criação de resumos, o TurboQuant teve um desempenho igual ou superior ao do algoritmo de compressão KIVI em todas as tarefas.

O TurboQuant demonstrou eficácia não apenas no processamento de grandes modelos de linguagem, mas também na busca vetorial — ele foi testado contra os conhecidos algoritmos de compressão Product Quantization e RabbiQ no conjunto de dados GloVe. Mesmo sem treinamento ou otimização, o algoritmo desenvolvido pelo Google apresentou resultados superiores aos de seus concorrentes, que foram especificamente ajustados para esse conjunto de dados. O TurboQuant não requer nenhum treinamento ou ajuste fino, e sua execução tem sobrecarga insignificante — ele está pronto para ser implementado mesmo sob alta carga.