Fragmentar tarefas complexas e inserir exemplos de referência nos prompts são práticas comuns entre usuários mais experientes de GenAI. Um estudo conduzido por pesquisadores do MIT revelou que uma técnica chamada treinamento em tempo de teste (TTT) pode melhorar de forma significativa o desempenho de modelos de linguagem (LLMs) em problemas desafiadores — especialmente aqueles que exigem raciocínio estruturado e aprendizado com poucos exemplos. O artigo com os resultados será apresentado em setembro, na Conferência Internacional sobre Aprendizado de Máquina.
Na abordagem tradicional, ao inserir um exemplo de referência, o prompt direciona a resposta para um conjunto de parâmetros (as correlações possíveis) mais restrito, o que resulta em conteúdo mais contextualizado. Com o TTT, novos parâmetros podem ser gerados.
Diferentemente do aprendizado no contexto (in-context learning), em que o modelo apenas recebe exemplos da tarefa como referência, o TTT atualiza temporariamente parâmetros internos do modelo durante a própria inferência. Esse ajuste rápido, feito com base em um pequeno conjunto de dados gerado a partir dos exemplos apresentados, permite que a IA se adapte melhor a tarefas para as quais não foi treinada originalmente.
Ganhos de desempenho
“O aprendizado genuíno – o que fizemos aqui com o treinamento em tempo de teste – é algo que esses modelos não podem fazer por conta própria depois de serem configurados”, afirma Ekin Akyürek, PhD e principal autor do estudo. “Mas mostramos que, se você forçar um pouco o modelo a realmente aprender, grandes melhorias no desempenho podem acontecer.”
Entre os resultados, o TTT aumentou em até seis vezes a precisão de um modelo de 8 bilhões de parâmetros no Abstraction and Reasoning Corpus (ARC) — um conjunto de quebra-cabeças visuais com regras abstratas — chegando a 53%. Combinado a técnicas de síntese de programas, a taxa subiu para 61,9%, próxima ao desempenho humano médio. No BIG-Bench Hard (BBH), que reúne 27 tarefas de linguagem natural voltadas ao raciocínio, a técnica superou o método padrão de few-shot prompting em 7,3 pontos percentuais, passando de 50,5% para 57,8%.
Segundo os autores, os ganhos são particularmente expressivos em desafios que envolvem regras formais ou mudanças inesperadas no padrão dos dados, como nas tarefas Dyck Languages e Ruin Names, com saltos de precisão entre 20 e 50 pontos percentuais. “Embora simplesmente fornecer exemplos possa aumentar modestamente a acurácia, atualizar o modelo com esses exemplos pode levar a um desempenho muito superior, principalmente em domínios desafiadores”, observa Mehul Damani, coautor do estudo.
A pesquisa detalha fatores críticos para o sucesso do TTT, como métodos de geração e aumento de dados (data augmentation), escolha da função de perda, número de passos de otimização e uso de adaptadores LoRA (Low-Rank Adaptation) específicos para cada tarefa.
O time planeja avançar para um modelo capaz de decidir, de forma autônoma, quando aplicar o treinamento em tempo de teste ou recorrer apenas ao aprendizado no contexto. “Não gostaríamos de fazer isso para todas as consultas do usuário, mas é útil se você tiver uma tarefa muito difícil que deseja que o modelo resolva bem”, afirma Akyürek.
O estudo é assinado também por Linlu Qiu, Han Guo, Jyothish Pari, Adam Zweiger, Yoon Kim e Jacob Andreas, todos vinculados ao MIT e ao CSAIL (Laboratório de Ciência da Computação e IA).
Formas de implementação
O artigo descreve o processo de implementação do TTT (Treinamento em Tempo de Teste) em detalhes. Ele discute as principais decisões de design e desafios na aplicação do TTT em LMs, incluindo:
- Geração de dados: O documento explica como construir um conjunto de treinamento em tempo de teste ($\mathcal{D}_{TTT}$) usando a abordagem de “leave-one-out” ou “direct I/O”. Também menciona a possibilidade de usar técnicas de aumento de dados (como rotações, reflexos e permutações de cores) para ampliar o material em tarefas com entradas estruturadas, como no ARC.
- Função de perda: São apresentadas três maneiras de calcular a perda no contexto de “leave-one-out”: apenas na saída de teste, em todas as saídas (incluindo as de demonstração) e em todas as entradas e saídas.
- Otimização e parametrização: O TTT é implementado com a otimização de adaptadores LoRA (Low-Rank Adaptation) por um pequeno número de passos de gradiente. O artigo compara a abordagem de adaptadores específicos por tarefa com a de um único adaptador compartilhado por várias tarefas.
- Inferência: Após o treinamento, o modelo atualizado é usado para fazer previsões com uma estratégia de inferência aumentada, que pode incluir um esquema de votação hierárquica para agregar as previsões.
O documento também fornece detalhes sobre os hiperparâmetros específicos usados para as experiências nos conjuntos de dados ARC e BBH.
