Poster
in
Workshop: Statistical Frontiers in LLMs and Foundation Models
Towards Probabilistically-Sound Beam Search with Masked Language Models
Anna Sappington · Robert Calef · Creston Brooks · Charlie Cowen-Breen
Keywords: [ protein engineering ] [ customizable ] [ masked language models ] [ zero-shot ] [ ancient text restoration ] [ text infilling ] [ foundation models ] [ beam search ] [ efficient ] [ BERT ]
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 adapts existing generalist MLMs for the specific task of text infilling and identifies potential shortcomings in existing practices.