Learning Extrapolative Sequence Transformations from Markov Chains

Sophia Hager, Aleem Khan, Andrew Wang, Nicholas Andrews

TL;DR: We can distill extrapolative sequence transformations from Markov chains.

Abstract: Most successful applications of deep learning involve similar training and test conditions. However, tasks such as biological sequence design involve searching for sequences that improve desirable properties beyond previously known values, which requires novel hypotheses that \emph{extrapolate} beyond training data. In these settings, extrapolation may be achieved by using random search methods such as Markov chain Monte Carlo (MCMC), which, given an initial state, sample local transformations to approximate a target density that rewards states with the desired properties. However, even with a well-designed proposal, MCMC may struggle to explore large structured state spaces efficiently. Rather than relying on stochastic search, it would be desirable to have a model that greedily optimizes the properties of interest, successfully extrapolating in as few steps as possible. We propose to learn such a model from the Markov chains resulting from MCMC search. Specifically, our approach uses selected states from Markov chains as a source of training data for an autoregressive model, which is then able to efficiently generate novel sequences that extrapolate along the sequence-level properties of interest. The proposed approach is validated on three problems: protein sequence design, text sentiment control, and text anonymization. We find that the autoregressive model can extrapolate as well or better than MCMC, but with the additional benefits of scalability and significantly higher sample efficiency.

Lay Summary: Neural networks have trouble "extrapolating", or generating sequences with qualities more extreme than they've seen in the past. For instance, a model trained on 2-4 star reviews is unlikely to generate a 5-star review. In some use cases, this ability would be very useful. It's possible to extrapolate by repeatedly changing one small part of a sequence, checking if the model thinks it's going to score better, and keeping the change if it does. However, because this involves so many small changes, it takes a long time to run for any sequence, which makes it impractical. We take these long chains of changed sequences and choose a handful of examples, then train a model on those examples to learn a model that generates much shorter chains which still extrapolate. We look at three settings and find that our model accomplishes extrapolation in much less time than generating the long chain, and in some cases extrapolates better than the long chain.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Link To Code: https://github.com/sophia-hager/learning-MCMC-extrapolation

Primary Area: Deep Learning->Generative Models and Autoencoders

Keywords: Large language model, Markov chain Monte Carlo

Submission Number: 8564