TAO: Usando o cálculo em tempo de teste para treinar LLMs eficientes sem dados rotulados

Grandes modelos de linguagem são desafiadores para adaptar a novas tarefas empresariais. O uso de prompts é propenso a erros e alcança ganhos de qualidade limitados, enquanto o ajuste fino requer grandes quantidades de dados humanos rotulados que não estão disponíveis para a maioria das tarefas empresariais. Hoje, estamos introduzindo um novo método de ajuste de modelo que requer apenas dados de uso não rotulados, permitindo que as empresas melhorem a qualidade e o custo para IA usando apenas os dados que já possuem. Nosso método, Otimização Adaptativa em Tempo de Teste (TAO), aproveita o cálculo em tempo de teste (popularizado por o1 e R1) e o aprendizado por reforço (RL) para ensinar um modelo a realizar uma tarefa melhor com base apenas em exemplos de entrada anteriores, o que significa que ele escala com um orçamento de cálculo de ajuste ajustável, não com o esforço de rotulagem humana. Crucialmente, embora o TAO use cálculo em tempo de teste, ele o usa como parte do processo para treinar um modelo; esse modelo então executa a tarefa diretamente com baixos custos de inferência (ou seja, não requer cálculo adicional no momento da inferência). Surpreendentemente, mesmo sem dados rotulados, o TAO pode alcançar melhor qualidade de modelo do que o ajuste fino tradicional, e pode trazer modelos de código aberto baratos como o Llama para dentro da qualidade de modelos proprietários caros como o GPT-4o e o3-mini.

TAO faz parte do programa de nossa equipe de pesquisa sobre Inteligência de Dados - o problema de fazer a IA se destacar em domínios específicos usando os dados que as empresas já possuem. Com TAO, alcançamos três resultados empolgantes:

• Em tarefas empresariais especializadas, como resposta a perguntas em documentos e geração de SQL, o TAO supera o ajuste fino tradicional em milhares de exemplos rotulados. Ele traz modelos de código aberto eficientes como Llama 8B e 70B para uma qualidade semelhante a modelos caros como GPT-4o e o3-mini¹ sem a necessidade de rótulos.
• Também podemos usar o TAO multi-tarefa para melhorar um LLM de maneira ampla em muitas tarefas. Sem usar rótulos, o TAO melhora o desempenho do Llama 3.3 70B em 2,4% em um amplo benchmark empresarial.
• Aumentar o orçamento de cálculo do TAO no momento do ajuste resulta em melhor qualidade do modelo com os mesmos dados, enquanto os custos de inferência do modelo ajustado permanecem os mesmos.

A Figura 1 mostra como o TAO melhora os modelos Llama em três tarefas empresariais: FinanceBench, DB Enterprise Arena e BIRD-SQL (usando o dialeto SQL do Databricks)². Apesar de ter acesso apenas a entradas LLM, TAO supera o ajuste fino tradicional (FT) com milhares de exemplos rotulados e traz Llama para a mesma faixa que modelos proprietários caros.

Figura 1: TAO em Llama 3.1 8B e Llama 3.3 70B em três benchmarks empresariais. TAO leva a melhorias substanciais na qualidade, superando o ajuste fino e desafiando LLMs proprietários caros.

TAO agora está disponível em pré-visualização para clientes Databricks que desejam ajustar Llama, e estará alimentando vários produtos futuros. Preencha este formulário para expressar seu interesse em testá-lo em suas tarefas como parte da pré-visualização privada. Neste post, descrevemos mais sobre como o TAO funciona e nossos resultados com ele.

Como o TAO funciona? Usando Test-Time Compute e Reinforcement Learning para Afinar Modelos

Em vez de exigir dados de saída anotados por humanos, a ideia chave no TAO é usar o cálculo em tempo de teste para que um modelo explore respostas plausíveis para uma tarefa, e então use aprendizado por reforço para atualizar um LLM com base na avaliação dessas respostas. Este pipeline pode ser escalado usando cálculo em tempo de teste, em vez de esforço humano caro, para aumentar a qualidade. Além disso, ele pode ser facilmente personalizado usando insights específicos da tarefa (por exemplo, regras personalizadas). Surpreendentemente, a aplicação desta escala com modelos de código aberto de alta qualidade leva a melhores resultados do que rótulos humanos em muitos casos.

Especificamente, TAO compreende quatro etapas:

