摘要
将深度学习模型从 PyTorch 的灵活面向对象设计迁移到 JAX 的无状态功能设计通常是一项手动且容易出错的任务。自动迁移面临挑战，因为大型语言模型（LLMs）在严格且动态的 API 对齐上表现不佳，且在精确操作时容易出错。
我们提出了一种完全自主的系统，结合了上下文学习（ICL）与基于预言机的自我调试。首先，我们整理了一个 ICL 上下文，作为 JAX 风格和测试用例生成的严格参考。其次，我们不依赖 LLM 推导数学输出，而是运行源 PyTorch 模块以获取其实际动态张量状态。这创建了一个不可更改的执行预言机。
然后，我们使用自主代理循环基于预言机数据合成测试。测试用例被反复执行，回溯信息发送回 LLM 进行自我修正。消融实验表明，将 ICL 参考与预言机基础和自我调试相结合，大大优于纯指令和基本代理基线。这种改进并没有增加过多的计算开销。我们的轻量级管道在神经模块上实现了 91% 的数值等价性（基线为 9%，指令 + 自我调试为 27%），为跨框架迁移提供了一个高度可靠、可扩展的蓝图。
这一方法已在多个最先进的模型上进行了验证，包括 SAM（Segment Anything）、T5、Code Whisper 等，显示出高数值等价性。
代码
博主点评： 本文提出的自主迁移系统通过结合上下文学习和自我调试，显著提高了深度学习框架间迁移的准确性与效率，展示了前沿技术的潜力与应用价值。特别是在复杂模型的迁移中，这种创新方法为研究人员和开发者提供了极大的便利。