A Multi-agent AI System for Deep Learning Model Migration from TensorFlow to JAX

1. O Problema

A Google opera milhares de modelos de aprendizado de máquina (ML) construídos principalmente no TensorFlow. Devido à necessidade de melhor desempenho em TPUs (Tensor Processing Units) e uma arquitetura mais funcional e stateless, a empresa está migrando seu stack de IA para o JAX (especificamente usando a biblioteca Flax Linen).

Os desafios principais dessa migração incluem:

• Divergências Arquiteturais: O TensorFlow é baseado em objetos e estado (stateful), enquanto o JAX é funcional e stateless. Isso exige uma reestruturação fundamental de como camadas, inicialização de estado e regularização são implementadas.
• Complexidade de Dependências: Muitos modelos de produção dependem de "custom ops" de baixo nível que não têm equivalentes diretos no JAX, exigindo reimplementação.
• Escala e Custo: Migrar manualmente milhares de modelos exigiria centenas de anos de engenharia especializada.
• Limitações de Agentes Únicos: Sistemas de IA de agente único (single-agent) falhavam em tarefas de longo horizonte, perdendo contexto, alucinando APIs e produzindo código que não compilava ou não seguia padrões internos rigorosos.

2. Metodologia: Sistema Multi-Agente

Os autores desenvolveram um sistema baseado em uma arquitetura de agentes colaborativos (Planner, Orchestrator e Coder), utilizando o modelo Gemini 3 Pro. O sistema não é apenas um gerador de código, mas um fluxo de trabalho estruturado.

A. Engenharia de Prompt e Playbooks

Para garantir consistência e qualidade, o sistema utiliza uma hierarquia de Playbooks (arquivos Markdown fornecidos como contexto ao agente):

• General Instruction: Regras do repositório monolítico da Google.
• Style Playbook: Padrões de codificação específicos da Google.
• Task-Specific Playbook: Diferenças sutis entre APIs do TensorFlow e JAX.
• Client-Specific Playbook: Gerado automaticamente a partir de exemplos "dourados" (modelos migrados manualmente). Um agente de IA decompõe esses exemplos em unidades semânticas para criar regras específicas do cliente (ex: YouTube), permitindo que o sistema generalize para novos modelos.

B. Arquitetura dos Agentes

O sistema divide o trabalho em três papéis especializados para gerenciar o contexto e a complexidade:

1. Planner (Planejador):
   • Realiza análise estática (usando a ferramenta Kythe) para descobrir dependências.
   • Cria um plano de migração detalhado em JSON, dividindo o modelo em "passos" discretos e gerenciáveis (ex: migrar uma camada específica, criar estrutura de diretórios).
   • Garante que cada passo seja independente e testável.
2. Orchestrator (Orquestrador):
   • Atua como ponte entre o plano de alto nível e a execução.
   • Agrupa passos lógicos em "sub-passos" para manter o contexto do agente codificador focado.
   • Gerencia falhas (retry, skip, abort) e injeta conhecimento dinâmico (apenas os playbooks relevantes para a tarefa).
3. Coder (Codificador):
   • Um agente ReAct que executa um sub-passo por vez.
   • Possui ferramentas para ler/escrever arquivos, buscar código, rodar builds e testes unitários.
   • Auto-correção: O agente deve garantir que o código gerado compile e passe nos testes antes de declarar o passo concluído.
   • Gera um resumo (summary) das alterações para um "banco de memória" compartilhado, permitindo que agentes subsequentes mantenham o contexto global do projeto.

C. Avaliação de Qualidade

Como a equivalência numérica exata é difícil de validar automaticamente em modelos complexos, os autores desenvolveram um Julgador baseado em Checklist (LLM Judge):

• Um agente LLM analisa o código migrado (JAX) contra um checklist "Golden" criado manualmente.
• O checklist é agnóstico ao framework (foca em lógica funcional, ex: "embeddings de usuário são normalizados L2", em vez de sintaxe específica).
• O julgamento é "cego": o juiz não vê o código original do TensorFlow, forçando a validação da lógica no novo código.

3. Principais Contribuições

1. Sistema Multi-Agente Especializado: Demonstração de que decompor tarefas complexas de migração em planejamento, orquestração e codificação supera significativamente os agentes únicos, especialmente em modelos de grande escala.
2. Geração de Playbooks Clientes-Specíficos: Uma metodologia para criar automaticamente guias de estilo e regras de migração a partir de poucos exemplos humanos ("golden examples"), criando um ciclo virtuoso onde a IA melhora seu próprio processo de desenvolvimento.
3. Métricas de Avaliação Inovadoras: Introdução de um sistema de avaliação baseado em checklist via LLM para validar a integridade arquitetural e a completude da migração em cenários onde testes de equivalência numérica são inviáveis.

4. Resultados

O sistema foi testado em modelos reais do YouTube, variando de pequenos a extremamente complexos (até ~20k linhas de código e centenas de camadas).

• Qualidade: A configuração completa (multi_agent_yt_specific) obteve pontuações de checklist significativamente superiores (até 0.90 em modelos complexos) comparado a agentes únicos (que ficaram abaixo de 0.45). Agentes únicos frequentemente falhavam em migrar funções de perda (loss functions) e métricas.
• Velocidade:
  • Para modelos pequenos, a aceleração foi de 1.7x.
  • Para modelos de produção complexos, a aceleração foi de 6.4x a 8x em relação à migração manual.
  • Exemplo: Um modelo que levaria ~32 horas para migrar manualmente foi migrado e revisado em ~5 horas pelo sistema.
• Dependência de Complexidade: A vantagem do sistema multi-agente aumenta exponencialmente com a complexidade do modelo. Em modelos simples, agentes únicos funcionam razoavelmente bem, mas falham catastróficamente em modelos grandes sem a decomposição do Planner e do Orchestrator.

5. Significado e Conclusão

O trabalho demonstra que a migração de frameworks de ML em escala industrial é viável e economicamente viável através de sistemas multi-agentes bem arquitetados.

• Ciclo Virtuoso: O uso de IA para migrar e melhorar a própria infraestrutura de IA (TensorFlow para JAX) acelera a adoção de novas tecnologias.
• Generalização: A abordagem não se limita a TensorFlow/JAX; as técnicas de decomposição de tarefas, uso de playbooks dinâmicos e avaliação baseada em checklist podem ser aplicadas a outras transformações de código complexas.
• Futuro: Os autores planejam expandir o sistema para incluir agentes de depuração autônoma para resolver problemas de estabilidade de treinamento e latência de inferência que surgem após a migração de código.

Em resumo, o artigo valida que a combinação de planejamento estático, orquestração dinâmica e conhecimento específico de domínio (via playbooks) é a chave para superar as limitações atuais dos LLMs em tarefas de engenharia de software complexas e de longo prazo.