Consistency LLM: transformando LLMs em decodificadores paralelos para acelerar a inferência em até 3,5x

 (hao-ai-lab.github.io)
2 pontos por GN⁺ 2024-05-09 | 1 comentários | Compartilhar no WhatsApp
  • Ao gerar respostas longas, a decodificação autorregressiva (AR), que produz tokens um a um, é o principal gargalo de latência; o CLLM busca reduzir isso com decodificação paralela em blocos de n tokens
  • Consistency Large Language Models (CLLMs) ajustam um LLM pré-treinado para aprender a mapear rapidamente qualquer estado arbitrário de n tokens para o mesmo ponto fixo do resultado greedy AR
  • O Jacobi decoding teoricamente converge para o mesmo resultado da geração greedy AR, mas em LLMs convencionais mostrou ganho limitado de velocidade, em média de apenas 1,05x, o que reduzia sua utilidade prática
  • Em experimentos com Spider, Human-Eval, GSM8k e MT-bench, o CLLM mostrou aceleração de geração entre 2,4x e 3,4x, sendo avaliado em nível semelhante ou superior a técnicas de inferência rápida como Medusa2 e Eagle
  • Na inferência, não exige componentes auxiliares de modelo nem mudanças na arquitetura do modelo-alvo, o que permite buscar ganhos de velocidade junto com eficiência de memória

O gargalo de decodificação que o CLLM quer atacar

  • LLMs são usados em várias áreas, como programação, direito e aconselhamento em saúde, mas na etapa de inferência normalmente geram tokens um a um com decodificação autorregressiva (AR)
  • Quanto mais longa a resposta, maior a latência causada por essa geração token a token, aumentando o tempo de espera percebido pelo usuário
  • Métodos anteriores de inferência rápida frequentemente exigem mudanças de arquitetura, componentes auxiliares ou um modelo rascunho para gerar vários tokens de uma vez

Como funciona o Jacobi decoding e quais são seus limites

  • O Jacobi decoding parte das iterações de ponto fixo de Jacobi e Gauss-Seidel, e foi provado que, na decodificação greedy, produz o mesmo resultado da geração AR
  • Ele transforma a geração sequencial em um sistema com n variáveis e n equações não lineares, permitindo cálculo paralelo por iteração de Jacobi
  • O fluxo concreto é o seguinte
    • A partir do prompt de entrada, faz uma suposição aleatória dos próximos n tokens
    • Insere o prompt e a sequência de n tokens no LLM e atualiza isso iterativamente
    • Quando a sequência deixa de mudar, chega-se a um ponto fixo (fixed point)
    • A sequência final de n tokens converge para a saída que a decodificação AR teria produzido com estratégia greedy
  • O caminho entre a suposição aleatória inicial e o resultado final da geração AR é chamado de trajetória de Jacobi (Jacobi trajectory)
  • O Jacobi decoding básico mostrou, em LLMs reais, apenas cerca de 1,05x de ganho médio de velocidade em relação à decodificação AR
    • Como os LLMs foram treinados em modo AR, quando há erro nos tokens anteriores eles têm dificuldade para prever corretamente os tokens seguintes
    • Na maioria das iterações de Jacobi, apenas um token é corrigido na sequência de n tokens, o que alonga a trajetória
  • Lookahead decoding e speculative decoding tentam reduzir as ineficiências do Jacobi decoding e da decodificação AR tradicional, mas introduzem custo extra de memória na inferência

O objetivo de treinamento do Consistency LLM

  • O CLLM adapta um LLM pré-treinado para que ele se mova de forma consistente de qualquer ponto da trajetória de Jacobi até o ponto fixo
  • Esse objetivo é semelhante ao dos consistency models, uma técnica de aceleração para modelos de difusão
  • Usando trajetórias de Jacobi coletadas do modelo-alvo, o modelo é treinado com uma perda que incentiva convergência em uma única etapa durante a iteração de Jacobi
  • O treinamento para transformar cada modelo-alvo em CLLM é composto por duas partes
    • Preparação das trajetórias de Jacobi
      • Para cada prompt, o Jacobi decoding é executado sequencialmente em todos os segmentos de corte em blocos de n tokens
      • A sequência completa de resposta passa a ser a concatenação de pontos fixos consecutivos
      • Cada sequência gerada ao longo da trajetória é contabilizada como um item de dado de treinamento
      • O corte em n tokens é usado para evitar a queda de velocidade causada por avaliar entradas muito longas em respostas longas
    • Otimização conjunta de consistency loss e AR loss
      • A consistency loss induz a previsão de vários tokens de uma vez
      • A AR loss ajuda a preservar a qualidade de geração, evitando que o CLLM se afaste da distribuição do LLM-alvo

