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.
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.
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.
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.