Plug-and-Play Controllable Generation for Discrete Masked Models

## Abstract

Abstract This article makes discrete masked models for the generative modeling of discrete data controllable. The goal is to generate samples of a discrete random variable that adheres to a posterior distribution, satisfies specific constraints, or optimizes a reward function. This methodological development enables broad applications across downstream tasks such as class-specific image generation and protein design. Existing approaches for controllable generation of masked models typically rely on task-specific fine-tuning or additional modifications, which can be inefficient and resource-intensive. To overcome these limitations, we propose a novel plug-and-play framework based on importance sampling that bypasses the need for training a conditional score. Our framework is agnostic to the choice of control criteria, requires no gradient information, and is well-suited for tasks such as posterior sampling, Bayesian inverse problems, and constrained generation. We demonstrate the effectiveness of our approach through extensive experiments, showcasing its versatility across multiple domains, including protein design.

## 1 Introduction

Modeling complex discrete probability distributions in high-dimensional spaces is a crucial challenge across multiple domains in generative AI, including language, vision, audio, and biology. Among the various approaches, discrete masked models have emerged as powerful tools, offering robust solutions for generating and understanding discrete data. Notable examples include BERT for language modeling (Devlin et al., ), MaskGIT for image synthesis (Chang et al., ), DNABERT for DNA modeling (Ji et al., ; Zhou et al., ), the ESM series for protein generation (Rives et al., ; Lin et al., ; Hayes et al., ), and the more recent masked discrete diffusion models (see, e.g., Lou et al. (); Ou et al. (); Sahoo et al. (); Shi et al. (); Zheng et al. ()). These models typically learn the conditional distribution of each masked position given a partially masked data sequence, allowing for iterative decoding to generate a full sequence during inference.

In many practical applications of masked models, the objective extends beyond generating realistic samples from the data distribution to doing to in a *controlled* manner. This involves generating samples that align with specific constraints, conditions, or prompts, often by sampling from a conditional data distribution or maximizing a reward function (Zhang et al., ). Controlled generation is crucial in tasks such as (1) posterior sampling, where samples are drawn from a posterior distribution conditioned on observed data; (2) constrained generation, where the aim is to produce samples that meet predefined constraints. These applications highlight the growing need for generative models that not only capture the underlying data distribution, but also allow for flexible control over the generated outputs.

For controllable generation, it is desirable to establish a framework of *plug-and-play samplers* that allows for efficient sampling from the desired controlled distribution without requiring further training or fine-tuning of the underlying pre-trained model, given any suitable control criterion. This approach is computationally advantageous, as it avoids the costly and time-consuming process of retraining the model for each new task, making it a highly scalable and adaptable solution for real-world applications.

However, existing work on plug-and-play controllable generation primarily focused on the continuous domain, such as continuous diffusion models (Chung et al., ; Song et al., ; Huang et al., ), and often requires the differentiability in the control criteria. In contrast, controllable generation for discrete generative models often relies on learning-based approaches (Dathathri et al., ; Nisonoff et al., ; Li et al., ), and to the best of our knowledge, there is no plug-and-play sampler in this domain.

In this paper, we address this gap by developing a plug-and-play framework for controllable generation using discrete masked models, eliminating the need for task-specific fine-tuning. Our algorithm operates through iterative unmasking and remasking. At each step, given the unmasked positions, we apply a mean-field approximation to estimate the conditional distribution of the masked positions, sample from it via Monte Carlo, and employ importance sampling to filter the most likely samples. We then remask a portion of the newly generated positions, and repeat this unmasking-remasking process for several times until all positions are unmasked. Since the complexity of querying the masked model to obtain conditional probabilities is typically much higher than querying the reward function in most of the real-world applications, the Monte Carlo estimation introduces minimal computational overhead, keeping the sampling and filtering process efficient. In our experiments, high-quality samples can be obtained with approximately $10$ queries to the masked model and around $1000$ Monte Carlo samples, demonstrating the effectiveness of our proposed algorithm.

In summary, our contributions are as follows:

- •
We tackle the problem of plug-and-play controllable generation for discrete masked models, introducing an efficient and economical paradigm for sampling from these models.
- •
We propose a novel framework based on mean-field approximation of multi-dimensional discrete distributions and iterative masking-unmasking. This fine-tuning-free approach enables sample generation that satisfies control criteria for any (potentially non-differentiable) reward functions. To the best of our knowledge, this is the first plug-and-play controllable sampler for discrete masked models.
- •
We demonstrate the versatility of our method through multiple experiments, including sampling sequence of integers with equality constraint and designing protein sequences, highlighting its adaptability and effectiveness across diverse tasks.

Figure: Figure 1: A demonstration of Algorithm 1 with vocabulary size $N=4$, sequence length $D=8$, and number of Monte Carlo estimate $K=6$.
Refer to caption: x1.png

Notations.
The indicator function $1_{A}$ of a statement $A$ is $1$ if $A$ is true and 0 0 if otherwise.
We will use superscripts for indexing a vector, e.g., $x=(x^{1},\dots,x^{D})$.
Given $x\in\mathbb{R}^{D}$ and a set $\Omega\subset\{1,\dots,D\}$, the slicing $x^{\Omega}$ is defined as $(x^{d}:~{}d\in\Omega)$.
Given a partition $(\Omega_{1},\Omega_{2})$ of $\{1,\dots,D\}$ (i.e., $\Omega_{1},\Omega_{2}$ are disjoint and their union is $\{1,\dots,D\}$), and two vectors $x=(x^{d}:d\in\Omega_{1})$ and $y=(y^{d}:d\in\Omega_{2})$, their concatenation $x\oplus y$ is a $D$-dimensional vector whose $d$-th entry is $x^{d}1_{d\in\Omega_{1}}+y^{d}1_{d\in\Omega_{2}}$.
Finally, $f(x)\propto_{x}g(x)$ means that $f(x)=c\cdot g(x)$ for some constant $c>0$ that does not depend on $x$.

## 2 Preliminaries and Problem Formulation

### 2.1 Discrete Masked Models