• Geração de Respostas: Esta etapa começa com a coleta de exemplos de prompts de entrada ou consultas para uma tarefa. No Databricks, esses prompts podem ser coletados automaticamente de qualquer aplicação de IA usando nosso AI Gateway. Cada prompt é então usado para gerar um conjunto diversificado de respostas candidatas. Um rico espectro de estratégias de geração pode ser aplicado aqui, variando de prompts simples de cadeia de pensamento a técnicas sofisticadas de raciocínio e prompts estruturados.
• Avaliação de Respostas: Nesta etapa, as respostas geradas são sistematicamente avaliadas. As metodologias de pontuação incluem uma variedade de estratégias, como modelagem de recompensa, pontuação baseada em preferências ou verificação específica da tarefa utilizando juízes LLM ou regras personalizadas. Esta etapa garante que cada resposta gerada seja quantitativamente avaliada em termos de qualidade e alinhamento com os critérios.
• Treinamento de Aprendizado por Reforço (RL): Na etapa final, uma abordagem baseada em RL é aplicada para atualizar o LLM, orientando o modelo a produzir saídas alinhadas de perto com as respostas de alta pontuação identificadas na etapa anterior. Através deste processo de aprendizado adaptativo, o modelo refina suas previsões para melhorar a qualidade.
• Melhoria Contínua: Os únicos dados de que o TAO precisa são exemplos de entradas LLM. Os usuários naturalmente criam esses dados interagindo com um LLM. Assim que seu LLM é implantado, você começa a gerar dados de treinamento para a próxima rodada de TAO. No Databricks, seu LLM pode melhorar quanto mais você o usa, graças ao TAO.

Crucialmente, embora o TAO use computação em tempo de teste, ele a usa para treinar um modelo que executa uma tarefa diretamente com baixos custos de inferência. Isso significa que os modelos produzidos pelo TAO têm o mesmo custo e velocidade de inferência que o modelo original - significativamente menos do que modelos de cálculo em tempo de teste como o1, o3 e R1. Como nossos resultados mostram, modelos de código aberto eficientes treinados com TAO podem desafiar os principais modelos proprietários em qualidade.

TAO fornece um novo método poderoso no kit de ferramentas para ajustar modelos de IA. Ao contrário da engenharia de prompts, que é lenta e propensa a erros, e do ajuste fino, que requer a produção de rótulos humanos caros e de alta qualidade, o TAO permite que os engenheiros de IA obtenham ótimos resultados simplesmente fornecendo exemplos de entrada representativos de sua tarefa.

O TAO é um método altamente flexível que pode ser personalizado se necessário, mas nossa implementação padrão no Databricks funciona bem diretamente da caixa em diversas tarefas empresariais. No cerne de nossa implementação estão novas técnicas de aprendizado por reforço e modelagem de recompensa que nossa equipe desenvolveu que permitem ao TAO aprender por exploração e então ajustar o modelo subjacente usando RL. Por exemplo, um dos ingredientes que alimenta o TAO é um modelo de recompensa personalizado que treinamos para tarefas empresariais, DBRM, que pode produzir sinais de pontuação precisos em uma ampla gama de tarefas.

Melhorando o Desempenho da Tarefa com TAO

Nesta seção, aprofundamos como usamos o TAO para ajustar os LLMs em tarefas empresariais especializadas. Selecionamos três benchmarks representativos, incluindo benchmarks de código aberto populares e internos que desenvolvemos como parte de nossa Suite de Benchmark de Inteligência de Domínio (DIBS).

Para cada tarefa, avaliamos várias abordagens:

• Usando um modelo Llama de código aberto (Llama 3.1-8B ou Llama 3.3-70B) diretamente da caixa.
• Ajuste fino em Llama. Para fazer isso, usamos ou criamos grandes conjuntos de dados de entrada e saída realistas com milhares de exemplos, que geralmente é o que é necessário para obter um bom desempenho com o ajuste fino. Estes incluíram:
  • 7200 perguntas sintéticas sobre documentos da SEC para FinanceBench.
  • 4800 entradas escritas por humanos para DB Enterprise Arena.
  • 8137 exemplos do conjunto de treinamento BIRD-SQL, modificados para corresponder ao dialeto SQL do Databricks.
• TAO no Llama, usando apenas os exemplos de entrada de nossos conjuntos de dados de ajuste fino, mas não as saídas, e usando nosso modelo de recompensa focado em empresas DBRM. O próprio DBRM não é treinado nessas referências.
• LLMs proprietários de alta qualidade - GPT 4o-mini, GPT 4o e o3-mini. 

Como mostrado na Tabela 3, em todos os três benchmarks e ambos os modelos Llama, o TAO melhora significativamente o desempenho do Llama de base, até mesmo além do ajuste fino. 

Como a computação clássica em tempo de teste, o TAO produz resultados de maior qualidade quando tem acesso a mais computação (veja a Figura 3 para um exemplo). Ao contrário do cálculo em tempo de teste, no entanto, este cálculo adicional é usado apenas durante a fase de ajuste; o LLM final tem o mesmo custo de inferência que o LLM original. Por exemplo, o o3-mini produz 5-10x mais tokens de saída do que os outros modelos em nossas tarefas, resultando em um custo de inferência proporcionalmente maior, enquanto o TAO tem o mesmo custo de inferência que o modelo Llama original. 

