Type-Constrained Code Generation with Language Models
作者: Niels Mündler, Jingxuan He, Hao Wang, Koushik Sen, Dawn Song, Martin Vechev
分类: cs.LG, cs.PL
发布日期: 2025-04-12 (更新: 2025-05-08)
DOI: 10.1145/3729274
💡 一句话要点
提出类型约束解码方法，提升语言模型代码生成的正确性和可编译性
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 代码生成 语言模型 类型约束 约束解码 类型系统
📋 核心要点
- 现有语言模型在代码生成中常出现类型错误，导致代码无法编译，约束解码方法难以有效处理。
- 论文提出类型约束解码方法，利用类型系统指导代码生成，确保生成代码的类型正确性。
- 实验表明，该方法能显著减少编译错误，提高代码合成、翻译和修复任务的功能正确性。
📝 摘要（中文）
大型语言模型（LLMs）在代码生成方面取得了显著成功。然而，由于其下一个token的推断过程没有对代码的形式化方面进行建模，因此仍然经常产生无法编译的输出。虽然约束解码是缓解此问题的一种有前途的方法，但它仅被应用于处理特定领域语言或通用编程语言的语法特征。然而，LLM经常生成带有类型错误的代码，这些错误超出了语法的范畴，并且通常难以充分约束。为了应对这一挑战，我们引入了一种利用类型系统来指导代码生成的类型约束解码方法。为此，我们开发了新颖的前缀自动机和对可居住类型的搜索，形成了一种可靠的方法来强制执行LLM生成的代码的良好类型性。我们在一个基础的简单类型语言上形式化了我们的方法，并将其扩展到TypeScript以证明实用性。我们在HumanEval和MBPP数据集上的评估表明，我们的方法将编译错误减少了一半以上，并显着提高了各种大小和模型系列的LLM（包括具有超过300亿参数的最先进的开放权重模型）在代码合成、翻译和修复任务中的功能正确性。结果表明，我们的方法在用类型系统的形式规则约束LLM代码生成方面的通用性和有效性。
🔬 方法详解
问题定义：现有的大型语言模型在代码生成任务中表现出色，但生成的代码经常包含类型错误，导致无法编译或运行。现有的约束解码方法主要关注语法层面的约束，难以有效处理类型错误，这限制了语言模型在代码生成方面的应用。
核心思路：论文的核心思路是利用类型系统来指导代码生成过程，通过类型约束来减少类型错误的产生。具体来说，就是在解码过程中，只允许生成符合类型规则的token，从而保证生成的代码具有良好的类型性。这样可以显著提高生成代码的可编译性和正确性。
技术框架：该方法主要包含两个关键组件：前缀自动机和可居住类型搜索。前缀自动机用于表示当前已生成代码的类型信息，并根据类型规则确定下一个可以生成的token集合。可居住类型搜索用于在类型空间中搜索与当前上下文兼容的类型，从而指导token的选择。整体流程是，在每一步解码时，首先利用前缀自动机和可居住类型搜索确定可行的token集合，然后从该集合中选择一个token进行生成。
关键创新：该方法最重要的创新点在于将类型系统引入到语言模型的解码过程中，通过类型约束来指导代码生成。与现有的约束解码方法相比，该方法能够更有效地处理类型错误，从而提高生成代码的质量。此外，论文还提出了新颖的前缀自动机和可居住类型搜索算法，用于高效地进行类型约束。
关键设计：论文在简单类型语言上形式化了该方法，并将其扩展到TypeScript。在前缀自动机的设计中，需要考虑如何高效地表示类型信息和进行类型推断。在可居住类型搜索中，需要设计合适的搜索策略，以在类型空间中快速找到与当前上下文兼容的类型。具体的参数设置和损失函数未知。
🖼️ 关键图片
📊 实验亮点
在HumanEval和MBPP数据集上的实验表明，该方法能够将编译错误减少一半以上，并显著提高代码合成、翻译和修复任务的功能正确性。该方法在不同大小和模型系列的LLM上都表现出良好的效果，包括具有超过300亿参数的最先进的开放权重模型。
🎯 应用场景
该研究成果可应用于各种需要代码生成的场景，例如自动化软件开发、代码补全、代码修复等。通过提高生成代码的正确性和可编译性，可以显著提高开发效率和软件质量。此外，该方法还可以应用于其他编程语言，具有广泛的应用前景。
📄 摘要（原文）
Large language models (LLMs) have achieved notable success in code generation. However, they still frequently produce uncompilable output because their next-token inference procedure does not model formal aspects of code. Although constrained decoding is a promising approach to alleviate this issue, it has only been applied to handle either domain-specific languages or syntactic features of general-purpose programming languages. However, LLMs frequently generate code with typing errors, which are beyond the domain of syntax and generally hard to adequately constrain. To address this challenge, we introduce a type-constrained decoding approach that leverages type systems to guide code generation. For this purpose, we develop novel prefix automata and a search over inhabitable types, forming a sound approach to enforce well-typedness on LLM-generated code. We formalize our approach on a foundational simply-typed language and extend it to TypeScript to demonstrate practicality. Our evaluation on the HumanEval and MBPP datasets shows that our approach reduces compilation errors by more than half and significantly increases functional correctness in code synthesis, translation, and repair tasks across LLMs of various sizes and model families, including state-of-the-art open-weight models with more than 30B parameters. The results demonstrate the generality and effectiveness of our approach in constraining LLM code generation with formal rules of type systems.