Discrete masked models (or simply masked models) are designed to learn the distribution of sequential discrete data. Let $\{1,\dots,N\}$ represent a vocabulary of $N$ tokens, and consider a random sequence $X=(X^{1},\dots,X^{D})\in\{1,\dots,N\}^{D}$ of length $D$. For example, $X$ can be a sentence with $D$ words, a protein composed of $D$ amino acids, or a discrete latent representation of images tokenized by a vector quantized variational autoencoder (VQ-VAE) (van den Oord et al., ). To learn the probability distribution $p$ of $X$ given i.i.d. samples of $X\sim p$, masked models are trained by randomly replacing certain positions in the sequence with a masked token ${\bf M}$, and learning the probability of the masked positions conditional on the unmasked portion of the sequence.

Throughout this paper, for a partially observed sequence $x\in\{1,\dots,N,{\bf M}\}^{D}$, we use $\Omega:=\{d:x^{d}\neq{\bf M}\}$ and ${\cal M}:=\{d:x^{d}={\bf M}\}$ to denote the unmasked and masked positions, respectively. Formally, a masked model $\widehat{p}$ takes such a sequence as input and output a matrix $\widehat{p}(x)\in\mathbb{R}_{\geq 0}^{D\times N}$ that approximates the following probability distribution:

$$ $\widehat{p}(x)_{d,n}\approx p(X^{d}=n|X^{\Omega}=x^{\Omega}).$ $$

The rows in $\widehat{p}$ corresponding to the unmasked positions in $\Omega$ are trivial, as the probabilities are either 0 0 or $1$. To learn the remaining entries, masked models are typically trained by minimizing the cross-entropy loss:

$$ $\min_{\widehat{p}}\operatorname{\mathbb{E}}_{X\sim p}\operatorname{\mathbb{E}} _{\text{random subset}~{}\Omega\subset\{1,\dots,D\}}\left(-\log\sum_{d\not\in \Omega}\widehat{p}\left(X^{\Omega}\oplus{\bf M}\right)_{d,X^{d}}\right),$ $$

where $X^{\Omega}\oplus{\bf M}$ represents the sequence obtained by replacing all entries not in $\Omega$ with ${\bf M}$.

During inference, a direct approach is to initialize with the fully-masked sequence, select an order $\sigma$ (a permutation of $\{1,\dots,D\}$), and autoregressively sample $X^{\sigma(t)}$ based on $X^{\sigma(<t)}$ for $t=1,2,\dots,D$. However, this scheme requires $D$ queries to the masked model, which is inefficient for long sequences. To balance accuracy and efficiency, instead of unmasking one position at a time, one may introduce a decreasing function $\gamma:[0,1]\to[0,1]$ known as an unmasking schedule, which determines the number of remaining masked tokens at each step (Chang et al., ). At the $t$-th step ($t=1,2,\dots,T$) out of $T$ total steps, given $x$ with observed positions $\Omega$, the following steps are implemented:

- 1.
Predict the probability of the masked positions $p(X^{d}=\cdot|X^{\Omega}=x^{\Omega})\approx\widehat{p}(x)_{d,\cdot}$, $d\notin\Omega$ using the masked model.
- 2.
Independently sample the masked values $x^{d}\sim\widehat{p}(x)_{d,\cdot}$, $d\notin\Omega$.
- 3.
Remask $\left\lfloor\gamma(t/T)D\right\rfloor$ newly-generated positions with the lowest predicted probability $\widehat{p}(x)_{d,x^{d}}$.

Such steps will also be utilized in the design of our algorithm.

There are several equivalent formulation of masked models. First, as demonstrated in the sampling procedure outlined above, they can be interpreted as any-order autoregressive models (Hoogeboom et al., ; Shih et al., ), where the joint distribution is factorized using any arbitrary order $\sigma$, and the model learns the conditional distributions $p(X^{\sigma(d)}|X^{\sigma(<d)}),~{}d\in\{1,\dots,D\}$. Second, recent research (Ou et al., ; Sahoo et al., ; Shi et al., ; Zheng et al., ) has shown that masked models are equivalent to the masked discrete diffusion model, which involves reversing a continuous-time Markov chain that transform any data distribution into the distribution concentrated on the sequence with all masked states.

### 2.2 Controllable Generation

This paper focuses on the following task of controllable generation for masked models: given a pretrained masked model $\widehat{p}(\cdot)$ that generates from a data distribution $p$, sample from a modified distribution

$$ $q(x)=\frac{1}{Z}r(x)p(x),~{}x\in\{1,\dots,N\}^{D},$ $$

where $r(x)$ is a non-negative reward function, and the normalizing constant $Z=\sum_{x}r(x)p(x)$ is unknown. This formulation encompasses various applications, including:

- •
Posterior sampling for Bayesian inference: Suppose $X\sim p(x)$ and we have a conditional distribution $p(y|x)$ for a related random variable $Y$, where $Y|X=x\sim p(y|x)$ serves as the reward function. Given an observation $Y=y$, the posterior distribution of $X$ is $p(x|y)\propto_{x}p(y|x)p(x)$. For instance, if $X$ is a tokenized image and $Y$ is its corresponding class label given by a classifier, then sampling from $X|Y=y$ would generate an image belonging to a specific class $y$.
- •
Constrained generation: Given a specific subset $S\subset\{1,\dots,N\}^{D}$ of interest, we define the reward function as the indicator function $r(x)=1_{x\in S}$. In this case, $q$ represents the distribution $p$ truncated to the set $S$. For example, $X$ represents DNA sequences, and $S$ is the set of DNA sequences whose percentage of A is less than $30\%$.

We also assume that evaluating the reward function $r(\cdot)$ is significantly less computationally expensive than querying the mask model $\widehat{p}(\cdot)$, which is a common scenario in most of the real-world applications, and serves as an important starting point for our algorithm’s design.

## 3 Controllable Generation for Masked Models

In this section, we present our framework for plug-and-play controllable generation of masked models. We begin by drawing insights from the plug-and-play conditional generation for continuous diffusion models in Section 3.1, and then propose our novel algorithms tailored specifically for discrete masked models in Section 3.2.

### 3.1 Existing Plug-and-play Samplers for Continuous Diffusion Models