Composição da função de perda

  • O LLM-alvo é denotado por p, e o CLLM por qθ; qθ é inicializado com os parâmetros de p
  • A global consistency (GC) loss incentiva o CLLM a produzir o ponto fixo y* quando recebe como entrada um estado arbitrário y da trajetória de Jacobi
  • A local consistency (LC) loss ajusta estados adjacentes y(j) e y(j+1) na trajetória de Jacobi para que produzam a mesma saída
  • A distância entre distribuições D(·||·) segue as opções discutidas no método GKD, e nos experimentos foi usado principalmente forward KL
  • A AR loss aplica a perda autorregressiva tradicional com base na saída gerada l do LLM-alvo p, para impedir desvio da distribuição do modelo-alvo
  • A perda total de treinamento é composta pela consistency loss somada à AR loss ponderada por w

Configuração experimental e resultados

  • Os experimentos incluem três tarefas de domínio específico e um benchmark de diálogo de domínio aberto
    • Spider: text-to-SQL
    • Human-Eval: completude de código Python
    • GSM8k: matemática
    • MT-bench: diálogo de domínio aberto
  • Dependendo da tarefa, os modelos-alvo usados foram um coder LLM ajustado por fine-tuning, Deepseek-coder-7B-instruct, LLaMA-2-7B e ABEL-7B-001
  • Treinamento e avaliação foram realizados em servidores com NVIDIA A100 40GB
  • Nos domínios específicos, comparado ao modelo-alvo original, Medusa2, speculative decoding e outros baselines, o CLLM mostrou o maior ganho de velocidade
  • No MT-bench, um CLLM treinado a partir do LLaMA2-7B com o dataset ShareGPT alcançou ganho de velocidade semelhante ao Medusa2 quando combinado com lookahead decoding
    • A pontuação no MT-bench também ficou em nível comparável
    • O CLLM não exige mudanças na arquitetura original do modelo-alvo
    • Sem componentes auxiliares, oferece maior eficiência de memória

Custo de treinamento

  • O custo de fine-tuning do CLLM é apresentado como intermediário
  • O LLaMA-7B alcançou ganho de 3,4x no dataset Spider após processar apenas cerca de 1 milhão de tokens
  • Em datasets grandes, como o CodeSearchNet-Python, usar apenas 10% do dataset para gerar trajetórias de Jacobi rendeu cerca de 2,5x de aceleração
  • O número total de tokens é estimado pela seguinte fórmula
    • número médio de trajetórias por prompt × comprimento médio da trajetória × número de prompts
  • O custo estimado de treinamento por dataset é o seguinte
    • Spider: 2 milhões de tokens, < 0,01% do custo de pré-treinamento
    • CodeSearchNet-Python: 100 milhões de tokens, ~0,1% do custo de pré-treinamento
    • GSM8K: 10 milhões de tokens, ~0,01% do custo de pré-treinamento
    • ShareGPT: 200 milhões de tokens, ~0,2% do custo de pré-treinamento

Fast forwarding e stationary tokens

  • O LLM-alvo normalmente gera apenas um token correto por iteração de Jacobi
  • No CLLM, aparece o fenômeno de fast forwarding, em que vários tokens consecutivos são previstos corretamente em uma única iteração de Jacobi
  • No modelo-alvo, até tokens gerados corretamente antes podem ser trocados por tokens incorretos em iterações posteriores
  • O CLLM mostra capacidade de prever corretamente tokens com antecedência, mesmo quando há erro nos tokens anteriores, e de mantê-los nas iterações seguintes
    • Esses tokens são chamados de stationary tokens
  • Fast forwarding e stationary tokens fazem o Jacobi decoding do CLLM convergir mais rápido, contribuindo para o ganho de velocidade de geração