Melhorando a Inteligência Multitarefa com o TAO

Até agora, usamos o TAO para melhorar os LLMs em tarefas individuais estreitas, como a geração de SQL. No entanto, à medida que os agentes se tornam mais complexos, as empresas precisam cada vez mais de LLMs que possam realizar mais de uma tarefa. Nesta seção, mostramos como o TAO pode melhorar amplamente o desempenho do modelo em uma variedade de tarefas empresariais.

Neste experimento, reunimos 175.000 prompts que refletem um conjunto diversificado de tarefas empresariais, incluindo codificação, matemática, resposta a perguntas, compreensão de documentos e chat. Em seguida, executamos o TAO em Llama 3.1 70B e Llama 3.3 70B. Finalmente, testamos um conjunto de tarefas relevantes para empresas, que inclui benchmarks populares de LLM (por exemplo. Arena Hard, LiveBench, GPQA Diamond, MMLU Pro, HumanEval, MATH) e benchmarks internos em várias áreas relevantes para empresas.

TAO melhora significativamente o desempenho de ambos os modelos^[t][u]. Llama 3.3 70B e Llama 3.1 70B melhoram em 2,4 e 4,0 pontos percentuais, respectivamente. TAO aproxima significativamente o Llama 3.3 70B do GPT-4o em tarefas empresariais^[v][w]. Tudo isso é alcançado sem custo de rotulagem humana, apenas dados de uso representativos de LLM e nossa implementação de produção do TAO. A qualidade melhora em todas as subpontuações, exceto na codificação, onde o desempenho é estático.

Usando TAO na Prática

TAO é um método de ajuste poderoso que funciona surpreendentemente bem em muitas tarefas, aproveitando o cálculo em tempo de teste. Para usá-lo com sucesso em suas próprias tarefas, você precisará:

• Entradas de exemplo suficientes para sua tarefa (vários milhares), coletadas de um aplicativo de IA implantado (por exemplo, perguntas enviadas a um agente) ou geradas sinteticamente.
• Um método de pontuação suficientemente preciso: para os clientes do Databricks, uma ferramenta poderosa aqui é nosso modelo de recompensa personalizado, DBRM, que alimenta nossa implementação do TAO, mas você pode aumentar o DBRM com regras de pontuação personalizadas ou verificadores se eles forem aplicáveis para sua tarefa.

Uma melhor prática que permitirá TAO e outros métodos de melhoria de modelo é criar um ciclo de dados para suas aplicações de IA. Assim que você implantar um aplicativo de IA, poderá coletar entradas, saídas de modelo e outros eventos por meio de serviços como Tabelas de Inferência Databricks. Você pode então usar apenas as entradas para executar TAO. Quanto mais as pessoas usam seu aplicativo, mais dados você terá para ajustá-lo, e - graças ao TAO - melhor será seu LLM.

Conclusão e Começando no Databricks

Neste blog, apresentamos a Otimização Adaptativa em Tempo de Teste (TAO), uma nova técnica de ajuste de modelo que alcança resultados de alta qualidade sem a necessidade de dados rotulados. Desenvolvemos o TAO para enfrentar um desafio chave que vimos os clientes empresariais enfrentando: eles não tinham os dados rotulados necessários para o ajuste fino padrão. O TAO usa cálculo em tempo de teste e aprendizado por reforço para melhorar os modelos usando dados que as empresas já têm, como exemplos de entrada, tornando fácil melhorar qualquer aplicativo de IA implantado em qualidade e reduzir custos usando modelos menores. O TAO é um método altamente flexível que mostra o poder do cálculo em tempo de teste para o desenvolvimento de IA especializado, e acreditamos que ele dará aos desenvolvedores uma ferramenta nova, poderosa e simples para usar ao lado de prompts e ajuste fino.

Os clientes do Databricks já estão usando o TAO no Llama em pré-visualização privada. Preencha este formulário para expressar seu interesse em testá-lo em suas tarefas como parte da pré-visualização privada. O TAO também está sendo incorporado em muitas de nossas próximas atualizações e lançamentos de produtos de IA - fique ligado!

¹ Autores: Raj Ammanabrolu, Ashutosh Baheti, Jonathan Chang, Xing Chen, Ta-Chung Chi, Brian Chu, Brandon Cui, Erich Elsen, Jonathan Frankle, Ali Ghodsi, Pallavi Koppol, Sean Kulinski, Jonathan Li, Dipendra Misra, Jose Javier Gonzalez Ortiz, Sean Owen, Mihir Patel, Mansheej Paul, Cory Stephenson, Alex Trott, Ziyi Yang, Matei Zaharia, Andy Zhang, Ivan Zhou 

² Usamos o3-mini-medium ao longo deste blog.
³ Este é o benchmark BIRD-SQL modificado para o dialeto e produtos SQL do Databricks.

(This blog post has been translated using AI-powered tools) Original Post