In continuous diffusion model, the data distribution is $X_{0}\sim p(x_{0})$, and a diffusion process $(X_{t}\sim p_{t})_{t\in[0,T]}$ transforms data to noise following $X_{t}|X_{0}=x_{0}\sim\operatorname{{\cal N}}\left(x_{0},\sigma^{2}_{t}I\right)$. Leveraging the data samples and the diffusion process, one can learn the score function $\nabla_{x_{t}}\log p_{t}(x_{t})$ for all $t>0$. Given a condition variable $Y$ with distribution $p(y|x_{0})$ conditional on $X_{0}=x_{0}$, the task of sampling from the posterior distribution $p(x_{0}|y)$ boils down to estimating $\nabla_{x_{t}}\log p_{t}(y|x_{t})$ for all $t>0$, where $p_{t}(y|x_{t})$ is the predicted distribution of $Y$ given *noisy* sample $X_{t}=x_{t}$.

To estimate this quantity in a training-free way, Chung et al. () observed that $p_{t}(y|x_{t})=\operatorname{\mathbb{E}}_{p(x_{0}|x_{t})}p(y|x_{0})$, and proposed to approximate the unknown distribution $p(x_{0}|x_{t})$ by the point mass at $\widehat{x}_{0}(x_{t}):=\operatorname{\mathbb{E}}(X_{0}|X_{t}=x_{t})=x_{t}+
\sigma^{2}_{t}\nabla_{x_{t}}\log p_{t}(x_{t})$, resulting in $p_{t}(y|x_{t})\approx p(y|\widehat{x}_{0}(x_{t}))$. Song et al. () later proposed a Gaussian approximation centered at $\widehat{x}_{0}(x_{t})$ with a suitable variance, estimating the expectation by the empirical mean over i.i.d. Gaussian samples. Finally, one step of backward propagation yields the gradient of $\log p(x_{0}|x_{t})$ with respect to $x_{t}$. Throughout this process, only one query to the score model $\nabla_{x_{t}}\log p_{t}(x_{t})$ is required, while the conditional probability $p(y|x_{0})$ is generally easy to obtain, minimizing computational overhead.

To sum up, the key ingredient of this approach is the approximation of the distribution $p(x_{0}|x_{t})$, which involves predicting the clean data $X_{0}$ given a noisy sample $X_{t}=x_{t}$.

### 3.2 Controllable generation for Masked Models

We first continue the exploration of conditional generation, and then extend our methodology to general controllable generation problems.

In the discrete case, unlike the Gaussian noise $X_{t}|X_{0}=x_{0}\sim\operatorname{{\cal N}}\left(x_{0},\sigma^{2}_{t}I\right)$, we now have a masking process $X_{t}|X_{0}=x_{0}$ obtained by independently masking each position with a probability that depends on $t$. As the required entity is not the score function here, we cannot directly apply the continuous diffusion model’s strategy. Moreover, the discrete state space lacks a Gaussian distribution, and the posterior mean $\operatorname{\mathbb{E}}(X_{0}|X_{t}=x_{t})$ is meaningless. However, a simple yet effective alternative exists: a mean-field approximation.

Our goal is to sample from $p(x_{0}|x_{t},y)$. By conditional independence of $X_{t}$ and $Y$ given $X_{0}$,

$$ $\displaystyle p(x_{0}|x_{t},y)$ $\displaystyle\propto_{x_{0}}p(x_{0},x_{t},y)=p(x_{0})p(x_{t}|x_{0})p(y|x_{0})= p(x_{0}|x_{t})p(x_{t})p(y|x_{0})$ $\displaystyle\implies p(x_{0}|x_{t},y)$ $\displaystyle\propto_{x_{0}}\operatorname{\mathbb{E}}_{p(x_{0}|x_{t})}p(y|x_{0 }).$ $$

As $X_{t}$ is by independently masking each position in $X_{0}$, we only need to predict the conditional probability of the masked positions in $x_{t}$ given the observed ones. As the masked model only predicts one-dimensional probability distributions, a straightforward approach is to iteratively unmasking one position at a time, requiring $|{\cal M}|$ queries to the masked model.

To avoid such computational overhead, we employ a mean-field approximation, i.e., assume that conditional on the observed part $x^{\Omega}$ of a sequence $x$, each remaining dimension in ${\cal M}$ is independent. Formally,

$$ $p(X^{\cal M}=u|X^{\Omega}=x^{\Omega})\approx\prod_{d\in{\cal M}}p(X^{d}=u^{d}| X^{\Omega}=x^{\Omega}),$ $$

which only requires one query to the masked model. While mean-field approximation introduces some error, it is effective for approximating multidimensional distributions when dependencies between different positions are relatively weak. Similar ideas have been incorporated to the sampling algorithm of MaskGIT (Chang et al., ) and the backward sampling of Continuous-Time Markov Chain (Lou et al., ; Ou et al., ; Zheng et al., ).

Building upon our insights into conditional generation, we are ready to present a solution to the general problem of controlled generation. Based on the problem settings in Section 2.2, we now illustrate our method for sampling from $q$ in a plug-and-play manner.

Suppose we have a partially observed sequence $x$ during the generation process, with observed and masked positions $\Omega$ and ${\cal M}$. We can calculate the conditional distribution of masked positions given the observed ones under $q$ as follows: for all $u\in\{1,\dots,N\}^{|{\cal M}|}$,

$$ $\displaystyle q(X^{{\cal M}}=u|X^{\Omega}=x^{\Omega})$ $\displaystyle\propto_{u}q(X^{\cal M}=u,X^{\Omega}=x^{\Omega})$ $\displaystyle\propto_{u}r(x^{\Omega}\oplus u)p(X^{\cal M}=u,X^{\Omega}=x^{ \Omega})$ $\displaystyle\propto_{u}r(x^{\Omega}\oplus u)p(X^{\cal M}=u|X^{\Omega}=x^{ \Omega})$ $\displaystyle\approx r(x^{\Omega}\oplus u)\prod_{d\in{\cal M}}p(X^{d}=u^{d}|X^ {\Omega}=x^{\Omega}),$ (1) $$

where the last line is obtained by mean-field approximation. As the normalizing constant of this distribution is unknown, one way to sample a $u$ from this distribution is by leveraging importance sampling. Let us first recall the following lemma.

###### Lemma 1 (Importance Sampling) .

Suppose two probability masses or densities $p,q$ are related through $q(x)=\frac{1}{Z}r(x)p(x)$. Then, with $x_{1},\dots,x_{K}$ i.i.d. samples from $p$, one can approximate $q$ with the following weighted empirical distribution:

