Consistency LLM: convertendo LLMs em decodificadores paralelos para melhorar a velocidade de inferência em 3,5x

Consistency Large Language Models (CLLM)

• LLMs têm sido tradicionalmente vistos como decodificadores sequenciais que funcionam decodificando um token por vez de forma contínua
• Este blog mostra que LLMs pré-treinados podem ser facilmente convertidos em decodificadores paralelos eficientes
• Consistency Large Language Models (CLLM) é uma nova família de decodificadores paralelos capaz de decodificar com eficiência sequências de n tokens em cada etapa da inferência, reduzindo a latência
• O CLLM imita o processo cognitivo humano de formar uma frase completa na mente antes de pronunciar palavra por palavra, e pode ser treinado de forma eficaz apenas com o fine-tuning de um LLM pré-treinado
• O CLLM é treinado realizando decodificação paralela de modo a mapear uma sequência de n tokens inicializada aleatoriamente para o mesmo resultado da decodificação autoregressiva (autoregressive, AR) em o menor número possível de etapas
• Segundo os resultados experimentais, o CLLM apresentou ganho de velocidade de geração entre 2,4x e 3,4x em comparação com decodificadores AR tradicionais, com desempenho comparável ou superior a técnicas de inferência rápida como Medusa2 e Eagle
• O CLLM consegue atingir esse ganho de desempenho sem custo adicional de memória

Contexto e limitações da decodificação Jacobi

• Como LLMs geram tokens um a um no modo de decodificação AR, respostas longas acabam tendo alta latência
• A decodificação Jacobi tem origem nos métodos Jacobi e Gauss-Seidel para resolver equações não lineares, e foi demonstrado que ela é equivalente à geração AR com greedy sampling
• A decodificação Jacobi reconstrói o processo de geração sequencial como um sistema de n equações não lineares com n variáveis baseado em iteração de Jacobi, permitindo processamento paralelo
• Em cada etapa de iteração, ela pode prever um ou mais tokens corretos, o que potencialmente acelera a decodificação AR
• Porém, na prática, como LLMs treinados com AR quase não conseguem gerar o token correto quando há erro nos tokens anteriores, a maioria das iterações Jacobi produz apenas uma correção na sequência de n tokens, gerando trajetórias (trajectories) Jacobi mais longas
• Técnicas como lookahead decoding ou speculative decoding tentam aliviar essa ineficiência da decodificação Jacobi, mas incorrem em custo extra de memória na inferência, enquanto o CLLM não

Método de treinamento do CLLM

• O treinamento do CLLM é composto principalmente por duas partes: preparação das trajetórias Jacobi e otimização conjunta das perdas de consistency e AR
• Na etapa de preparação das trajetórias Jacobi, a decodificação Jacobi é executada sequencialmente em blocos de n tokens até que toda a sequência de resposta de l tokens seja gerada, e a sequência produzida em cada trajetória é tratada como um item de dados
• Durante o treinamento, são otimizadas em conjunto a perda de consistency e a perda AR: a perda de consistency garante a previsão de vários tokens de uma vez, enquanto a perda AR evita que o CLLM se afaste do LLM de referência, preservando a qualidade da geração
• A perda de global consistency (GC) minimiza a distância entre um ponto arbitrário da trajetória Jacobi e o ponto fixo, incentivando o CLLM a prever o ponto fixo a partir de qualquer ponto da trajetória
• A perda de local consistency (LC) induz estados adjacentes da trajetória Jacobi a produzirem a mesma saída
• A perda AR inclui a perda AR tradicional com base nos resultados gerados pelo LLM de referência, com o objetivo de evitar que o CLLM se desvie da distribuição do modelo-alvo

Resultados experimentais

