Learning Bug Context for PyTorch-to-JAX Translation with LLMs

Resumo Técnico: Contexto de Aprendizado de Bugs para Tradução de PyTorch para JAX com LLMs

Definição do Problema
Embora os Grandes Modelos de Linguagem (LLMs) demonstrem um forte desempenho na tradução de código entre linguagens de programação de uso geral, sua confiabilidade diminui significativamente ao lidar com código de domínio específico, como frameworks de aprendizado profundo. Especificamente, traduzir código de PyTorch para JAX apresenta desafios únicos devido a diferenças fundamentais de semântica de execução, incluindo execução de grafo dinâmico, compilação Just-In-Time (JIT), abordagens de diferenciação automática e fluxos de vetorização. Trabalhos anteriores, como o UniTrans, tentam melhorar a tradução por meio de correção iterativa usando entradas de teste ou erros de compilação. No entanto, esses métodos frequentemente falham em retificar traduções de forma eficaz, particularmente quando dependem de LLMs menores ou "fracos" (por exemplo, aqueles com menos de 10 bilhões de parâmetros), resultando em código não funcional ou mudanças comportamentais sutis que se desviam da intenção original.

Metodologia: O Framework T2J
Os autores introduzem o T2J, um framework de aumento de prompt projetado para melhorar a tradução de PyTorch para JAX baseada em LLM, aproveitando conjuntos de dados curados de bugs e suas correções. O framework opera através de vários módulos principais:

1. Construção de Dados (Dataset de Bugs Corrigidos):

   • Os autores partem de 20 trechos de código PyTorch do dataset TorchLeet (código de resolução de problemas).
   • Esses trechos são traduzidos para JAX usando um LLM "fraco" (gpt-4o-mini).
   • Desenvolvedores de software profissionais foram contratados para depurar e reparar as implementações JAX geradas. Este processo envolve múltiplas rodadas de verificação manual para garantir que o código JAX seja compilável, executável e produza saídas equivalentes ao código PyTorch original.
   • Isso resulta em um dataset curado de 163 pares de bug–solução, documentando erros de tradução prevalentes (ex: incompatibilidades de API, discrepâncias de gradiente) e suas respectivas correções.

2. Aprendizado em Contexto via Aumento de Prompt:

   • Em vez de realizar o ajuste fino (fine-tuning) do modelo, o T2J emprega aprendizado em contexto (in-context learning).
   • O prompt de tradução é aumentado com orientação estruturada derivada do dataset de bugs corrigidos. Especificamente, o prompt inclui uma lista formatada em JSON de padrões de erro comuns e suas soluções como contexto.
   • O LLM "fraco" (gpt-4o-mini) utiliza este contexto aumentado para gerar o código JAX corrigido sem a necessidade de um retreinamento caro.

3. Framework de Avaliação:

   • Datasets: A avaliação é conduzida em dois datasets: um Intrínseco (20 trechos de TorchLeet com ground truth verificado por humanos) e um Extrínseco (100 trechos de PyTorch do GitHub, onde o ground truth é aproximado por um LLM "forte", gpt-4o, seguido de verificação humana).
   • Métricas: O artigo propõe três métricas inovadoras além do tradicional CodeBLEU:
     • T2J CodeTrans Score: Uma métrica de "LLM-como-juiz" que avalia a Utilidade e a Corretude Funcional (com e sem código de referência), inspirada no ICE-score.
     • T2J FixCost Score: Quantifica o esforço necessário para corrigir o código contando o número de etapas de correção humana.
     • T2J Comparison Score: Um julgamento binário por um LLM comparando duas saídas de tradução para determinar qual é superior.

Principais Contribuições

1. Dataset T2J: A criação do primeiro dataset de bugs corrigidos especificamente para traduções de PyTorch para JAX, contendo anotações detalhadas de padrões de erro e correções para facilitar melhorias de confiabilidade na tradução de código por LLM.
2. Framework T2J: Uma técnica novel de aumento de prompt que preenche lacunas de domínio específico ao integrar contexto estruturado de correção de bugs, oferecendo uma abordagem escalável para migrações entre ecossistemas sem o custo computacional do fine-tuning.
3. Validação Empírica: Avaliação rigorosa demonstrando que o framework proposto melhora significativamente a qualidade da tradução para LLMs fracos.

Resultos

• Avaliação Intrínseca: Usando o T2J, o LLM fraco (gpt-4o-mini) alcançou uma melhoria de 20% no proposto T2J CodeTrans Score (especificamente em corretude funcional e utilidade) em comparação com o prompt padrão de baseline. O T2J Comparison Score indicou que 100% das traduções geradas pelo T2J foram julgadas superiores à baseline.
• Eficiência: O framework reduziu o esforço humano necessário para depuração. O T2J FixCost Score caiu de 163 etapas (baseline) para 87 etapas, representando uma redução de aproximadamente 50% no esforço de correção.
• Desempenho: O código corrigido gerado pelo T2J rodou aproximadamente 2,5 vezes mais rápido que a saída da baseline em testes de tempo de execução.
• Avaliação Extrínseca: No dataset do GitHub, o T2J mostrou melhorias na corretude funcional e utilidade (até 1,2 pontos), embora os scores de CodeBLEU tenham sido ligeiramente inferiores aos da baseline, destacando as limitações das métricas tradicionais para este domínio específico.
• Correlação: O estudo descobriu que as métricas propostas de T2J CodeTrans mostraram a correlação mais forte com os custos de correção humana em comparação com métricas tradicionais como CodeBLEU.

Significância e Alegações
O artigo afirma que o T2J fornece uma solução de baixo custo e escalável para melhorar a tradução de código entre frameworks. Ao curar um dataset de "contexto de aprendizado de bug" e utilizar o aprendizado em contexto, os autores demonstram que mesmo LLMs mais fracos e de baixo custo podem alcançar traduções de alta qualidade que se aproximam do desempenho de modelos mais fortes, desde que sejam guiados por exemplos estruturados de erros comuns e suas resoluções. Os autores posicionam este trabalho como um passo para modernizar sistemas de aprendizado profundo legados e facilitar a migração entre frameworks, reconhecendo limitações como o foco atual em datasets de resolução de problemas e a falta de aplicação a LLMs de código aberto devido a restrições orçamentárias.