$$ $q(x)\approx\widetilde{q}(x):=\sum_{k=1}^{K}q_{k}\delta_{x_{k}}(x),~{}\text{ where}~{}q_{k}=\frac{r(x_{k})}{\sum_{j=1}^{K}r(x_{j})}.$ $$

###### Proof.

Since $p(x)\approx\widetilde{p}(x):=\frac{1}{K}\sum_{k=1}^{K}\delta_{x_{k}}(x)$, we have the following approximation of $q$:

$$ $q(x)=\frac{r(x)p(x)}{\operatorname{\mathbb{E}}_{p}r}\approx\frac{r(x) \widetilde{p}(x)}{\operatorname{\mathbb{E}}_{\widetilde{p}}r}=\frac{r(x)\sum_{ k=1}^{K}\delta_{x_{k}}(x)}{\sum_{k=1}^{K}r(x_{k})}=\widetilde{q}(x).$ $$

∎

Thanks to Lemma 1, we can approximately sample from the distribution in Equation 1 by independently sampling each dimension $d\in{\cal M}$ conditional on the observed sequence $x^{\Omega}$ and then obtain their corresponding weights, using which we can sample a $u$.

Although this process generates a full sequence $x$, its alignment to the target distribution $q$ may be sub-optimal due to the lost of dependencies in mean-field approximation, especially when $|{\cal M}|$ is large (i.e., many positions are unmasked in a row). To address this, we further introduce a remasking step as discussed in Section 2.1: remask some of positions in ${\cal M}$ according to a remasking schedule $\gamma$ (e.g., by sampling a random subset of ${\cal M}$), so that in the $t$-th step among the $T$ total steps ($t=1,\dots,T$), after remasking, the remaining number of masked tokens is $\left\lfloor\gamma(t/T)D\right\rfloor$. In our experiments, we found that this uniform remasking strategy performs well. However, exploring more advanced remasking strategies is a promising area for future research.

We summarize the entire procedure in Algorithm 1. Notably, our algorithm can be easily extended for the inpainting task, i.e., given a subset $\overline{\Omega}$ of $\{1,\dots,D\}$ that has known values $x^{\overline{\Omega}}$, we aim to sample the remaining positions according to $q(X^{\overline{\Omega}^{\complement}}=\cdot|X^{\overline{\Omega}}=x^{\overline
{\Omega}})$. This modified version of the algorithm is provided in Algorithm 2.

Figure: Algorithm 1 Plug-and-play controllable sampler of discrete masked models

###### Lemma 1 (Importance Sampling) .

###### Proof.

## 4 Related Work

Conditional generation based on guidance. As an important task in controllable generation, conditional generation aims to generate a random variable $X\sim p(x)$ (e.g., image) given another random variable $Y$ known as the condition (e.g., the text description of an image). A popular approach is guidance: for example, in continuous diffusion model, the classifier guidance (Dhariwal & Nichol, ) learns a classifier $p(y|x)$ that predicts the condition $Y$ given a *noisy* sample of $X$ and leverages its information to generate $X$ conditional on $Y=y$. Classifier-free guidance (Ho & Salimans, ) trains a score model that approximates both the conditional and unconditional score functions, using a combination of them for conditional generation. These approaches can be extended to the discrete diffusion model (see Nisonoff et al. ()).

Controllable generation for discrete generative models. The study of controllable generation is an emerging area in language modeling (see, e.g., Zhang et al. () for a review). A notable work, Dathathri et al. (), proposed applying gradient updates to the key-value cache in transformers, a task-agnostic approach but requiring fine-tuning during inference. For diffusion models, a recent work Li et al. () introduced soft value-based decoding, a derivative-free algorithm that requires pre-sampled trajectories $x_{0},x_{1},\dots,x_{T}$ of discrete diffusion model to estimate a conditional expectation. This method does not exploit the special properties of the masking process. To the best of our knowledge, there are no fine-tuning-free samplers for controllable generation in discrete masked models.

## 5 Experimental Results

### 5.1 Toy Experiment on Sampling Equality-constrained Sequences

We first apply our controllable generation framework to constrained sampling, demonstrating its outstanding efficiency in challenging tasks where the constraint set is extremely sparse.

Consider the state space $\{1,\dots,N\}^{D}$ comprising sequences of integers $x=(x^{1},\dots,x^{D})$, where we fix the sequence length $D$ to $10$. Let the unconditional distribution be the uniform distribution. In this case, the masked model has a closed form, requiring no training:

$$ $\widehat{p}(x)_{d,n}=p(X^{d}=n|X^{\Omega}=x^{\Omega})=\begin{cases}1/N,&d\in{ \cal M}\\ \delta_{x^{d},n},&d\in\Omega\end{cases}.$ $$

We study sampling from the uniform distribution restricted to the following set:

$$ $S_{N}=\left\{x\in\{1,\dots,N\}^{D}:\phi(x):=x^{1}-x^{2}x^{3}-x^{4}+x^{5}x^{6}x ^{7}+x^{8}+x^{9}-x^{10}=0\right\}.$ $$

The ratio between the cardinalities of $S_{N}$ and the entire state space $\{1,\dots,N\}^{D}$ is extremely low: approximately $1.07\%$ when $N=10$, $0.20\%$ when $N=20$, and $0.07\%$ when $N=30$.

We use Algorithm 1 to sample from the uniform distribution constrained on $S_{N}$, and present the results in Figure 2. We observe that the number of Monte Carlo samples, $K$, significantly impacts the quality of generated samples, while the number of steps of unmasking $T$ has a less pronounced effect, likely due to the short sequence length. Further experimental details are provided in Section B.1.

Figure: (a) $N=10$.
Refer to caption: x2.png

### 5.2 Controllable Protein Generation

We employ ESM3 (Hayes et al., ) as the underlying protein generation model, which is a pioneering generative model for protein, capable of reasoning simultaneously over multiple modalities including sequence, structure, and function, achieving state-of-the-art performance in multiple protein generation tasks. In our experiments, we focus on generation in the sequence domain, and fix the length of sequence to $50$. The vocabulary considered is the set of $20$ standard amino acids(^1^11They are: A, C, D, E, F, G, H, I, K, L, M, N, P, Q, R, S, T, V, W, Y.). We will consider three tasks of controlled generation: generating hydrophilic and hydrophobic proteins, sampling proteins with a propensity for alpha-helices, and protein inpainting for higher percentage of alpha-helices. In all tasks, we impose a constraint based on the instability index (Guruprasad et al., ), a positive number that estimates the stability of a protein sequence. A protein with value between 0 0 and $40$ is predicted as stable. To evaluate all the required metrics of protein sequences, we use the Biopython package (Cock et al., ), a widely used tool in computational molecular biology.

