Resumo Técnico: Estrutura Agêntica para Migração de Cargas de Trabalho de Deep Learning via Aprendizado em Contexto

1. Declaração do Problema

A migração de modelos de deep learning do paradigma orientado a objetos, eager e de estado mutável do PyTorch para o paradigma funcional, sem estado e compilado via JIT (XLA) do JAX é um gargalo de engenharia significativo. Esta transição é crítica para aproveitar otimizações específicas de hardware, como as TPUs do Google, que oferecem maior eficiência de custo e velocidades de multiplicação de matrizes. No entanto, a migração manual é propensa a erros devido a diferenças fundamentais em:

• Modelos de Execução: Execução eager do PyTorch vs. compilação Just-In-Time (JIT) do JAX.
• Gerenciamento de Estado: Estados de objetos mutáveis do PyTorch vs. abstrações funcionais e sem estado do flax.linen do JAX.
• Fluxo de Controle: Loops Python padrão vs. estruturas amigáveis ao compilador, como jax.lax.scan.
• Discrepâncias Numéricas: Diferenças em layouts de tensores padrão (NCHW vs. NHWC), precisões (fp32 vs. bfloat16) e comportamentos de preenchimento (padding) em camadas de atenção.

Embora os Grandes Modelos de Linguagem (LLMs) tenham demonstrado proficiência em geração de código geral, eles têm dificuldade com tradução estrita de frameworks. Sem verificação confiável, os LLMs tendem a alucinar operações matemáticas e lógica semântica, particularmente no que diz respeito à aritmética de tensores de alta dimensão, levando a códigos que compilam, mas falham em produzir resultados numericamente equivalentes.

2. Metodologia

Os autores propõem um pipeline agêntico totalmente autônomo projetado para traduzir módulos PyTorch para JAX. O sistema trata a tradução de framework como um "problema de geração fundamentada", utilizando um fluxo de trabalho de quatro fases (ilustrado na Figura 1 do artigo) que combina Aprendizado em Contexto (ICL) com um oráculo baseado em execução.

Fase 1: Ancoragem via Aprendizado em Contexto (ICL)

Para mitigar o fenômeno "Lost in the Middle" e reduzir o ruído de métodos de recuperação amplos (como RAG), o sistema emprega uma estratégia de ICL rigidamente controlada. Em vez de recuperar fragmentos de código grandes e complexos, o contexto é limitado a um conjunto pequeno e denso de referências contendo:

• O módulo PyTorch de origem.
• A tradução idiomática em JAX.
• Casos de teste relevantes.
  Estes exemplos servem como âncoras estruturais, guiando o LLM para adotar formatação sintática e padrões funcionais corretos sem sobrecarregar a janela de contexto.

Fase 2: O Oráculo de Execução Dinâmica

Reconhecendo que os LLMs não conseguem deduzir confiavelmente a aritmética de tensores de alta dimensão, o sistema estabelece um oráculo de execução imutável.

• O módulo PyTorch de origem é executado com tensores de entrada aleatórios.
• O estado computacional dinâmico resultante — incluindo pesos inicializados (state_dict), tensores de entrada e ativações de saída — é serializado em um artefato fixo (arquivos .pkl).
• Este oráculo captura comportamentos que a análise estática ignora, como métodos de preenchimento silenciosos e mudanças de forma dinâmica, fornecendo a verdade matemática absoluta para o código JAX de destino.

Fase 3: Tradução Agêntica e Geração de Testes

O agente traduz o código fonte PyTorch para abstrações flax.linen, sendo restringido pelas âncoras de ICL. Crucialmente, o agente tem a tarefa de gerar harnesses de teste personalizados baseados diretamente no Oráculo da Fase 2. A avaliação segue um processo hierárquico de três estágios:

1. Verificação de Compilação: Verifica se o módulo JAX inicializa e realiza o trace sem erros de tempo de execução.
2. Correção de Forma (Shape): Confirma a consistência dimensional, lidando com mudanças de layout (ex: NCHW para NHWC).
3. Equivalência Numérica: O script desserializa os pesos do Oráculo, mapeia-os para o modelo JAX e verifica a equivalência numérica usando numpy.allclose com um limite de erro absoluto estrito de $10^{-7}$. Estes são chamados de casos de teste "Silver".

Fase 4: Autodebugging Iterativo

Se o código gerado falhar nos testes condicionados pelo Oráculo, o sistema entra em um loop de autocorreção. O feedback de execução, incluindo rastreios de pilha (stack traces), erros padrão e diferenças numéricas precisas, é alimentado de volta no contexto do agente. Guiado pelos templates de ICL e por este feedback objetivo, o LLM refina iterativamente o código até que ele passe pelas restrições matemáticas estritas.

