用开发者修复的错误数据提升LLM从PyTorch到JAX的翻译准确率。
Learning Bug Context for PyTorch-to-JAX Translation with LLMs
- 构建包含20个核函数及对应修复方案的T2J基准数据集。
- 在gpt-4o-mini上使用上下文学习，翻译准确率提升最高达20%。
- 专为深度学习框架迁移设计，适合研究代码生成与框架互操作者。
大型语言模型在通用编程语言间代码翻译中表现良好，但在领域特定代码（如从PyTorch到JAX）中可靠性下降，因框架特有接口和执行语义导致细微错误或非惯用用法。现有研究指出，通过人工标注的错误修复数据可提升生成质量，但此类资源在PyTorch-to-JAX场景中仍稀缺。本文提出T2J基准，包含20个来自TorchLeet数据集的核函数，经gpt-4o-mini翻译后由开发者修复。利用T2J进行上下文学习，使gpt-4o-mini在新提出的T2J-CodeTrans-Score指标上最高提升20%。
原文摘要 · Abstract (English)
Large language models (LLMs) have shown strong performance on code translation between widely used programming languages. However, translation becomes much less reliable for domain-specific code, where correctness depends on framework-specific APIs and execution semantics. One example is translating deep-learning code from PyTorch to JAX, where LLM outputs often contain subtle bugs or non-idiomatic usage that prevents execution or changes behavior. Prior work suggests that curated bug-fix data from LLM-generated code can help improve code generation quality, but such resources are still limited for PyTorch-to-JAX translation. In this work, we introduce T2J, a benchmark of LLM translation bugs paired with developer-written fixes for PyTorch-to-JAX code. We start from 20 kernels in the TorchLeet dataset, translate them to JAX using the weak LLM gpt-4o-mini, and hire software developers to debug and repair the generated JAX implementations. We then use T2J to improve PyTorch-to-JAX translation for the weak LLM gpt-4o-mini via in-context learning. Our evaluation shows that using T2J yields up to 20% improvement of our proposed metric T2J-CodeTrans-Score.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。