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

Technische Samenvatting: Learning Bug Context voor PyTorch-naar-JAX Vertaling met LLM's

Probleemstelling
Hoewel Large Language Models (LLM's) sterke prestaties laten zien bij het vertalen van code tussen algemene programmeertalen, neemt hun betrouwbaarheid aanzienlijk af wanneer ze te maken krijgen met domeinspecifieke code, zoals deep learning-frameworks. Specifiek vormt het vertalen van code van PyTorch naar JAX unieke uitdagingen vanwege fundamentele verschillen in executie-semantiek, waaronder dynamische graaf-executie, just-in-time (JIT) compilatie, methoden voor automatische differentiatie en vectorisatie-stromen. Eerder werk, zoals UniTrans, probeert vertaling te verbeteren via iteratieve correctie met behulp van testinputs of compilatiefouten. Deze methoden slagen er echter vaak niet in om vertalingen effectief te herstellen, vooral wanneer ze vertrouwen op kleinere of "zwakkere" LLM's (bijv. die met minder dan 10 miljard parameters), wat resulteert in niet-functionele code of subtiele gedragsveranderingen die afwijken van de oorspronkelijke intentie.

Methodologie: Het T2J-framework
De auteurs introduceren T2J, een prompt-augmentatie-framework ontworpen om LLM-gebaseerde PyTorch-naar-JAX vertaling te verbeteren door gebruik te maken van gecureerde datasets van bugs en hun fixes. Het framework werkt via verschillende kernmodules:

1. Data Constructie (Fixed-Bug Dataset):

   • De auteurs beginnen met 20 PyTorch-codesnippets uit de TorchLeet dataset (probleemoplossende code).
   • Deze snippets worden vertaald naar JAX met behulp van een "zwakke" LLM (gpt-4o-mini).
   • Professionele softwareontwikkelaars worden ingehuurd om de gegenereerde JAX-implementaties te debuggen en te repareren. Dit proces omvat meerdere ronden van handmatige verificatie om ervoor te zorgen dat de JAX-code compileerbaar en uitvoerbaar is, en outputs produceert die equivalent zijn aan de oorspronkelijke PyTorch-code.
   • Dit resulteert in een gecureerde dataset van 163 bug–oplossing paren, die veelvoorkomende vertaalfouten (bijv. API-mismatches, gradiënt-verschillen) en hun bijbehorende fixes documenteren.

2. In-Context Learning via Prompt Augmentatie:

   • In plaats van de model te finetunen, gebruikt T2J in-context learning.
   • De vertaalprompt wordt geaugmenteerd met gestructureerde begeleiding afgeleid van de fixed-bug dataset. Specifiek bevat de prompt een in JSON-formaat weergegeven lijst van veelvoorkomende foutpatronen en hun oplossingen als context.
   • De "zwakke" LLM (gpt-4o-mini) gebruikt deze geaugmenteerde context om de gecorrigeerde JAX-code te genereren zonder dat dure hertraining vereist is.

3. Evaluatie-framework:

   • Datasets: Evaluatie vindt plaats op twee datasets: een Intrinsieke set (20 TorchLeet snippets met menselijk geverifieerde ground truth) en een Extrinsieke set (100 PyTorch snippets van GitHub, waarbij de ground truth benaderd wordt door een "sterke" LLM, gpt-4o, gevolgd door menselijke verificatie).
   • Metrics: Het paper stelt drie nieuwe metrics voor naast de traditionele CodeBLEU:
     • T2J CodeTrans Score: Een LLM-als-rechter metric die Bruikbaarheid en Functionele Correctheid evalueert (met en zonder referentiecode), geïnspireerd door ICE-score.
     • T2J FixCost Score: Kwantificeert de inspanning die nodig is om code te corrigeren door het aantal menselijke fix-stappen te tellen.
     • T2J Comparison Score: Een binaire beoordeling door een LLM die twee vertaaloutputs vergelijkt om te bepalen welke superieur is.

Kernbijdragen

1. T2J Dataset: De creatie van de eerste fixed-bug dataset specifiek voor PyTorch-naar-JAX vertalingen, die gedetailleerde annotaties van foutpatronen en fixes bevat om de betrouwbaarheid van LLM-codevertaling te verbeteren.
2. T2J Framework: Een nieuwe prompt-augmentatie techniek die de domeinspecifieke kloof overbrugt door gestructureerde bug-fix context te integreren, wat een schaalbare aanpak biedt voor cross-ecosysteem migraties zonder de computationele kosten van fine-tuning.
3. Empirische Validatie: Rigoureuze evaluatie die aantoont dat het voorgestelde framework de kwaliteit van vertaling voor zwakke LLM's aanzienlijk verbetert.

Resultaten

• Intrinsieke Evaluatie: Met behulp van T2J bereikte de zwakke LLM (gpt-4o-mini) een verbetering van 20% in de voorgestelde T2J CodeTrans Score (specifiek in functionele correctheid en bruikbaarheid) vergeleken met de baseline standaard prompt. De T2J Comparison Score gaf aan dat 100% van de door T2J gegenereerde vertalingen als superieur werd beoordeeld ten opzichte van de baseline.
• Efficiëntie: Het framework verminderde de menselijke inspanning die nodig is voor debugging. De T2J FixCost Score daalde van 163 stappen (baseline) naar 87 stappen, wat een reductie van ongeveer 50% in correctie-inspanning betekent.
• Prestaties: De gecorrigeerde code gegenereerd door T2J draaide ongeveer 2,5 keer sneller dan de baseline output in executietijd-tests.
• Extrinsieke Evaluatie: Op de GitHub-dataset toonde T2J verbeteringen in functionele correctheid en bruikbaarheid (tot 1,2 punten), hoewel de CodeBLEU-scores iets lager waren dan de baseline, wat de beperkingen van traditionele metrics voor dit specifieke domein benadrukt.
• Correlatie: De studie vond dat de voorgestelde T2J CodeTrans metrics de sterkste correlatie vertoonden met menselijke fix-kosten vergeleken met traditionele metrics zoals CodeBLEU.

Betekenis en Claims
Het paper claimt dat T2J een kosteneffectieve en schaalbare oplossing biedt voor het verbeteren van cross-framework codevertaling. Door een "learning bug context" dataset te cureren en in-context learning te gebruiken, demonstreren de auteurs dat zelfs zwakkere, goedkopere LLM's kwalitatief hoogwaardige vertalingen kunnen bereiken die die van sterkere modellen benaderen, mits ze worden begeleid door gestructureerde voorbeelden van veelvoorkomende fouten en hun oplossingen. De auteurs positioneren dit werk als een stap naar het moderniseren van legacy deep learning systemen en het faciliteren van migratie tussen frameworks, terwijl ze beperkingen erkennen zoals de huidige focus op probleemoplossende datasets en het gebrek aan toepassing op open-source LLM's vanwege budgettaire beperkingen.