Os grandes modelos de linguagem (LLMs, na sigla em inglês) de raciocínio são projetados para resolver problemas complexos, decompondo-os em uma série de etapas menores. Esses modelos poderosos são particularmente eficazes em tarefas desafiadoras, como programação avançada e planejamento em várias etapas.

Mas o desenvolvimento de modelos de raciocínio exige uma quantidade enorme de computação e energia devido às ineficiências no processo de treinamento. Enquanto alguns dos processadores de alta potência trabalham continuamente em consultas complexas, outros no grupo permanecem ociosos.

Pesquisadores do MIT e de outras instituições descobriram uma maneira de usar esse tempo ocioso computacional para acelerar com eficiência o treinamento de modelos de raciocínio.

O novo método treina automaticamente um modelo menor e mais rápido para prever as saídas do modelo de raciocínio LLM maior, que é então verificado pelo modelo maior. Isso reduz a quantidade de trabalho que o modelo de raciocínio precisa realizar, acelerando o processo de treinamento.

A chave deste sistema reside na sua capacidade de treinar e implementar o modelo menor de forma adaptativa, entrando em ação apenas quando alguns processadores estão ociosos. Ao aproveitar recursos computacionais que de outra forma seriam desperdiçados, ele acelera o treinamento sem incorrer em sobrecarga adicional.

Ao ser testado em múltiplos modelos de lógica de aprendizagem (LLMs) de raciocínio, o método dobrou a velocidade de treinamento, mantendo a precisão. Isso pode reduzir o custo e aumentar a eficiência energética do desenvolvimento de LLMs avançados para aplicações como previsão de tendências financeiras ou detecção de riscos em redes elétricas.

“As pessoas querem modelos que consigam lidar com tarefas mais complexas. Mas se esse é o objetivo do desenvolvimento de modelos, então precisamos priorizar a eficiência. Encontramos uma solução sem perdas para esse problema e, em seguida, desenvolvemos um sistema completo que pode proporcionar ganhos de velocidade bastante expressivos na prática”, afirma Qinghao Hu, pós-doutorando do MIT e coautor principal de um artigo sobre essa técnica .

O artigo é acompanhado pelo coautor principal Shang Yang, estudante de pós-graduação em Engenharia Elétrica e Ciência da Computação (EECS); Junxian Guo, também estudante de pós-graduação em EECS; pelo autor sênior Song Han, professor associado em EECS, membro do Laboratório de Pesquisa em Eletrônica e cientista de destaque da NVIDIA; bem como por outros pesquisadores da NVIDIA, ETH Zurich, MIT-IBM Watson AI Lab e Universidade de Massachusetts em Amherst. A pesquisa será apresentada na Conferência Internacional da ACM sobre Suporte Arquitetural para Linguagens de Programação e Sistemas Operacionais.

Os desenvolvedores desejam que os LLMs (Learning Learning Machines) raciocinem para identificar e corrigir erros em seu processo de pensamento crítico. Essa capacidade permite que eles resolvam com maestria consultas complexas que um LLM padrão não conseguiria solucionar.

Para ensinar essa habilidade, os desenvolvedores treinam modelos de aprendizado de máquina (LLMs) de raciocínio usando uma técnica chamada aprendizado por reforço (RL). O modelo gera múltiplas respostas potenciais para uma consulta, recebe uma recompensa pela melhor resposta e é atualizado com base na melhor resposta. Essas etapas se repetem milhares de vezes à medida que o modelo aprende.

Mas os pesquisadores descobriram que o processo de geração de múltiplas respostas, chamado de rollout, pode consumir até 85% do tempo de execução necessário para o treinamento de RL.

“A atualização do modelo — que é a parte de 'treinamento' propriamente dita — consome muito pouco tempo em comparação”, diz Hu.

Esse gargalo ocorre nos algoritmos de RL padrão porque todos os processadores no grupo de treinamento precisam terminar suas respostas antes de poderem passar para a próxima etapa. Como alguns processadores podem estar trabalhando em respostas muito longas, outros que geraram respostas mais curtas ficam aguardando a conclusão das respostas anteriores.

“Nosso objetivo era transformar esse tempo ocioso em aumento de velocidade sem custos desnecessários”, acrescenta Hu.

Eles buscaram usar uma técnica já existente, chamada decodificação especulativa, para acelerar o processo. A decodificação especulativa envolve o treinamento de um modelo menor, chamado de rascunho, para adivinhar rapidamente as saídas futuras do modelo maior.

O modelo maior verifica as suposições do desenhista, e as respostas que ele aceita são usadas para treinamento.

Como o modelo maior consegue verificar todas as suposições do projetista de uma só vez, em vez de gerar cada resultado sequencialmente, ele acelera o processo.

Mas na decodificação especulativa, o modelo do desenhista é normalmente treinado apenas uma vez e permanece estático. Isso torna a técnica inviável para aprendizado por reforço, já que o modelo de raciocínio é atualizado milhares de vezes durante o treinamento.

Um desenhista estático se tornaria obsoleto e inútil rapidamente após algumas etapas.

Para superar esse problema, os pesquisadores criaram um sistema flexível conhecido como "Domando a Cauda Longa", ou TLT.

A primeira parte do TLT é um treinador de desenho adaptativo, que usa o tempo livre em processadores ociosos para treinar o modelo de desenho em tempo real, mantendo-o bem alinhado com o modelo alvo sem usar recursos computacionais extras.

O segundo componente, um mecanismo de implementação adaptativa, gerencia a decodificação especulativa para selecionar automaticamente a estratégia ideal para cada novo lote de entradas. Esse mecanismo altera a configuração da decodificação especulativa com base nas características da carga de trabalho de treinamento, como o número de entradas processadas pelo modelo preliminar e o número de entradas aceitas pelo modelo alvo durante a verificação.

Além disso, os pesquisadores projetaram o modelo preliminar para ser leve, permitindo um treinamento rápido. O TLT reutiliza alguns componentes do processo de treinamento do modelo de raciocínio para treinar o desenhista, resultando em ganhos adicionais de aceleração.

“Assim que alguns processadores terminam suas consultas curtas e ficam ociosos, nós os direcionamos imediatamente para o treinamento de modelos preliminares usando os mesmos dados que estão sendo usados no processo de implementação. O mecanismo principal é nossa decodificação especulativa adaptativa — esses ganhos não seriam possíveis sem ela”, diz Hu.

Eles testaram o TLT em vários modelos de lógica de aprendizagem (LLMs) de raciocínio treinados com conjuntos de dados do mundo real. O sistema acelerou o treinamento entre 70 e 210 por cento, preservando a precisão de cada modelo.

Como um bônus adicional, o modelo de desenhista compacto poderia ser facilmente utilizado para uma implantação eficiente como um subproduto gratuito.

No futuro, os pesquisadores pretendem integrar a TLT em mais tipos de estruturas de treinamento e inferência e encontrar novas aplicações de aprendizado por reforço que possam ser aceleradas usando essa abordagem.

Este trabalho foi financiado pelo MIT-IBM Watson AI Lab, pelo MIT AI Hardware Program, pelo MIT Amazon Science Hub, pela Hyundai Motor Company e pela National Science Foundation.