3. Principais Contribuições

O artigo descreve três contribuições principais:

1. Ancoragem Estrutural via ICL: Demonstra que um conjunto pequeno e cuidadosamente curado de traduções de referência e testes ancora efetivamente o processo de geração, reduzindo significativamente o risco de alucinação em comparação com abordagens de recuperação ampla ou apenas instrução.
2. O Oráculo de Execução Dinâmica: Introduz um método para contornar as limitações aritméticas dos LLMs usando a implementação PyTorch de origem como a verdade fundamental. Isso alinha-se com a mudança para a avaliação de código baseada em execução, garantindo fidelidade matemática através de computação real em vez de estimativa teórica.
3. Guia de Ablação para Fluxos de Trabalho Agênticos: Fornece evidência empírica através de estudos de ablação que isolam o impacto de instruções, ICL e autodebugging condicionado à execução, oferecendo um modelo para futuros sistemas de tradução de código agênticos.

4. Resultados

O sistema foi avaliado em um conjunto de teste dividido em três níveis de dificuldade:

• Nível 1: Operações matemáticas fundamentais (9 tarefas).
• Nível 2: Módulos neurais e abstrações arquitetônicas (11 tarefas, incluindo MLPs, Transformers, GANs, etc.).
• Nível 3: Migração de repositório completo (10 repositórios, incluindo SAM, T5, Code Whisper, FinGPT).

Desempenho Quantitativo (Tabela 1):

• Baseline (Apenas Instrução): Alcançou apenas 9% de equivalência numérica em módulos de Nível 2.
• Instrução + Autodebugging: Melhorou para 27% de equivalência numérica no Nível 2, mas a correspondência de forma (shape) permaneceu em 55%.
• Pipeline Completo (ICL + Oráculo + Autodebugging): Alcançou 100% de equivalência numérica no Nível 1 e 91% no Nível 2.

Migração de Repositório (Tabela 2):
Em 10 repositórios de Nível 3, o pipeline completo alcançou 100% de completude na conversão de código. A equivalência numérica variou conforme a complexidade do modelo:

• Alta Fidelidade (93%–100%): Code Whisper, Multimodal Transformer, previsão de séries temporais, StyleGAN2, T5, modelos Two-tower.
• Fidelidade Moderada: FinGPT (86%), SAM2 (80%), DETR (57%).
• Legibilidade: Avaliadores humanos classificaram a legibilidade dos arquivos convertidos entre 3,75 e 4,9 em uma escala Likert de 5 pontos.

Análise de Falhas:
O artigo identifica três modos principais de falha nos baselines: Alucinações Sintáticas (falta de classes wrapper), Incompatibilidades de Assinatura de API (kwargs incorretos) e Discrepâncias de Mapeamento de Estado (erros de KeyError durante a transferência de pesos). O loop de debug iterativo do pipeline completo aborda especificamente estes problemas, forçando o agente a resolver a lógica estrutural e de mapeamento de pesos contra o Oráculo.

5. Significância e Alegações

Os autores alegam que seu pipeline agêntico e fundamentado oferece um modelo altamente confiável e escalável para migração entre frameworks. A significância deste trabalho reside em:

• Eliminação de Testes Manuais: O sistema remove a necessidade de criação manual de testes ao gerar casos de teste "Silver" fundamentados em oráculos de execução.
• Superação das Limitações dos LLMs: Ao desacoplar o papel do LLM da dedução matemática (tratada pelo Oráculo) e focar sua função na tradução estrutural e no debug, o sistema alcança resultados de alta fidelidade onde abordagens puras de LLM falham.
• Escalabilidade: A abordagem migrou com sucesso modelos complexos e de última geração (ex: SAM, T5), validando sua utilidade para desafios de engenharia do mundo real.

Limitações:
Os autores notam modestamente duas limitações atuais:

1. O método ainda não foi expandido para migrações de nível de repositório totalmente automáticas que exijam gerenciamento completo de dependências estruturais.
2. O impacto específico de variar os exemplos usados no contexto de ICL não foi explorado sistematicamente.

Apesar dessas limitações, o artigo conclui que combinar ancoragem estrutural com feedback dinâmico baseado em execução aborda efetivamente o gargalo de tradução de frameworks, permitendo a migração de cargas de trabalho de deep learning para ambientes JAX otimizados com intervenção manual mínima.