# Transformer Models for Type Inference in the Simply Typed Lambda Calculus: A Case Study in Deep Learning for Code

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 29 references that are in the map.
- First posted on arXiv 15 Mar 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [On the Variance of the Adaptive Learning Rate and Beyond](https://mlsanity.com/paper/1908.03265/) (2019)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Generative Language Modeling for Automated Theorem Proving](https://mlsanity.com/paper/2009.03393/) (2020)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [Learning to Prove Theorems via Interacting with Proof Assistants](https://mlsanity.com/paper/1905.09381/) (2019)
- [TreeGen: A Tree-Based Transformer Architecture for Code Generation](https://mlsanity.com/paper/1911.09983/) (2019)
- [Deep Network Guided Proof Search](https://mlsanity.com/paper/1701.06972/) (2017)
- [Neural-Guided Deductive Search for Real-Time Program Synthesis from Examples](https://mlsanity.com/paper/1804.01186/) (2018)
- [On the adequacy of untuned warmup for adaptive optimization](https://mlsanity.com/paper/1910.04209/) (2019)
- [Neural Program Synthesis By Self-Learning](https://mlsanity.com/paper/1910.05865/) (2019)

## Built on by (most cited first, 1 of 1)

- [Is Pre-training Truly Better Than Meta-Learning?](https://mlsanity.com/paper/2306.13841/) (2023)

## Abstract (arXiv)

Despite a growing body of work at the intersection of deep learning and formal languages, there has been relatively little systematic exploration of transformer models for reasoning about typed lambda calculi. This is an interesting area of inquiry for two reasons. First, typed lambda calculi are the lingua franc of programming languages. A set of heuristics that relate various typed lambda calculi to effective neural architectures would provide a systematic method for mapping language features (e.g., polymorphism, subtyping, inheritance, etc.) to architecture choices. Second, transformer models are widely used in deep learning architectures applied to code, but the design and hyperparameter space for them is large and relatively unexplored in programming language applications. Therefore, we suggest a benchmark that allows us to explore exactly this through perhaps the simplest and most fundamental property of a programming language: the relationship between terms and types. Consequently, we begin this inquiry of transformer architectures for typed lambda calculi by exploring the effect of transformer warm-up and optimizer selection in the task of type inference: i.e., predicting the types of lambda calculus terms using only transformers. We find that the optimization landscape is difficult even in this simple setting. One particular experimental finding is that optimization by Adafactor converges much faster compared to the optimization by Adam and RAdam. We conjecture that such different performance of optimizers might be related to the difficulties of generalization over formally generated dataset.