Design of reward function. To begin with, we first introduce a flexible way of designing the reward function $r(\cdot)$ for a given task. Suppose we have $M$ metrics $m_{i}:\{1,\dots,N\}^{D}\to\mathbb{R}$, $i=1,2,\dots,M$ that we are interested in, and our constraint is expressed as

$$ $\bigcap_{i=1}^{M}\{x:m_{i}(x)\in A_{i}\},$ $$

where $A_{i}\subset\mathbb{R}$ is an interval. We propose the following reward function:

$$ $r(x)=\exp\left(-\sum_{i=1}^{M}w_{i}\operatorname{dist}(m_{i}(x),A_{i})^{\alpha _{i}}\right),$ $$

where $\operatorname{dist}(a,A)$ is the distance from $a\in\mathbb{R}$ to the interval $A\subset\mathbb{R}$, $w_{i}>0$ is the weight of the $i$-th constraint ($m_{i}(x)\in A_{i}$), and $\alpha_{i}>0$ determines the shape of penalty (e.g., linear or quadratic).

Sampling hydrophilic/hydrophobic proteins. Solubility is one of the key traits of protein and it is a joint consequence of amino acid sequence composition as well as 3D structure. Protein solubility is often quantified through the hydropathy index, with high hydropathy values indicating water-repelling (hydrophobic) and low hydropathy value indicating water-attracting (hydrophilic). Designing hydrophilic or hydropathic proteins have numerous important applications, yet the tasks also pose unique challenges (Qing et al., ). Low hydropathy value is often desirable for therapeutic purposes as it’s central to the protein expression and purification process. Hydrophilicity is also critical for a longer circulation time in the human bloodstream, potentially indicating a better therapeutic efficacy (Garidel, ). Proteins with high hydropathy value are often designed for transmembrane proteins, important examples of which include cell surface receptors. Transmembrane proteins are often targets of drugs, especially receptors like G-protein-coupled receptors. Designing these proteins can help create more effective disease treatments by improving the capability to modulate cellular signals (Yang et al., ).

In the following, we consider the challenging task of generating hydrophilic and hydrophobic protein sequences. We set the length of protein sequence to $50$, and use GRAVY (Grand Average of Hydropathy) value (Kyte & Doolittle, ) to quantify the hydropathy level of the generated sequence, which is defined as the average hydropathy value of a peptide or protein. Negative GRAVY value indicates hydrophilic, while positive one means hydrophobic.

We compare the GRAVY values of the samples generated by ESM3 model in both uncontrolled and controlled way in Table 1. For controlled generation, we fix the hyperparameters $w_{2}$ and $A_{2}$ for promoting low instability index, and experimented with multiple choices of $w_{1}$ and $A_{1}$ (see further details in Appendix B). The table presents the best generation result. Using our proposed controllable generation method, we successfully sample protein sequences with the desired values of GRAVY while maintaining protein stablility. We also observe a significant reduction in standard deviation among the controlled generated samples, highlighting the reliability of our method.

**Table 1: Controlled generation for high and low GRAVY values. The metrics are presented in the form of mean $\pm$ std.**
| Task | $m_{1}=\text{GRAVY}$ | $m_{2}=\text{instability}$ | GRAVY | Instability $\downarrow$ |
| --- | --- | --- | --- | --- |
| Uncontrolled | / | / | $-0.207\pm 0.552$ | $41.467\pm 17.906$ |
| Controlled,<br>high GRAVY | $w_{1}=30$, $\alpha_{1}=1$,<br>$A_{1}=[1,\infty)$ | $w_{2}=5$, $\alpha_{2}=2$,<br>$A_{2}=[0,40]$ | $1.209\pm 0.223$ | $25.176\pm 10.094$ |
| Controlled,<br>low GRAVY | $w_{1}=35$, $\alpha_{1}=1$,<br>$A_{1}=(-\infty,-1]$ | $w_{2}=5$, $\alpha_{2}=2$,<br>$A_{2}=[0,40]$ | $-1.261\pm 0.260$ | $31.569\pm 8.204$ |

Sampling alpha-helix-rich protein. Proteins fold naturally into unique three-dimensional structures based on their amino acid sequence composition. This spatial conformation further determines its molecular and cellular function. Among the possible spatial structures, alpha helix is one of the most common secondary structures in protein, where the amino acids in the polypeptide chain form into a coil (helix) through hydrogen bonding. Particularly, designing proteins rich in alpha helices is of special interest for engineering certain protein functions (Kortemme, ), such as binding and self-assembly (Sakuma et al., ). These functions are central to applications such as therapeutic protein discovery (Walsh & Jefferis, ) and alpha-helical nanofiber design (Zhang, ).

Motivated by the versatile purpose of protein with rich alpha-helix structures, in the following, we focus on generating protein sequences with secondary structure being all or predominant alpha-helices. We compare the metrics of the generated proteins by our proposed method with those generated by the unconditional model in Table 2. Our controlled generation method successfully produces sequences with a high predicted percentage of alpha-helices ($\text{helix}\%$). To verify the structural accuracy of the generated sequences, we use the folding algorithm provided by ESM3 to predict their 3D structures given the sequences. The result is displayed in Figure 3, where we randomly sample five generated sequences from both the uncontrolled and controlled generated sets and predict their structure. The controlled generated samples exhibit a higher frequency of alpha-helices in the predicted structure, confirming the effectiveness of our algorithm. For further details on the figures, please refer to Section B.3.

We also investigate the influence of the hyperparameters $w_{i}$ and $A_{i}$ in the design of reward function $r(\cdot)$. Our empirical findings suggest an optimal choice for these parameters, providing valuable insights into designing reward function for general controlled generation problems. Additional details are presented in Section B.2.

