Learning Extrapolative Sequence Transformations from Markov Chains

Thank you for your careful and detailed review of our paper. We are grateful for the feedback, and address your main points below:

“In section 2, you suggest your model can be useful even if you only have an approximation of s. Could you demonstrate this?”

In our extrapolation settings (sentiment and protein design), we intentionally train imperfect approximations of the true property we are attempting to optimize; the extrapolation setting assumes that we do not have data outside a certain range, and that we can train a scorer to approximate the oracle, but that that scorer will not be reliable outside of the range of values it was trained on; therefore reviews outside the training region are OOD for our guide, which may return unreliable predictions in this extrapolation range (see e.g. https://arxiv.org/abs/2006.10108). Specifically, the sentiment scorer is trained on reviews between 2 and 4 stars, and the protein scorer is trained with values above -5, despite the objective being to generate reviews or proteins beyond that range . We therefore believe that our results demonstrate that our model is useful even with an imperfect approximation of s.

Could you compare to this paper that also trains on improving pairs of data?

Thank you for the reference; this is a very interesting approach that we will include in our discussion. The key difference with our work is that PropEn performs optimization in latent space and subsequently decodes from that latent space. However, latent space models such as VAE's are well-known to suffer from posterior collapse when using highly flexible generative models such as the large language models we use in our experiments. To sidestep this issue, we perform both search (MCMC) and optimization ($q_{\theta}$) in discrete space. As such, we compare with other pair-learning based methods that operate in discrete space. Our primary baseline is ICE, which also trains on improving pairs of data and is state-of-the-art for the tasks we consider. Additionally, PropEn focuses on single-property enhancement, while our tasks as implemented require multi-attribute optimization: protein stability and similarity to original protein for protein engineering, fluency and sentiment for the sentiment task, and EER and semantic similarity for anonymization.

Why don't you compare with other discrete optimization methods? Ex. https://openreview.net/forum?id=ZMP0Bki9aK

While in principle this approach is complementary to ours, since we could benefit from these same proposals to improve the efficiency of our MCMC step, in practice it would not be straightforward to apply this approach to the language modeling applications we consider here, due to the size of the vocabulary and sequence lengths. Generally speaking, we view improvements to MH for language models as a promising area for future work, which is complementary to our direction here about learning efficient extrapolation models $q_{\theta}$ on the basis of the Markov chains from MH. We expect that more efficient exploration will lead to a more efficient exploration step and a more effective $q_{\theta}$ for the same amount of compute.

Why did you not anneal in your MCMC?

We employ MCMC to explore sequence transformations that improve the objective, and then use our learned $q_{\theta}$ to greedily maximize this objective. While annealing MCMC could help it attain a better local optima, this could actually be counterproductive to the goal of exploring productive sequence-to-sequence transformations. Compared to other works such as https://openreview.net/forum?id=ZMP0Bki9aK above, we seek to efficiently optimize using $q_\theta$ in as few steps as possible, while annealing would require a much larger number of steps.

“One could spend the compute training $q_{\theta}$ on running more MCMC.”

While it is true that the time spent training $q_{\theta}$ could be used to run MCMC for longer, the training cost is a fixed cost that occurs offline, while running further iterations of MCMC is an online cost which would add more computational expense for each inference example. We emphasize that once $q_{\theta}$ is trained, it can be applied to new instances without further MCMC steps. Training $q_{\theta}$ allows for rapid inference on unseen examples, while MCMC requires an expensive sampling process for each inference example.

We appreciate the detailed and thoughtful feedback, and aim to address remaining questions and concerns here.

“While intermediate state generation improves performance in the protein domain, its impact on anonymization and sentiment tasks remains unclear.”

In the case of our sentiment task, there is no clear benefit to adding states: the task is sufficiently simple that a single edit is the most effective way to extrapolate, hence why we consider first/best to be the best method in this case. In more complex tasks we find it to be more important to have additional steps as there can be a generalization benefit, as in protein synthesis and anonymization. Additional edits consistently decrease semantic similarity for anonymization but improve EER up to a certain point, after which the SBERT decreases without consistent improvement in EER. We select our number of states based on this point.

Missing information on computational burden: While this may be a minor issue, it would have been helpful to include more details on the additional computational cost associated with model training and inference.

Thanks for the suggestion! Indeed, the computational burden reported in our tables is purely inference-time computation. We recognize that our model has the additional computational burden of generating training data and fine-tuning a model on that data. However, this is a single fixed cost at training time. $q_{\theta}$ takes seconds for each inference example where MCMC takes ten minutes or longer for a single batch at inference time. The number of iterations in our tables are meant to provide an understanding of the scale of the ongoing computational cost at inference time, and does not account for the fixed (“offline”) training cost. We will be sure to include more details on the offline training costs in our revisions.

Thanks for the thoughtful response. We address some of the raised questions here.

”If the scorer is mis-specified or the MCMC exploration is very limited, the AR model will learn a suboptimal strategy…”

These are valid concerns. A strength of our approach is that it benefits from the wealth of more sophisticated MCMC approaches; while we choose to use a basic sampler in this work, approaches such as the suggested adaptive sampling could lead to improved MCMC chains and thus better performance for $q_{\theta}$. We consider this to be a strength of our approach, as the theoretical properties of MCMC allow for more exploration and better fit of the target distribution as MCMC is run for more steps, offering a mechanism to improve generalization at the expense of further offline training time.

“... In the context of model extrapolation, there exists a risk of mode collapse…”

We acknowledge that there is a possibility of mode collapse for extrapolation tasks, and that we should examine the diversity of our outputs. In consideration of this point, for our protein task, we counted the number of unique outputs, finding that 100% of the 10k proteins generated using $q_{\theta}$ were unique. For sentiment and anonymization, we ran corpus BLEU score between all pairs of generated sentences, finding that we achieve only 1.39 BLEU for sentiment and 0.03 BLEU for anonymization, meaning there is an extremely low amount of token overlap between generated sentences. Furthermore, in Appendix E we show randomly selected examples of our generated outputs for sentiment and anonymization, which demonstrate significant sample variance. In light of this evidence, we believe it is unlikely that mode collapse is occurring.

We sincerely appreciate the feedback offered in this review, and hope to address some concerns.

”A theoretical perspective of… what makes an optimal $q_{\theta}$ for the MCMC algorithm is missing.”

We agree theoretical grounding is important. Nonetheless, compared to previous approaches, we believe our approach has demonstrated its robustness across three different settings, with minimal problem-specific tuning. Furthermore, though the learned extrapolation model $q_{\theta}$ lacks theoretical guarantees despite its empirical success, the MCMC search procedure that is the basis for fitting $q_{\theta}$ inherits the usual benefits of MCMC. Since extrapolative generation is an understudied area of deep learning, we hope our contribution can motivate more theoretical results in the future.

“...the MCMC sampling to generate data for the training process of the model should also be considered…”

While we note the reviewer’s concern, we emphasize that although we use MCMC as a source of training data for $q_{\theta}$, MCMC is not required on new problem instances; $q_{\theta}$ can be applied to new starting sequences without MCMC if the extrapolation criteria are the same. Thus, besides the extrapolation benefits of $q_{\theta}$, it also effectively amortizes the sampling process.

“an ablation on how the number of iterations for data generation in model training….”

We appreciate this suggestion and present our ablation of limited MCMC exploration on our anonymization task. We limit the MCMC chains to 50% and 25% of their length before constructing training episodes. Below, we show that increasing resources spent during data generation leads to improvements in our trained $q_{\theta}$ model. We consider this to be a strength of our approach; the theoretical properties of MCMC allow for more exploration and better fit of the target distribution as MCMC is run for more steps, offering a mechanism to improve generalization at the expense of further offline training time. In particular, the table below suggests that running MCMC for longer (e.g., 150%, 200%) could lead to better results than presented in the paper.

“...compare to demonstrate the key components that lead to the best $q_{\theta}$…”

Regarding ablations of $q_{\theta}$, thank you for the suggestion. We do analyze several important components of $q_{\theta}$ in the text of the paper, most notably the data selection strategy (Section 3.4) and reward choice (Appendix A). After following reviewer suggestions, we can now additionally present results for the effects of the number of iterations for data generation (see above) and for the effects of episode length on the anonymization task:

We find that additional edits improve EER at the expense of decreasing semantic similarity; this is only true up to a certain point, after which the SBERT decreases without consistent improvement in EER, demonstrating the importance of this component in $q_{\theta}$. Multiple iterations are similarly useful for protein synthesis.

We are happy to discuss additional factors should there be further concerns. For example, the paper as it stands focuses on a basic autoregressive architecture for $q_{\theta}$. However, it is likely that our results could be improved, for example by initializing from larger pre-trained models for $q_{\theta}$ or employing different training strategies.

Regarding comparisons to further baselines, we agree this would strengthen the work. To our knowledge, our comparisons include SOTA extrapolative baselines for the challenging extrapolation tasks we consider, but of course there are many other tasks to consider, such as the design of cell-type-specific promoters for gene delivery considered in [1]. However, the complex approach described in [1] involves a very application-specific pipeline involving 5 different steps which would be hard to adapt to our more general setting. On the other hand, our approach is largely application-agnostic provided a suitable pre-trained LM is available.

Regarding comparisons to RL-based methods such as [1], our approach may be viewed as performing off-policy reinforcement learning, even though we do not explicitly cast it that way (see e.g. https://arxiv.org/pdf/2106.02039). We considered using a more overt framing of the work as RL, but felt the additional notational baggage made it harder to understand the main idea. We view our contribution as a first effort to adapt off-policy RL to extrapolation settings by bridging RL and MCMC. More discussion of this point will be included in the paper.