Aprendizado de padrões linguísticos

  • Observou-se que o CLLM aprende, durante o treinamento, o conceito linguístico de colocações (collocations)
  • Colocações são sequências de palavras ou termos que aparecem juntas com frequência maior do que o acaso sugeriria
  • Há exemplos tanto em linguagem natural quanto em código
    • Linguagem natural: talk to, remind … of …
    • Estrutura verbo + substantivo: make a decision, catch a cold
    • Estruturas gramaticais específicas de domínio: SELECT … FROM …, if … else
  • O objetivo de consistency generation faz com que o CLLM infira essas estruturas em qualquer ponto da trajetória de Jacobi, incentivando a prever várias palavras ao mesmo tempo para reduzir o número de iterações

Material e código

Leituras relacionadas

1 comentários

 
GN⁺ 2024-05-09
Opiniões no Hacker News

  • Parece com a experiência que tive quando fiz uma aula de "desenho livre" (sem aulas expositivas)
    Desde criança me diziam que eu desenhava bem, mas, na prática, lembro que eu repetia desenhos detalhados parecidos com os que já tinha feito antes, ou então passava muito tempo desenhando. Acho que, com tempo e paciência, qualquer pessoa consegue desenhar uma cena de forma bastante convincente
    Na aula não havia regras nem explicações; cada um levava o material que quisesse. Alguns levavam tinta, outros lápis; eu levei carvão, e a única coisa definida era o tempo entre as poses do modelo. As primeiras poses eram bem curtas, cerca de 1 minuto, e iam ficando mais longas até chegar a poses de 5 minutos; a qualquer momento você podia rasgar o desenho e desenhar a mesma pose de novo
    O aquecimento curto, na verdade, forçava a acertar proporções e contornos na primeira tentativa e, ao contrário da ideia comum de que a pressa estraga tudo, parece que, ao aprender ou refinar uma técnica, a urgência funciona como um fator de estresse que induz atenção e aprendizado
    Antes da aula eu provavelmente já conseguiria fazer desenhos de qualidade parecida, mas tenho certeza de que levaria de 5 a 10 vezes mais tempo. Funcionou ser forçado a não ficar rodeando o problema e sentir o custo dos erros apressados
    Ainda assim, acho uma pena chamar essa técnica de Consistency. Combina com melhoria de desempenho, mas parece um nome menos adequado para melhoria de velocidade de inferência; entendo que o sentido seja "consistência com o resultado que acabaria saindo se fosse gerado um token por vez". Eu preferiria chamá-la de Proficiency LLM, esperando a mesma saída, mas sem a inibição de chegar à mesma conclusão aos tropeços

    • Como autor do CLLM, obrigado por compartilhar sua experiência e seus insights. O processo de aprimorar a habilidade de desenho parece parecido com o processo de treinamento do CLLM, embora o fator de estresse no treinamento atual do CLLM não seja algo que vá ficando cada vez mais rígido
      No desenho, é possível definir um tempo permitido para cada tentativa e torná-lo cada vez mais curto. No CLLM, parece possível tornar o processo de treinamento mais difícil fazendo com que estados cada vez mais distantes na trajetória de Jacobi sejam mapeados para o estado final
      O termo "consistency" vem da similaridade entre o consistency model da geração de imagens por difusão e o consistency LLM, porque os processos de treinamento são parecidos
    • Tive uma experiência interessante em uma aula prática de zoologia de invertebrados em certo verão
      Os alunos entravam no laboratório, recebiam uma amostra, e a instrução era apenas: "desenhem isto em 30 minutos. Comecem"
      Não havia nada como "é assim que se desenha", "faça isto e não faça aquilo"; era, na prática, algo mais próximo de "não importa se você está ansioso ou se acha que não sabe desenhar. Não dê desculpas, apenas desenhe. Agora"
      Todos nós desenhamos e, ao longo do verão, com mais animais chegando e o mesmo treino se repetindo, todos melhoraram enormemente
      O que aquela aula ensinou foi que qualquer pessoa, realmente qualquer pessoa, consegue desenhar. A atitude coletiva mudou de "nem sei se isso é possível" para "claro que consigo. É fácil, comum, nada demais"
      É uma abordagem altamente recomendável, e foi uma das aulas mais libertadoras e surpreendentes que fiz na universidade
    • Sistemas, em geral, ficam mais eficientes quando submetidos a estresse. Ao mesmo tempo, também podem ser empurrados para ótimos locais; tudo tem seus prós e contras

  • Os autores dizem que a decodificação de Jacobi é equivalente à decodificação autorregressiva gulosa, mas, na prática, muitas vezes não queremos usar uma temperatura de amostragem maior que 0 para evitar repetições e respostas genéricas demais?
    Não conheço nada dessa estratégia de decodificação, então talvez eu esteja deixando passar uma forma simples de refletir isso

    • Boa pergunta. Estamos trabalhando ativamente para também dar suporte a outras estratégias de amostragem que não sejam amostragem gulosa
      No contexto do treinamento de CLLM, em vez de mapear para o objetivo de treinamento o ponto fixo estático obtido pela decodificação de Jacobi, estamos chamando isso de ponto fixo dinâmico. Para acompanhar os novos avanços, basta ficar de olho no repositório do GitHub
    • Concordo. É fácil verificar se um token foi a escolha de valor máximo, mas parece difícil verificar se um token apareceu com a probabilidade desejada
      Uma etapa de ajuste fino que treine a trajetória a se aproximar de uma conclusão de n tokens com as estatísticas desejadas ainda parece possível, mas não sei bem como substituir a etapa de verificação do ponto fixo. Talvez seja algo como "verificar se estava acima deste limiar fixo de verossimilhança"

  • Parece uma otimização bastante arriscada para fazer antes de realmente entendermos o que acontece dentro de um LLM. Por exemplo, quem acredita em uma interpretação geométrica provavelmente teria algo a dizer, e, nos casos em que se usam tokens de "preenchimento", isso também poderia ser prejudicial
    Além disso, a suposição de que "a pessoa forma uma frase completa na cabeça e depois fala palavra por palavra" não é um fato universal, é apenas uma suposição, e parece simplificar demais a atividade que ocorre na nossa mente. Será que realmente temos um plano completo antes de falar ou digitar? Como budista, vejo isso quase como uma ilusão. Indo além, e quanto aos pensamentos simultâneos? Pensamos linearmente em unidades de frases?
    De qualquer forma, a matemática é bem legal

    • Essa otimização não afeta os resultados do LLM e tem garantia de gerar um resultado equivalente ao da decodificação direta
      Não vamos tratar LLMs como algo mágico parecido com a nossa mente. São apenas mais um programa que gera frases que fazem sentido
    • Essa suposição pode ser útil neste contexto, mas parece bastante claro que ela não é verdadeira
      Quando se pede a alguém que explique um evento passado complexo em vários ramos, as pessoas rapidamente inserem fragmentos, complementos e digressões no meio das frases para cobrir toda a extensão do acontecimento. Acho que nunca vi a hipótese de granularidade no nível da frase em um contexto científico sério
    • Não posso dizer que isso valha para todo mundo, mas pelo menos eu não formo frases completas na cabeça antes de falar
      Às vezes fico preso em um beco sem saída gramatical no meio da frase e preciso concluir o pensamento com uma palavra ou expressão estranha, ou simplesmente paro e começo de novo do início
    • Embora palavras possam ser decompostas, em várias línguas, em unidades menores que carregam significado, não parecemos formar palavras sequencialmente a partir de seus componentes inferiores
      Também não parece haver um motivo claro para que esse fenômeno se quebre de repente no nível da frase
    • Fiquei curioso para saber o que é a interpretação geométrica

  • É surpreendente que isso não esteja recebendo mais atenção. Parece um ganho claro em desempenho de inferência
    O custo desse ajuste fino é razoável, cerca de 0,01% do custo do pré-treinamento original. Os ganhos de desempenho também parecem bastante consistentes

    • Parece um resultado muito grande para o desempenho de LLMs
      Não conheço bem outro artigo que tenha sugerido que é possível aumentar tanto o desempenho de inferência de LLMs. Houve algum antes?
      Ainda mais sob as condições de pelo menos manter a qualidade da saída, melhorar não só a latência das consultas como também a vazão total, não exigir computação adicional, e ter uma implementação relativamente prática sem adicionar grande complexidade
      Como isso é construído sobre trabalhos em decodificação paralela/Jacobi, o insight em si pode ser visto como incremental. Os resultados anteriores também eram necessários e importantes, mas talvez este seja o resultado que extraiu valor real do mundo real a partir do potencial da decodificação paralela
    • Ganhos de inferência semelhantes ou maiores já são obtidos com a decodificação especulativa, que é amplamente usada
      Então este trabalho é realmente interessante e, até onde sei, já houve tentativas anteriores menos bem-sucedidas, mas ainda não está claro qual será seu impacto real
    • Obrigado pelo interesse no nosso trabalho. Mesmo treinando apenas uma parte do conjunto de dados com consistency loss + AR loss, obtivemos uma aceleração considerável, e o custo ficou em torno de 0,01% do pré-treinamento
      Treinar com mais dados aumenta ainda mais a velocidade. Isso porque o modelo consegue aprender com colocações e frases que aparecem com mais frequência
      Para mais detalhes, veja o artigo; também é possível observar que, conforme o tamanho dos dados de treinamento aumenta, o ganho de velocidade acaba saturando

  • No começo achei que fosse um artigo do tipo Medusa, usando mais unembed heads para adivinhar tokens subsequentes, mas não era nada disso
    É realmente impressionante. Não usa parâmetros adicionais; apenas adiciona uma perda auxiliar de treinamento

    • A única coisa em comum entre Medusa e CLLM é que ambos treinam/adaptam LLMs para inferência rápida
      As técnicas de treinamento e de decodificação são completamente diferentes e, como você apontou, CLLM não precisa de parâmetros adicionais nem de configuração de máscaras de atenção para verificação baseada em árvore

  • Acho que em breve vamos perceber que não é necessariamente preciso treinar modelos
    O que é necessário é uma boa indexação e amostragem
    Em essência, em certo nível, todo LLM pode ser visto como algo parecido com um banco de dados do dataset, com uma ótima interface em linguagem natural por cima
    Ambos são apenas maneiras diferentes de navegar pelos dados armazenados

    • LLMs conseguem criar facilmente dados que não existem no conjunto de dados de treinamento
      LLMs não navegam por dados armazenados. LLMs não são um banco de dados dos dados de treinamento
    • Mas indexação também é treinamento. Só não usa descida de gradiente de ponta a ponta
    • Como o modelo é várias ordens de magnitude menor que uma versão comprimida dos dados de treinamento, ele não pode ser equivalente a esse banco de dados
    • Então talvez você goste do artigo Infinigram. Ele foi discutido recentemente
      https://news.ycombinator.com/item?id=40266791

  • Existe algum lugar onde uma pessoa leiga como eu possa "perguntar a um especialista em IA"?
    Por exemplo, eu queria perguntar por que um LLM não responde de forma determinística ao receber o mesmo prompt
    Quero aprender isso, e talvez eu devesse acompanhar vídeos do tipo "construindo um LLM em 1 hora" no YouTube

    • O software literalmente contém um gerador de números aleatórios que escolhe um entre os candidatos a próximo token ponderados que o modelo produz
      O processo de escolha pode ter vários controles para ajustar a resposta. Se você quiser torná-lo determinístico e tiver acesso direto ao software, dependendo do software usado, definir top-k = 1 ou temperature = 0.0 o torna determinístico
      Normalmente a configuração padrão não é determinística, porque, quando você a torna totalmente determinística, a qualidade dos resultados tende a não ser muito boa
    • Para essa resposta, veja os vídeos do 3blue1brown
      Um modelo de LLM produz um vetor de probabilidades sobre tokens, e o usuário do LLM usa aleatoriedade para escolher um token da lista dos mais prováveis
    • Porque LLMs são, em essência, matrizes de probabilidade
      Quando você insere um prompt, ele calcula a probabilidade da próxima palavra e repete esse processo até formar uma frase. As probabilidades aprendidas se baseiam nos dados de treinamento
      Por causa desse modelo probabilístico de base, ele não é 100% determinístico. Além disso, modelos como o ChatGPT introduzem intencionalmente um parâmetro temperature para adicionar aleatoriedade a todo o processo
      Se quiser ler mais, esta resposta se baseou no seguinte artigo: The Matrix: A Bayesian learning model for LLMs, https://arxiv.org/abs/2402.03175
    • Na maioria dos sistemas, isso pode ser controlado por um parâmetro de configuração de inferência chamado temperature
      Mas, ao definir a temperatura no valor mais baixo possível, a qualidade das respostas tende a ficar muito ruim. O sistema acaba preso em algum ótimo local e continua se repetindo. Esse tipo de resposta pode ser "determinística", mas não é boa
    • Este texto foi um bom ponto de partida: explica de forma bastante sistemática sem perder a visão geral
      https://writings.stephenwolfram.com/2023/02/what-is-chatgpt-...

  • Não existe almoço grátis, então, a meu ver, há também alguma perda de caminhos aqui
    Por exemplo, algumas trajetórias de Jacobi, por definição, excluem caminhos de temperatura mais alta. Do ponto de vista de recuperação de dados, isso pode até ser positivo, mas, se você quiser maximizar a criatividade, pode ser negativo

    • Existem algoritmos melhores e algoritmos piores
      Não sei se "não existe almoço grátis" sempre se aplica de forma especialmente significativa. Algumas coisas não estão na fronteira de Pareto

  • Fiquei curioso sobre a explicação detalhada da parte que diz que “a decodificação especulativa tem custo adicional de memória durante a inferência”
    Na decodificação especulativa, um modelo menor gera “ramificações” rápidas, mas que podem ser imprecisas, e depois um modelo maior valida essas ramificações. Mas a decodificação especulativa só precisa da memória correspondente a um único token, e os tokens das outras ramificações são simplesmente mascarados durante a inferência. Se o tamanho do contexto for 1000 e houver cerca de 30 ramificações de 5 tokens, o overhead de memória é de 3%, algo desprezível. Se o tamanho do contexto for muito menor que o número de ramificações, fico pensando se um usuário de LLM generativo com uma janela de contexto de apenas 50 tokens vai mesmo se importar com a velocidade de geração
    Além disso, técnicas de decodificação especulativa não se limitam à amostragem gulosa. Elas devem se comportar exatamente como o modelo original e amostrar com a probabilidade esperada. A maior parte da literatura sobre decodificação especulativa já relata ganhos de velocidade de 2,6~3,5 vezes. Este post do blog relata uma velocidade de geração 2,4~3,4 vezes maior, então não sei se isso é uma atualização tão grande assim
    Acima foi mencionada a decodificação especulativa, e as técnicas que o autor comparou parecem ser Medusa2 e Eagle, mas o problema central continua o mesmo. Seja qual for o método usado para prever tokens antecipadamente, há certos pontos em que o token anterior é absolutamente necessário antes de prever o próximo. Isso não é uma questão do modelo ou da técnica, mas do que é matematicamente possível. Se a distribuição de probabilidade do quinto token seguinte depende fortemente dos quatro tokens anteriores, como prever 5 tokens de uma vez? Vale tanto para decodificação especulativa quanto para decodificação de Jacobi ou decodificação paralela de múltiplos tokens
    Se esse método só oferece suporte à amostragem gulosa, fico em dúvida sobre qual é a vantagem. Ainda mais considerando que outras técnicas já estão alcançando os ganhos de velocidade esperados. Comparar aceleração em amostragem gulosa com aceleração em amostragem aleatória é como comparar maçãs com laranjas, e desconfio que, mesmo depois de adaptar esse método para amostragem aleatória, o mesmo ganho de velocidade não vá permanecer por causa do problema central mencionado acima

    • Talvez a parte de que “o token anterior é absolutamente necessário antes de prever o próximo token” seja justamente a contribuição central deste paper
      Pode ser que, por meio do treinamento de consistência, eles tenham mostrado que um LLM consegue prever os próximos n tokens mesmo quando há suposições incorretas nos tokens anteriores
      Por outro lado, matematicamente é verdade que p(x_t|x_1,...,x_t-1) depende de tudo de x_1 até x_t-1, mas, na prática, para prever x_t talvez só sejam necessários x_1 até x_t-2, e a atenção a x_t-1 pode ser mínima. Portanto, pode ser possível prever x_t usando x_1 até x_t-2 e um x_t-1 impreciso
    • A decodificação especulativa precisa carregar um modelo menor na memória e executar inferência com esse modelo

  • É um conteúdo interessante. É uma ideia que muita gente provavelmente já pensou, mas o texto e a apresentação estavam bem organizados

    • Sim. Meu colega de quarto e eu conversamos sobre isso há um ano. Dá para fazer algo parecido também com direcionamento de LLMs