**Table 2: Controlled generation for high alpha-helix percentage. The metrics are presented in the form of $\text{mean}\pm\text{std}$.**
| Task | $m_{1}=\text{helix\%}$ | $m_{2}=\text{instability}$ | Helix% $\uparrow$ | Instability $\downarrow$ |
| --- | --- | --- | --- | --- |
| Uncontrolled | / | / | $0.315\pm 0.072$ | $41.467\pm 17.906$ |
| Controlled | $w_{1}=50$, $\alpha_{1}=1$,<br>$A_{1}=[0.8,\infty)$ | $w_{2}=5$, $\alpha_{2}=2$,<br>$A_{2}=[0,40]$ | $0.600\pm 0.116$ | $27.387\pm 12.025$ |

Figure: Figure 3: Structure of protein sequences predicted by ESM3. The upper row are the uncontrolled generated sequences, and the lower row are the controlled generated sequences. The sequences are randomly chosen.
Refer to caption: x5.png

Protein inpainting. We consider the following inpainting problem: using a protein chain with high beta-sheet percentage as a prompt to generate the remaining positions in a controlled manner to maximize the percentage of alpha-helices. In particular, the prompt is chosen as a $35$-amino-acid slice from the SM-related protein of P. Abyssi (PDB ID: 1H64), chain A, which consists of $71$ amino acids and has a secondary structure composed almost entirely of beta-sheets. The cyan motif in the subfigure wrapped in red box in Figure 4 highlights the prompt. We use Algorithm 2 to generate alpha-helix-rich proteins based on this prompt, and display the predicted 3D structure of three randomly generated sequences in the blue box in Figure 4, demonstrating the presence of alpha-helices in the generated portions. Further experimental details can be found in Section B.4.

Figure: Figure 4: Visualization of the generated sequence from protein inpainting.
Refer to caption: extracted/5897534/figs/inpaint.png

## 6 Conclusion and Future Work

In this paper, we study the task of controllable generation for discrete masked models and introduce an efficient paradigm for plug-and-play sampling. Our approach is based on mean-field approximation and iterative masking and remasking, demonstrating promising potential for real-world applications. Future research directions include exploring alternative remasking strategies beyond uniform remasking, rigorously analyzing the error bounds of our method, and extending its application to other domains such as image, molecule, and DNA.

## Appendix A Algorithm for Inpainting

See Algorithm 2. Its difference with Algorithm 1 is marked in blue.

Figure: Algorithm 2 Plug-and-play controllable sampler of discrete masked models, the inpainting version.

## Appendix B Supplementary Experimental Results

### B.1 Implementation Details of the Toy Example

