Customização Serverless no Amazon SageMaker AI acelera chamadas de ferramentas com RLVR

Desafios na chamada de ferramentas por agentes de IA

Agentes de inteligência artificial que operam em ambientes produtivos dependem da chamada correta de ferramentas para consultar bancos de dados, acionar fluxos de trabalho e recuperar dados em tempo real. Contudo, modelos base frequentemente cometem erros como hallucination de ferramentas inexistentes, passagem incorreta de parâmetros e tentativas de ação indevidas, o que compromete a confiança e inviabiliza a implantação em produção.

Abordagem com customização serverless e RLVR no SageMaker AI

Para mitigar esses problemas, a AWS apresenta a customização serverless de modelos no Amazon SageMaker AI, que elimina a necessidade de gerenciar infraestrutura pesada. A técnica usada é o Reinforcement Learning with Verifiable Rewards (RLVR), onde o modelo gera múltiplas respostas candidatas, recebe uma recompensa baseada na qualidade dessas respostas e atualiza seu comportamento para favorecer as melhores.

O processo é simplificado: basta selecionar um modelo base, configurar a técnica RLVR, apontar para os dados de treinamento e a função de recompensa. O SageMaker AI gerencia o restante, incluindo o treinamento distribuído, orquestração de memória e checkpointing.

Preparação do conjunto de dados para três comportamentos distintos

O treinamento considerou três comportamentos fundamentais para agentes que chamam ferramentas:

Executar: O usuário fornece todos os parâmetros necessários e o modelo deve chamar a ferramenta correta.
Esclarecer: O usuário não fornece parâmetros suficientes, exigindo que o modelo solicite mais informações.
Recusar: O pedido é prejudicial ou fora do escopo, e o modelo deve recusar educadamente.

Foram gerados 1.500 exemplos sintéticos utilizando o Kiro, o IDE da AWS com suporte a IA, distribuídos em cinco esquemas de ferramentas (previsão do tempo, busca de voos, tradução, conversão de moeda e estatísticas). Essa geração abrange variações formais, informais e sucintas, garantindo diversidade na linguagem.

Exemplo de prompt para execução:

{
  "prompt": [
    {"role": "system", "content": "Você é um assistente útil. Ao usar ferramentas, responda com: [...]"},
    {"role": "user", "content": "Obter previsão do tempo para São Paulo"}
  ],
  "reward_model": {
    "ground_truth": "[{\"name\": \"get_weather_forecast\", \"arguments\": {\"city\": \"São Paulo\"}}]"
  }
}

Função de recompensa com pontuação em três níveis

A função de recompensa, implementada em Python, avalia as respostas do modelo comparando as chamadas de ferramentas extraídas com a verdade de solo. A pontuação é:

1.0: nome da função e todos os parâmetros corretos;
0.5: função correta, mas parâmetros incorretos ou incompletos;
0.0: função errada ou chamada indevida.

Essa granularidade permite que o modelo aprenda de forma mais precisa, reforçando respostas parcialmente corretas e evitando penalizações severas que poderiam dificultar o aprendizado.

Configuração e execução do treinamento

Na interface do SageMaker AI Studio, escolheu-se o modelo Qwen 2.5 7B Instruct para customização via RLVR. Os hiperparâmetros usados foram:

Tamanho do batch: 128
Taxa de aprendizado: 5e-6
Épocas: 3
Rollouts por prompt: 8 (geração de oito respostas candidatas para cada entrada)

O treinamento durou cerca de 40 minutos, com métricas acompanhadas via MLflow. O gráfico de recompensa média indicou que o modelo mais que dobrou a qualidade das respostas em aproximadamente 30 passos, estabilizando após o passo 20.

Avaliação em dados não vistos e ferramentas inéditas

Para validar a generalização, avaliou-se o modelo em um conjunto de 300 exemplos exclusivos, incluindo ferramentas não apresentadas no treinamento, como busca de restaurantes, preço de ações e cálculo de desvio padrão, além de casos de recusa para solicitações perigosas.

O resultado foi uma melhora de 57% na métrica customizada Tool Call Reward, saltando de 0,35 para 0,55. Outras métricas relevantes, como F1 Score Quasi e ROUGE, também apresentaram avanços entre 14 e 19 pontos percentuais. O Exact Match dobrou, passando de 11% para 21%, indicando que o modelo acertou com mais precisão a sintaxe das chamadas.

Exemplos de respostas do modelo fine-tuned:

Execução correta: Pergunta "Qual o preço da ação da AMZN?" → resposta correta chamando a ferramenta get_stock_price com parâmetro ticker "AMZN".
Pedido de esclarecimento: Pergunta "Encontre um restaurante próximo" → resposta solicitando mais detalhes sobre tipo de cozinha ou localização.
Quase acerto: Pergunta sobre desvio padrão de números → ferramenta correta chamada, porém parâmetros formatados como string ao invés de array, recebendo pontuação intermediária.

Implantação e próximos passos

Com a validação positiva, o modelo pode ser implantado diretamente do SageMaker AI Studio, seja em endpoints SageMaker AI ou Amazon Bedrock. Também é possível baixar os pesos para implantação autogerenciada.

Para melhorar ainda mais a precisão, recomenda-se ampliar o dataset com mais ferramentas, casos extremos e conversas multi-turno, além de refinar a função de recompensa para penalizar falhas específicas, como erros de formatação de parâmetros.

Além do uso em chamadas de ferramentas, RLVR pode ser aplicado em outras tarefas de raciocínio com objetivos verificáveis, como planejamento multi-etapas, extração estruturada de dados e geração de código.