MIT cria método que dobra velocidade de treinamento de inteligência artificial

Pesquisadores do Massachusetts Institute of Technology (MIT), nos EUA, e de outras instituições apresentaram uma técnica que promete reduzir significativamente o custo computacional e energético no treinamento de modelos de linguagem de grande porte voltados a raciocínio complexo. Em testes, o método chegou a dobrar a velocidade de treinamento sem perda de precisão.

Os chamados reasoning large language models (LLMs) são projetados para resolver problemas complexos ao dividi-los em etapas menores, o que os torna mais eficazes em tarefas como programação avançada e planejamento com múltiplos passos. No entanto, esse desempenho exige grande volume de processamento, especialmente durante o treinamento com aprendizado por reforço (reinforcement learning, RL).

Nesse processo, o modelo gera múltiplas respostas para uma mesma consulta, recebe uma recompensa pela melhor alternativa e é atualizado com base nesse resultado. O ciclo se repete milhares de vezes. Segundo os pesquisadores, apenas a etapa de geração dessas múltiplas respostas — conhecida como rollout — pode consumir até 85% do tempo total de execução do treinamento. Já a atualização efetiva dos parâmetros do modelo ocupa parcela reduzida desse tempo.

O gargalo ocorre porque, nos algoritmos tradicionais de RL, todos os processadores envolvidos precisam concluir suas respostas antes que o sistema avance para a próxima rodada. Quando algumas unidades lidam com respostas longas e complexas, outras que finalizaram tarefas mais curtas permanecem ociosas.

Para aproveitar esse tempo de inatividade, a equipe desenvolveu um sistema chamado Taming the Long Tail (TLT). A proposta combina uma versão adaptativa da técnica de speculative decoding, que utiliza um modelo menor — denominado “drafter” — para prever rapidamente as saídas do modelo maior. O modelo principal verifica essas previsões e aceita aquelas que considera corretas, acelerando o processo ao validar vários resultados de uma vez, em vez de gerá-los sequencialmente.

Diferentemente das abordagens convencionais, nas quais o modelo menor é treinado apenas uma vez e permanece estático, o TLT atualiza o “drafter” de forma contínua e adaptativa. Sempre que processadores ficam ociosos, eles passam automaticamente a treinar esse modelo auxiliar com os mesmos dados utilizados no rollout. Assim, o sistema mantém o alinhamento entre os dois modelos sem demandar recursos computacionais adicionais.

Além do treinador adaptativo, o TLT inclui um mecanismo que ajusta dinamicamente a estratégia de speculative decoding conforme as características de cada lote de treinamento, como o número de previsões feitas pelo modelo auxiliar e quantas são aceitas pelo modelo principal. O modelo menor também foi projetado para ser leve, permitindo treinamento rápido e reaproveitamento de componentes do próprio processo do modelo maior.

Em experimentos com diferentes LLMs de raciocínio treinados com bases de dados do mundo real, o sistema acelerou o treinamento entre 70% e 210%, preservando a acurácia. Como efeito adicional, o modelo auxiliar pode ser aproveitado posteriormente para implantação mais eficiente das aplicações.

O estudo será apresentado na ACM International Conference on Architectural Support for Programming Languages and Operating Systems e reúne autores do MIT, da NVIDIA, da ETH Zurich, do MIT-IBM Watson AI Lab e da University of Massachusetts Amherst.

Segundo Qinghao Hu, pesquisador de pós-doutorado no MIT e um dos autores principais do trabalho, a busca por modelos capazes de lidar com tarefas cada vez mais complexas torna a eficiência um fator central no desenvolvimento de inteligência artificial. O projeto também contou com financiamento da National Science Foundation, além de programas de pesquisa e parcerias industriais ligadas ao MIT.