# Towards Probabilistically-Sound Beam Search with Masked Language Models

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 35 references that are in the map.
- First posted on arXiv 22 Feb 2024.

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [BERTScore: Evaluating Text Generation with BERT](https://mlsanity.com/paper/1904.09675/) (2019)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [A large annotated corpus for learning natural language inference](https://mlsanity.com/paper/1508.05326/) (2015)
- [Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization](https://mlsanity.com/paper/1808.08745/) (2018)
- [It's Not Just Size That Matters: Small Language Models Are Also Few-Shot Learners](https://mlsanity.com/paper/2009.07118/) (2020)
- [DiffusionBERT: Improving Generative Masked Language Models with Diffusion Models](https://mlsanity.com/paper/2211.15029/) (2022)
- [BERT has a Mouth, and It Must Speak: BERT as a Markov Random Field Language Model](https://mlsanity.com/paper/1902.04094/) (2019)
- [Enabling Language Models to Fill in the Blanks](https://mlsanity.com/paper/2005.05339/) (2020)
- [Blank Language Models](https://mlsanity.com/paper/2002.03079/) (2020)
- [Latin BERT: A Contextual Language Model for Classical Philology](https://mlsanity.com/paper/2009.10053/) (2020)
- [Text Infilling](https://mlsanity.com/paper/1901.00158/) (2019)
- [Deriving Language Models from Masked Language Models](https://mlsanity.com/paper/2305.15501/) (2023)
- [Bidirectional Beam Search: Forward-Backward Inference in Neural Sequence Models for Fill-in-the-Blank Image Captioning](https://mlsanity.com/paper/1705.08759/) (2017)
- [What Do You Get When You Cross Beam Search with Nucleus Sampling?](https://mlsanity.com/paper/2107.09729/) (2021)

## Abstract (arXiv)

Beam search with masked language models (MLMs) is challenging in part because joint probability distributions over sequences are not readily available, unlike for autoregressive models. However, estimating such distributions has important domain-specific applications such as ancient text restoration and protein engineering. Here we present probabilistically-sound methods for beam search with MLMs. First, we clarify the conditions under which it is theoretically sound to perform text infilling with MLMs using standard beam search. When these conditions fail, we provide a probabilistically-sound inference time modification with no additional computational complexity and demonstrate that it is superior to the aforementioned beam search in the expected conditions. We then present empirical results comparing several infilling approaches with MLMs across several domains. Notably, our method probes the inductive biases of MLMs and explores the surprising contextual sensitivity of mask tokens for text infilling.