We first demonstrate how we choose the reward function $r(\cdot)$. Recall that the equality constraint is $\phi(x)=0$. Thus, we propose to use the reward function $r(x)=\exp(-w\max(|\phi(x)|,m)$, where the weight $w>0$ and the truncation threshold $m>0$. We choose $w=5$ and $m=10$ throughout all of our experiments.

We fix the remasking schedule $\gamma(t)=\cos\left(\frac{\pi}{2}r\right)$ as suggested by Chang et al. ().

The values of $K$ experimented in Figure 2 are $10,20,50,100,200,500,1000,2000,5000,10000$, while the values of $T$ experimented are $2,5,8,10$.

### B.2 Tuning the Hyperparameters for Controlled Protein Generation

In this section, we study how the hyperparameters $w_{i}$ and $A_{i}$ in the reward function influences the quality of controllable generation. In particular, we focus on the task of sampling alpha-helix rich protein, fix the hyperparameters for $m_{2}=\text{instability}$, and vary the choice of $w_{1}$, $A_{1}$ while fixing $\alpha_{1}=1$ throughtout the experiments. We experiment $w_{1}\in\{10,15,\dots,45,50\}$ and $A_{1}=[a_{1},\infty)$ for $a_{1}\in\{0.5,0.6,\dots,1.3,1.4\}$, and for each pair of $(w_{1},A_{1})$, generate $16$ protein sequences and evaluate the metrics helix% and instability. The results are displayed in Figures 5 and 6. We find that

- •
For a fixed $a_{1}$, as $w_{1}$ grows, the helix% of the generated sequences would grow at first, but may decline when $w_{1}$ is larger than some threshold. This threshold becomes smaller when $a_{1}$ gets larger.
- •
For small $a_{1}$, the instability index of the generated sequences does not vary significantly among different choices of $w_{1}$. But for large $a_{1}$, the generated sequences is prone to become more unstable when $w_{1}$ is large.

These results show that there may exist an optimal choice of the hyperparameters $w_{1}$ and $A_{1}$ that maximizes the helix% of the generated sequences while keeping low instability index, which may provide insights in designing the reward function $r(\cdot)$ for general controlled generation problems.

Figure: Figure 5: Influence of $w_{1}$ and $A_{1}$ on helix% in sampling alpha-helix rich proteins.
Refer to caption: x6.png

Figure: Figure 6: Influence of $w_{1}$ and $A_{1}$ on instability in sampling alpha-helix rich proteins.
Refer to caption: x7.png

### B.3 Protein Sequences in Figure 3 .

From left to right, the uncontrolled generated sequences are:

- 1.
RAGPRAPPRSDAGRTRGVGRKGQLLVTGKLDAPTLLSLPAAVKSTGATRS
- 2.
MGFPNVPATAAPCPAAPTYEDYAAARGGSLPQVIQHALPVIFTAPLRKST
- 3.
MNQQSTADIRMLIEIGSFMNDPNMMTLINLLILSNVFILLIVIYYRWRSL
- 4.
MKFLARSTAKTEQLRERYLKTDIQILVYETIQGDFESIRLLPASVYNVSL
- 5.
MLPSPAFAISEAQATVESGSIAGPELLAVAVEAPSTQDHRVFAGEETYGV

From left to right, the controlled generated sequences are

- 1.
MINAEAADKDECRLADLLEAKELEMLELKALYLRLEEENKALKELARAMA
- 2.
TACVEKPTHGNPTLHLAAKAFNAEIILDLAFLGQKREKLTQSNLRVISEK
- 3.
MNNDEDLWWKSKGKLINKDKYKNLDNTIMYMKQNMKDIKELKGLETILNA
- 4.
MNDQTREKLLKPGAAEVFAKKYRREKEAIESRAIARVADIDEALKLAAQL
- 5.
MLAEPIGNIVTYAYVIILSILLLVKLGLAENMETSVALTTLLFSNIWQLR

These sequence are randomly sampled from the batch of generated sequences and are not cherry-picked.

### B.4 More Details of Protein Inpainting

The whole sequence of the protein is
\seqsplitERPLDVIHRSLDKDVLVILKKGFEF\seqsplitRGRLIGYDIHLNVVLADAEMIQDGEVVKRYGKIVI\seqsplitRGDNVLAISPT, where the cyan part with 35 amino acids is the prompt we use for inpainting. We choose the same hyperparameters as in Table 2. From left to right, the three sequences displayed in the blue box in Figure 4 are

- 1.
\seqsplit
MSLAVLKNSEDTLVKAELKGDVSV\seqsplitRGRLIGYDIHLNVVLADAEMIQDGEVVKRYGKIVI\seqsplitRGDSVVTVHLLTALESQIHEIEDEKAKADRAVKARTKAIKA
- 2.
\seqsplit
MEGIALKALMDFQVVMKLKGGKEL\seqsplitRGRLIGYDIHLNVVLADAEMIQDGEVVKRYGKIVI\seqsplitRGTVITLIHIPEEVDFEAALKLLEKKPKKRIRRLKAEKSKK
- 3.
\seqsplit
SMSLAMQNLMGKEMKIRLAGGMCM\seqsplitRGRLIGYDIHLNVVLADAEMIQDGEVVKRYGKIVI\seqsplitRGNCIVYLDLPDSLKDELQSHERVHQYRGLKGAHAVKEKKR

## Appendix C Codes for Algorithms 1 and 2

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 1
presentation: 1
questions: I have several questions regarding the content of the paper:
*  What was the metric used for computing/conditional generation when optimizing the helical fraction?
* How is the reward function on page 8 derived?  Additionally, why are the intervals for the metrics sometimes closed (i.e. instability with $A = [0, 40]$, and sometimes unbounded (i.e. helix % with $A = [0.8, \infty)$), and in what settings is bounded/unbounded preferable?
* Are the helical fractions correct in the protein experiment? In figure 3, the bottom two proteins seem almost identical, yet one has a helix fraction of 0.78, and the other 0.44. This does not seem quite correct.
soundness: 1
strengths: * The paper tackles a broad category of problem; namely plug-and-play conditional generation using discrete masked models without the need for fine-tuning. Additionally, they lay out in which settings their methodology would be advantageous (for example, they indicate that this method is useful when evaluating the masked model is much more expensive to evaluate than the reward function).
* The authors make a good effort at making the paper reproducible by including source code of the algorithm (as well as detailed algorithm descriptions) in the appendix.
* Figure 1 is quite good and greatly facilitates understanding of the proposed methodology. In addition, the paper clearly describes the problem which they aim to solve, subsequently provides a concrete approach which makes the problem tractable, and performs some preliminary empirical validation.
summary: This paper tackles the problem of performing conditional generation of discrete structures via masked generative models. They propose a general-purpose approach for optimizing arbitrary objective functions during the generation process. Subsequently, they provide several simplifications and concrete modelling approaches to make the problem tractable and computationally efficient. Finally, they apply the methodology to a toy problem as well as a protein generation task.
weaknesses: __Theoretical Concerns__:

* Several key aspects of the paper lack a theoretical justification or are not derived in a principled manner. For example, the proposed reward equation $r(x) = \exp\left({-\sum w_i \text{dist}(m_i(x), A_i)^{\alpha_i}}\right)$ is provided with no theoretical grounding or explanation. As best I can tell, the definition of the sampling distribution $q(z) = Z^{-1} r(x)p(x)$ would require $r(x) \geq 0$ in order for $q(x)$ to be a valid distribution. However, this is not mentioned, and many alternative reward functions could be used. I would like to see a more detailed explanation of why this reward function was chosen (either theoretical justification or empirical results).
* Similarly, the use of the mean-field approximation and importance sampling present several practical challenges which are not addressed. In the case of importance sampling, the results are heavily dependent on samples obtained from regions of high density, and thus may require many monte-Carlo samples if the proposal distribution is far from the true distribution. Furthermore, the mean-field-approximation assumes that the probabilities of the masked inputs are independent conditioned on the observed values. This is clearly not the case for domains such as images, which exhibit strong local structure. The paper would be much improved with additional analysis of the performance of the proposed methodology when the assumptions are violated and/or on larger-scale problems more representative of real-world use.
* The authors mention that the proposed method is beneficial when the complexity of querying the masked model is much higher than evaluating the reward function. Unfortunately, this is only true for trivial objective functions. For example, protein structures are typically optimized for a complex objective that is computed by another deep learning model (i.e. predicting biological activity, folding structure, etc.). This calls in question the applicability of the method to wider categories of problems, as most problems of interest will not have a closed form/cheap objective function.

__Experimental Concerns__:
* In terms of the experimental validation, the experiments performed do not provide sufficient evidence that the methodology works as intended. First, the experiment using the toy problem uses a uniform masked model with a linear objective function. As expected, the proposed approach performs well given that the problem is explicitly formulated satisfy the mean-field approximation and importance sampling schemes. No attempt is made to characterize how the method performs as assumptions are violated. Furthermore, the protein experiments are conducted using objectives which are much too simple. GRAVY (Grand Average of hydropathy) is a simple sum of values per individual amino acid. Similarly, the instability index (Guruprasad et al., 1990), consists of summing values from a lookup table for pairs of amino acids. These objectives are simple enough that the assumptions of MFA and importance sampling are not violated, but are not representative in terms of computational costs or complexity of typical protein design tasks. Finally, an experiment is performed to optimize the helical fraction of the peptides. The objective used is not clearly defined in the paper, but validation is performed using ESM3. Consequently, if ESM3 is used for the helical fraction objective, then the objective would not be cheap to evaluate, and the initial assumptions made by the paper are violated. Overall, the paper would benefit from more extensive and principled empirical validation in settings more representative of how the methodology would be used in practice.
* Another aspect of the experimental results is that both the toy problem and the protein design task consist of relatively simple 1-dimensional discrete structures. I would need to see this methodology applied to more complex discrete structures such as 2D image generation or graph structures (such as per-atom molecule design) in order to validate some of the wider-scope claims made.
* In terms of presentation, many of the figures would benefit from more detailed captions to clearly present what is being shown. For example, figure 2 seems to imply that additional monte-carlo samples enable the algorithm to achieve a high degree of success when optimizing the objective, however this is only briefly touched upon in the main text, and not at all addressed in the caption. Additionally, figures 5/6 are quite visually crowded and hard to parse. As these figures occur in the appendix, the authors could take more space to make sure that the results are clearly and unambiguously presented.

__Contribution Concerns__
* The main contribution of the paper seems to be the introduction of the sampling distribution $Z^{-1} r(z)p(x)$, and then using MFA and importance sampling to sample from this distribution. This is not a novel methodology and is well known in various Bayesian settings. To accept the paper, there would need to be a more significant theoretical contribution. Additionally, there exists pre-existing plug-and-play samplers for continuous diffusion models, this paper extends plug-and-play samplers to discrete masked models, and this does not present a significantly novel framework for conditional generation.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: 1) Why did you choose to look only at protein sequences and not natural language for controllable generation tasks?
2) Is there a relationship between the number of MC samples needed and the mean field approximation or the remasking schedule? For example, if gamma is too high (or too low), does it take more samples to achieve high final reward?
3) What are some of the limitations of this model w.r.t. the reward design? What characterizes a controllable generative task for which reward design / success would be easy / hard?
soundness: 4
strengths: Overall the paper is very well written. The motivation of the problem, controllable discrete masked model generation without training, is good, as this implies flexible controllable generation without additional computational overhead of training for each controlled generation task. The theory appears to be sound to me without any errors, arriving at the mean field approximation with importance sampling, which seems to be a reasonable approach and yields decent results on both the toy task and protein generation tasks. The paper does an excellent job presenting the work as close to diffusion models, which makes the theory sections easy to read. The experiments for the most part are well motivated and the results do support the usefulness of the approach to some controllable settings.
summary: This model presents a method for controllable generation given a pre-trained (discrete) masking model without further training. Given a reward function, the method iteratively applies masking and remasking along with a mean-field approximation and importance sampling to perform controlled (e.g., conditioned on a class variable) generation in a "plug-and-play" manner. The work lays out the grounding theory and connections to (continuous and discrete) diffusion models, motivates the approach, and demonstrates the approach both on toy sequential data and protein generation (inpainting and class conditioned) tasks.
weaknesses: No limitations are presented in the paper, and it seems like there may be some worth discussing. One is reward function design, as it's unclear whether some tasks may not have difficult to design reward functions or if there's a high dependence on reward function on success. The next is that the Monte Carlo samples seem to be quite high, the performance in figure 10 seems to indicate that even at 10k samples the model is still improving. There really should be more of a discussion about this limitation, which I believe is likely due to either the mean field approximation or the remasking schedule, but neither of these limitations / issues are discussed to any significant degree. Finally, I wonder why we're only looking at protein sequence generation as the task: why not also look at some natural language applications?

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: Q: One of the motivations is that by setting r to be the p(y|x), one can sample from the Bayesian posterior p(x|y). How accurately can this method sample the Bayesian posterior and what kind of Bayesian inference problems can it be applied to?