• Os experimentos incluíram tarefas de domínios especializados como Spider (text-to-SQL), Human-Eval (completação de código Python) e GSM8k (matemática), além de desafios amplos de diálogo aberto como o MT-bench
• O CLLM apresentou o maior ganho de velocidade em relação ao modelo-alvo e, sem custo extra na inferência, obteve aceleração semelhante ou melhor que a do Medusa2
• No MT-bench, o CLLM alcançou praticamente o mesmo ganho de velocidade que quando combinado com Medusa2, mas oferecendo maior adaptabilidade e eficiência de memória
• O custo de fine-tuning do CLLM é moderado e, mesmo quando o conjunto de dados é grande, foi possível obter cerca de 2,5x de ganho de velocidade usando apenas cerca de 10% do dataset para gerar trajetórias Jacobi
• O CLLM consegue prever corretamente vários tokens consecutivos em uma única iteração Jacobi por meio do fenômeno de fast forwarding
• O CLLM também demonstrou capacidade de tokens estacionários (stationary tokens), prevendo com antecedência tokens corretos e mantendo-os inalterados mesmo quando há erro nos tokens anteriores
• Por meio do treinamento, o CLLM aprende conceitos linguísticos centrais como collocation, o que lhe permite inferir estrutura em qualquer ponto da trajetória Jacobi e prever várias palavras simultaneamente para minimizar o número de iterações

Opinião do GN⁺

• O CLLM parece resolver de forma eficaz o problema de alta latência do método de decodificação AR dos LLMs existentes ao usar decodificação Jacobi. Em especial, impressiona o fato de alcançar ganho de velocidade com decodificação paralelizada sem custo adicional de memória

• O método de treinamento do CLLM parece relativamente simples, pois consiste em fazer fine-tuning de LLMs existentes com perda de consistency, mas ainda assim é significativo por ter aprendido collocation, uma das propriedades importantes da linguagem, e com isso melhorado bastante o desempenho da decodificação paralela

• Ainda assim, como o CLLM assume greedy sampling, parece ser necessário mais estudo para verificar se ele também funciona bem com estratégias de decoding mais diversas. Além disso, como os experimentos atuais se limitam ao inglês, também será preciso validar a possibilidade de generalização para vários idiomas

• O CLLM parece ser uma abordagem prática para acelerar o tempo de resposta de LLMs. Deve poder ser aplicado bem a tarefas que exigem resposta em tempo real, como busca na web e chatbots

• Pessoalmente, há expectativa de que o método de aprendizado por consistency do CLLM possa ser aplicado não apenas a LLMs como GPT, mas também a outros modelos generativos, como modelos de geração de imagem ou de síntese de voz. Seria ótimo se a ideia do CLLM viesse a contribuir para melhorar a eficiência de diversos modelos generativos no futuro

[GN⁺]

Comentários do Hacker News

• Em aulas de desenho livre, praticar desenhos repetidos com limite de tempo melhora bastante a habilidade. Treinar desenho rápido acaba desenvolvendo melhor a capacidade de acertar proporções e contornos.
• Jacobi decoding é igual ao greedy autoregressive decoding, mas na prática é preciso definir a temperatura de amostragem acima de 0 para evitar repetições e respostas excessivamente genéricas.
• Em breve vamos perceber que não é necessário treinar modelos. Só é preciso um bom indexamento e amostragem. LLMs são basicamente equivalentes a um banco de dados de datasets com uma excelente interface de NLP.
• O ganho de desempenho de inferência do artigo sobre CLLM é notável. O custo de fine-tuning é razoável (cerca de 0,01% do custo de pré-treinamento) e a melhora de desempenho também é bem consistente.
• Como as trajetórias de Jacobi excluem caminhos de alta temperatura, isso pode ser positivo para recuperação de dados, mas negativo para maximizar a criatividade.
• Seria bom ter um recurso de "perguntar a um especialista em IA" para explicar por que LLMs não respondem de forma determinística ao mesmo prompt.
• Fico curioso se o ganho de velocidade da Groq(https://groq.com/) vem desse método ou se eles usam outra abordagem.