# Learning to Guarantee Type Correctness in Code Generation through Type-Guided Program Synthesis

arXiv: https://arxiv.org/abs/2510.10216 · HTML: https://mlsanity.com/paper/2510.10216/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 20 of its 42 references that are in the map.
- First posted on arXiv 11 Oct 2025.

## Builds on (most cited first, 20 of 20)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation](https://mlsanity.com/paper/2109.00859/) (2021)
- [LongBench v2: Towards Deeper Understanding and Reasoning on Realistic Long-context Multitasks](https://mlsanity.com/paper/2412.15204/) (2024)
- [A Systematic Evaluation of Large Language Models of Code](https://mlsanity.com/paper/2202.13169/) (2022)
- [CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X](https://mlsanity.com/paper/2303.17568/) (2023)
- [Beam Search Strategies for Neural Machine Translation](https://mlsanity.com/paper/1702.01806/) (2017)
- [SynCode: LLM Generation with Grammar Augmentation](https://mlsanity.com/paper/2403.01632/) (2024)
- [Bugs in Large Language Models Generated Code: An Empirical Study](https://mlsanity.com/paper/2403.08937/) (2024)
- [TreeGen: A Tree-Based Transformer Architecture for Code Generation](https://mlsanity.com/paper/1911.09983/) (2019)
- [What's Wrong with Your Code Generated by Large Language Models? An Extensive Study](https://mlsanity.com/paper/2407.06153/) (2024)
- [Type-Constrained Code Generation with Language Models](https://mlsanity.com/paper/2504.09246/) (2025)
- [Program Synthesis from Polymorphic Refinement Types](https://mlsanity.com/paper/1510.08419/) (2015)
- [Efficient Avoidance of Vulnerabilities in Auto-completed Smart Contract Code Using Vulnerability-constrained Decoding](https://mlsanity.com/paper/2309.09826/) (2023)
- [Grammar-Based Code Representation: Is It a Worthy Pursuit for LLMs?](https://mlsanity.com/paper/2503.05507/) (2025)
- [Program Synthesis by Type-Guided Abstraction Refinement](https://mlsanity.com/paper/1911.04091/) (2019)
- [RbSyn: Type- and Effect-Guided Program Synthesis](https://mlsanity.com/paper/2102.13183/) (2021)
- [GramTrans: A Better Code Representation Approach in Code Generation](https://mlsanity.com/paper/2510.02887/) (2025)

## Abstract (arXiv)

Language models have shown remarkable proficiency in code generation; nevertheless, ensuring type correctness remains a challenge. Although traditional methods, such as constrained decoding, alleviate this problem by externally rejecting untypable code, the model itself does not effectively learn type reasoning internally, which ultimately limits its overall performance. This paper introduces TyFlow, a novel system that internalizes type reasoning within code generation to guide the model to learn the type system. The core of our approach is a novel type-guided program synthesis system that maintains an isomorphism between type derivation trees and synthesis derivation trees, enabling a new code representation based on synthesis decision sequences rather than traditional text-based token sequences. By offloading the complexity of type system learning to the representation itself, models can redirect their computational resources toward higher-level program semantics. Our evaluation shows that TyFlow not only eliminates type errors but also significantly improves functional correctness, highlighting the importance of aligning LMs with type systems internally.