The reasoning for my score is that I find the claims of effectiveness and versatility lack evidence.
* Effectiveness: The experiments have no baselines therefore its difficult to evaluate if the method is effective or not.
* Versatility: The method is only evaluated on a single domain (not counting the toy example).

I am willing to increase my score if the authors argue or provide further evidence in support of these two claims in the rebuttal.
soundness: 3
strengths: The paper presents a well-justified method from conditional sampling the presence of a reward function. It details the assumptions it makes and it gives an intuition when/why someone would use this method for conditioning.

In terms of novelty, SIR is not novel, but its application to masked generative models for controllable generations is. I am not aware of other works that use this idea for masked generative models.

The paper is very well written and easy to understand. The motivation is clear, the method is well-explained and detailed. Figure 1 and Algorithm 1 give an excellent overview that makes it easy to implement.

The experimental results on protein generation are extensive. They show convincing results on two benchmarks: solubility and alpha-helix percentage. Furthermore, they include a qualitative assessment of protein in-painting.
summary: The paper proposes a method for generating conditional samples from a masked generative model. Assuming the existence of a reward model, the method draw approximate samples from the unnormalized density r(x)p(x) without requiring the generative model to be retrained.

The method applies the Sampling Importance Resampling (SIR) trick to obtain approximate samples from the target distribution over the course of the generative process.

Experimental results demonstrate the concept on a toy problem as well as showcasing impressive results on a protein generation benchmark.
weaknesses: A main weakness of the paper is the experimental results. The work is motivated by the versatility of the approach: they claim strong performance across multiple domains. However, experimental results only include protein generation benchmarks. There are not experiments on text, images or audio with the masked models that are discussed in the introduction.

Regarding the protein benchmarks, there is no baseline to compare against and there are no ablation experiments.
* Baselines: It would be good to see how the method compares to naive fine-tuning approaches (while acknowledging that the proposed method is much lighter computationally).
* Ablations: The method does not have many hyperparameters to set, but it would be good to see how the generation quality depends on the number of Monte-Carlo samples used.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 3
questions: See Weaknesses
soundness: 2
strengths: * The paper is generally well-written, and clear
* The method is relatively simple, and easy to implement
* The method only requires an unconditional model, and can be used to controllably generate from any conditional distribution given its corresponding reward function
summary: This paper proposes a method to enable controllable generation given any unconditional generative model and a reward function as conditioning signal. This is done through computing importance weights using Monte-Carlo estimates and evaluates resulting samples using the given reward function. The authors demonstrate the effectiveness of their method using a toy dataset and in the context of conditional protein generation.
weaknesses: * The novelty is relatively low, as importance sampling has been very well studied in prior works. Although to my knowledge, I have not seen it applied it in the context of controllable generation, the experiments do not well demonstrate the effectiveness of the proposed method
* Core experiments are on relatively easy (low-dim) distributions, and it is unclear as to how this method scales. How well does the method work for more complex distributions, e.g. for images, longer sequence proteins, etc? Do you need significantly more Monte Carlo samples?
* The method quite heavily relies on a good reward function -- which, in general may be difficult to properly specify. How does performance depend on how well the reward function is shaped?