Cliqueformer: Model-Based Optimization With Structured Transformers

Critique

1. The proposed architecture introduces new hyper-parameters ($N_{clique}$ and $d_{clique}$), and the authors claim that $d_{clique}=3$ consistently achieves good performance. However, the conducted ablation reveals that the model is sensitive to these hyper-parameters, and some settings of them can achieve bad performance.

The reviewer is right in pointing this out. Let us first clarify our statement about $d_{clique}=3$. As we described in 5.2 Ablations and Appendix B (Appendix C in the revised version), in our experiments, we fixed $dim(z) = d_{knot} + N_{clique} \cdot (d_{clique} - d_{knot})$ to match $dim(x)$ as closely as we can, and thus by setting $d_{clique}=3$ and $d_{knot}=1$ we have also prescribed the value of $N_{clique}$. We suggest setting $d_{clique}=6$ for tasks with $dim(x) > 100$ to prevent computing attentions greater than $50\times 50$ - we can highlight this point in the final version. We indeed get good (though not always optimal) results with this scheme. Note that on 7 of our tasks, Cliqueformer uses these values, and on Superconductor where $3$ was sub-optimal, it does not underperform COMs that much (Figure 5b, $N_{clique}=20$). We can use this result if the reviewer suggests so.

The ablation study that the reviewer refers to, first of all, serves the purpose to prove that our novel mechanism that decomposes the latent variable into cliques does play a big role in the performance. We imagine that the reviewer agrees that failure to demonstrate this dependence would prove the mechanism useless. Indeed, the deep learning literature is abundant in valuable techniques that do not bring improvement under every single setting. For example, Dropout [1] can help generalization, but at a certain (mostly task-dependent) threshold, it leads to underfitting (see [1] Figure 9). Similarly, the NeurIPS-published work that the reviewer cited work on BDI reports in its ablation that removing one of their losses $L_{l2h}$ and $L_{h2l}$ can sometimes improve the performance. Thus, we kindly ask the reviewer to re-consider this comment.

2. The authors claim that Cliqueformer achieves state-of-the-art performance while some of the recent baselines have not been compared against.

We agree with the reviewer that such a claim requires an enlarged list of baselines. In the updated version, we added a comparison against recent diffusion-based DDOM [2], as well as gradient ascent with Transformer backbone, to provide further evidence that the performance of our model does not solely come from the backbone. The results strengthen evidence of our claim. We apologize for our prior negligence. Due to issues with DesignBench benchmark described in Appendix B (now, Appendix C), we had to spend extra time searching for alternative data which left us with less time to implement baselines. We settled on COMs due to its strong performance even against more recent methods. We will be extremely grateful if the reviewer takes this into consideration while re-examining our paper.

Questions

1. Do you have any intuition on how to set $N_{clique}$ and $d_{clique}$?

This is a very valid question. As of now, we can observe what we stated in the paper, which is that $dim(z) \approx dim(x)$ and setting $d_{clique}=3$, or $d_{clique}=6$ for larger designs, tends to work well. We belive that the reason for that is that having more cliques - and thus smaller cliques - allows the decoder to learn more expressive attention patters, similarily to what was found in Diffusion Transformer [3] (Figure 6 bottom row). However, the cliques shouldn't be too small (like size 1) to enable more expressivity in the predictive model too. Our choice seems to find a sweet spot in this trade-off. A large-scale study of these hyper-parameters, both empirically and theoretically, is an exciting avenue of future work.

References

[1] N. Srivastave et al. 2014; Dropout: A Simple Way to Prevent Neural Networks from Overfitting

[2] S. Krishnamoorthy et al. 2023; Diffusion Models for Black-Box Optimization

[3] W. Peebles and S. Xie 2023; Scalable Diffusion Models with Transformers

Thank you for your effort, and I apologize for the delayed response.

Critique

1-a. Your responses have addressed my concerns regarding the hyperparameters. However, in the current version, I found it somewhat challenging to compare the performances presented in Figure 5 and Table 1. It would be helpful to include a brief comparison with the results from COMs either in the figures, captions, or main text. Additionally, the text in the figures is quite small, which makes it difficult to read.

1-b. I agree that the ablation studies don't need to demonstrate improvement in every setting. Nonetheless, I believe the comparison between $N_{\text{clique}}=1$ and $N_{\text{clique}}>1$ is important to understand the proposed method (though I don't expect Cliqueformer to outperform in every setup). Do you have any intuition for the results? For instance, I expected that using more cliques would be beneficial for high-dimensional tasks, but the results show that increasing $N_{\text{clique}}$ tends to give better results in TFBind-8 and lower scores in Superconductor.

2. Thank you for including a new baseline despite the limited time. It might be better to moderate expressions like "outperforms all baselines" in favor of something less strong than "state-of-the-art performance." Again, I don't mean to suggest that the performance of Cliqueformer is insufficient.

Questions

1. Thank you for the clarification.

I will consider raising my score if my remaining concerns in 1-a are addressed.

Thank you for your continuing work in helping us improve the paper. We are excited to respond to your suggestions and questions. We just submitted a revised version accordingly.

1a.

To improve the readability of Figure 5, we made new figures, in which we increased the font size of axis labels. We hope the figures are more legible now. The Reviewer asked about the relation of our results from Table 1 and the scores from the COMs paper. The key difference is that, in COMs, the reported results are in their raw form, without normalization (in the ICML version - in the Arxiv version, they are normalized like in DesignBench). This might make interpretation of the results difficult (e.g., what does it mean to score 333.4 in DKittyMorphology?), and that's why newer papers, like DesignBench (Trabucco et al. 2022) and DDOM (Krishnamoorthy et al. 2023), normalize their scores. To provide more clarity in reading the results in Table 1, in the caption under the table, we explain the reported scores as:

The values were normalized with the min-max scheme, where the minimum and the maximum are taken from the dataset, so that the scores of the designs in the dataset are in range $[0, 1]$. We note that, unlike Trabucco et al. 2022, we take the maximum from the data available for the MBO model training (the union of the train and test data), and not from the oracle training data, to make the results more interpretable (see Appendix C).

1b.

We agree with the reviewer that these observations seem tricky. For now, we can offer the following intuition: when we increase the number of cliques, we end up having a more expressive attention layer in the decoder, which can result in improved performance. On the other hand, note that the increased number of cliques implies more aggressive decomposition assumptions on the target function $f(z)$, potentially limiting the expressivity of our predictive model. We believe that, for different functions, there is a threshold into how many sub-functions it can be decomposed. Thus, we should expect the performance to improve until that threshold, and then start degrading once it's reached. Thus, we think that the shape of the curve in Figure 5 for TFBind8 is as expected. Probably, we did not get the same shape in Figures (b) and (c) because of the large exponential leaps in the number of cliques tested. Note that we made those leaps because we needed $N_{clique}$ to divide $dim(z) - d_{knot}$ (see Appendix C, Line 860).

We appreciate the reviewer's suggestion about the most proper choice of words. We have replaced the "state-of-the-art performance" with "outperforms the baselines" in both abstract and the experiment section. Does that seem like a reasonable change?

Critique

1. The baselines seem quite outdated. Shouldn't the authors compare against more recent methods, like the diffusion-based DDOM?

The author is right in pointing this out. We implemented and evaluated DDOM on our set of tasks and added the results to the revised paper. We also evaluated gradient ascent with Transformer backbone to provide evidence of the impact of the novel mechanisms in Cliqueformer.

2. Superconductor task uses a surrogate function instead of the true oracle, which was found not to be reliable.

It is true that a surrogate oracle is used in Superconductor, as well as in all continuous tasks. This is an absolute necessity, since evaluation of continuous designs would require synthesizing them in a wet lab. The accuracy of this oracle does not really matter that much since the training data are also labeled with it, which keeps the target function of our MBO algorithm consistent. Lastly, we investigated the evidence of that inaccuracy for Superconductor provided by the cited work of Nguyen et al. 2023 and found out that only a line plot, without alpha transparency, of several of thousands of examples was shown. This makes the graph unreadable and potentially misleading. On the contrary, it is known that this oracle achieves a very high Spearman-rank correlation of 0.927.

Questions

1. Figure 5 shows that the performance varies with different values of $N_{clique}$. How did the authors pick this value?

Finding a hyper-parameter setting scheme was part of our work. We conducted this investigation on LRBF tasks which, being synthetic, do not require that we collect external data, but instead allow us to study problems of arbitrary dimensionality. We have found the following scheme to work well reasonably well: we fix $d_{knot}=1$ and, (1) for tasks where $dim(x)<100$, we fix $d_{clique}=3$ and set $N_{clique}$ so that $dim(z)=N_{clique} \cdot (d_{clique} - d_{knot}) + d_{knot} \approx dim(x)$; (2) for tasks where $dim(x) \geq 100$, we fix $d_{clique}=6$, and repeat the procedure to find $N_{clique}$ - the reason for doing that is to decrease the cost of computing attention layers which is quadratic in $N_{clique}$. In Table 1, we made a single exception from this rule - for Superconductor, where another setting was indeed much better.

Generally, however, we believe that these values can be tuned in practice, in a very similar fashion MBO researchers evaluate their methods in papers. Namely, one can pre-train a surrogate oracle on a large dataset, label the data with the surrogate, and train the MBO model on a subset of the data with the new labels. This way, one can sweep over hyper-parameters to tackle the real dataset with.

2. Why does the decoder have to be a transformer? What happens if it is replaced with an MLP?

In our early experiments, we tried to implement our model with MLPs. This proved very difficult. We believe that the reason for that is the difficulty posed onto the model - to learn representations that provide good reconstructions, satisfy Gaussianity properties, and can predict the target function. After we switched to Transformers, we could learn such representations. The only place where we did not see benefit of such a change was the predictive model $f_{\theta}(z)$, where we stack with MLPs.

3. Why does Cliqueformer suffer less from distribution shift problems? Can adding an additional conservative regularization benefit the model?

Cliqueformer has an auto-encoder, as a sub-module, which produces representations $z(x)$. The representations are regularized with a KL divergence loss, so that a distribution of such representations can correspond to a single design $x$ [1]. The reviewer can think of it as partitioning the $z$-space into pieces which correspond to designs $x$. Note that we optimize our designs by optimizing representations $z$ that are later decoder to actual designs $x$. A helpful property is that, while a single gradient step cen create an invalid design, it is difficult to move from a valid-design region of $z$-space to an invalid region. In order to stay extra catious, though, we use weight decay while optimizing $z$: that is, we bring $z$ closer to $0$ with every gradient step. The reason for that is that, under standard Gaussian priors, the region of $z$-space near $0$ corresponds to valid designs, thus preventing the distribution shift. We hope that this helped. If not, would the reviewer be so kind to ask a follow-up question?

[1] D. Kingma and M. Welling 2013; Auto-encoding Variational Bayes

Questions

1. Was the transformer backbone the main source of improvement?

Thank you for pointing this out. To answer this question unequivocally, in the revised version, we evaluated gradient ascent with the transformer backbone. It achieves only slightly better performance than gradient ascent with MLP, and thus the answer to the reviewer's question is no. Similarly, as we found in Section 5.2, appropriate setting of the clique decomposition has big impact, even though the transformer backbone is employed for all runs. However, as we explained in lines 257-260 of the original submission, as well as we responded to Reviewer XTE4, the expressive transformer backbone is essential to enable learning representations that satisfy our desiderata.

2. Have the authors provided evidence that the proposed FGM compoenents were effective at exploiting the structure of the target black-box function?

Yes, we have. In addition to strong performance of the method as a whole, we swept over different configurations of the decomposition in Section 5.2. In each of the three sweeps, we covered the case $N_{clique}=1$, which corresponds to lack of the decomposition. In each of the tasks, this configuration was sub-optimal, loosing to configurations with many cliques.

References

[1] A. Krizhevsky et al. 2012; ImageNet Classification with Deep Convolutional Neural Networks

[2] A. Vaswani et al. 2017; Attention Is All You Need

[3] J. Ho et al. 2020; Denoising Diffusion Probabilistic Models

[4] Y. Song et al. 2021; Score-Based Generative Modeling Through Stochastic Differential Equations

[5] J. Kaplan et al. 2020; Scaling Laws For Neural Language Models

[6] S. Krishnamoorthy et al. 2023; Diffusion Models for Black-Box Optimization

Critique

1. The paper does not follow the proper scientific steps of (1) identifying a narrow problem in prior art, (2) proving its severity, and (3) proposing a solution to it. Is is, instead, focused on proposing a solution.

The approach the reviewer delineated sounds perfectly reasonable and we are sorry that the reviewer felt forced to give such a comment. However, we must contest it. The path that the reviewer offered is not the only valid way to conduct research. In fact, much of improvement in deep learning comes from projects that started by approaching a considered task from a new angle. Examples include AlexNet [1], Transformer [2], Diffusion Models [3], etc. Nevertheless, our project did tackle a specific problem: most of the research in MBO focused on algorithmic tenchiques, employing simple network architectures, to solve design tasks. Our goal was to show that such intricate algorithms are not necessary with careful adoption of the most performant deep learning architectures. And indeed, with Cliqueformer, we were able to optimize our designs with vanilla AdamW optimizer.

2. The experiments focus on the final performance of the model. Specific components of it lack empirical justification.

To begin with, we agree that we did not include experimental evidence for every compoenent of our architecture. However, in Section 3, we did propose reasoning that invokes evidence from prior work and our theoretical results. Note that Desideratum 1 is supported by (prior) Theorem 1 and our Theorem 2. Desideratum 2 also follows from Theorem 1, as well as it is intuitively clear: optimizing designs with respect to a model that was trained on too narrow of a set of designs is futile; the only question is if we can be less aggressive in enforcing the coverage requirements. Theorem 1 answers this question affirmatively, and ablations in Figure 5 provide empirical evidence for it. Additionally, in the revised version we add a Transformer baseline (without the regularized latent space and the clique decomposition) to our experiments and obtain largely inferior performance. We believe it provides additional support for our architecture.

We understand the reviewer's uncertainty about the empirical evidence for all components of Cliqueformer. Similarly to the previous question, however, we want to highlight that introducing a new approach to a considered task makes authors spend more time describing the approach - the remaining questions show we could spend even more time on that - and thus less time to properly study individual pieces. We want to be very honest: studying each step in isolation would be impossible here - see for example Line 47 in models/cliqueformer.py, where we state the best order of layers used in the implemented step. Studying each such piece properly would easily make us run out of space. Note that, for example, the Transformer paper [2] does not study much how important the order of layers is in a single block, or the attention head dimension is ([4], Fig 5, later showed that it is not that important). Also, the Diffusion Models paper [3] introduces many components, like the variance schedule, that it gets to work without a deeper analysis - this arrived later in follow-up works [5]. Similarly, we believe that a deeper empirical and theoretical analysis of Cliqueformer-like methods is an exciting avenue of future work.

3. Did the method actually result in better coverage of the clique data?

If we understood the question correctly, it actually asks about the effectiveness of VIBs/VAEs. To provide an instance-level piece of evidence, we ran an empirical study of the distribution of design representations in the revised version (Appendix B). The results indicate that the latent space did attain good coverage.

4. The current experimental breadth is awfully inadequate and the authors should run more baselines, e.g. DynaPPO, PEX, etc.

We are sorry that the reviewer feels so strongly about our experiments. Nevertheless, they are right that we shall update our range of baselines. In the revised version, we benchmarked a recent diffusion-based DDOM [6], and gradient ascent with the transformer backbone. When it comes to the methods suggested by the reviewer, they make us wonder if the problem setting of our paper is clear. Our work focuses on offline model-based optimization. Instead, the proposed methods come from the literature on protein/DNA-sequence design specifically. Crucially, they don't assume the offline scenario either - but the opposite online one. For example, in DynaPPO, a reinforcement learning-like ability to recollect new data is assumed (see Line 7 in Algorithm 1 of the paper). Also, PEX literally stands for proximal exploration, and contains wet-lab experiments as sub-routine (Algorithm 1). Could the reviewer re-consider our grade from the perspective of general-purpose offline MBO?

Dear Reviewer,

In addition to the rebuttal we responded to you above with, we added a new section More Related Work in Appendix D, in which we described the position of our paper with respect to works you have cited, such as DynaPPO and PEX. You have hinted a possibility of raising your score, given a strong rebuttal. As of now, we have run additional baseline experiments, including an ablative Transformer baseline, conducted extra experimental investigation in Appendix, upgraded our figures, responded to your questions, and added related work in the appendix. We would, therefore, greatly appreciate your updated feedback.

Best wishes,

Authors

Critique

1. How would you choose hyper-parameters for a new problem? More systematic study of the hyper-parameter configuration would be of value.

In our experiments, we found the following scheme to work generally well: (1) for problems with $dim(x) < 100$: we fix $d_{clique}=3$ and $d_{knot}=1$, and choose $N_{clique}$ so that $dim(z)=N_{clique}\cdot (d_{clique} - d_{knot}) + d_{knot} \approx dim(x)$; (2) for problems with $dim(x) \geq 100$, we fix $d_{clique}=6$ and repeat the procedure, to decrease the computational cost of the attention layer, which is quadratic in $N_{clique}$. We reported the results of these scheme in Table 1 with the exception of Superconductor, where a sweep over these values brought a substantial improvement. However, the result for $d_{clique}=3$ (and $N_{clique}=40$) remains in Figure 5 (b). Would the reviewer like us to put it in the table?

We understand that a systematic study of these hyper-parameters would be helpful. However, this paper introduces a new transformer architecture that is meant to enable optimization of designs with vanilla optimization techniques, like AdamW (without need for CbAS or conditional generation). We thus spent the majority of time describing the novel model, leaving little room for a systematic hyper-parameter study, and thus providing only the esssential ablation. In fact, this is quite normal for works that introduce new models: AlexNet has not ablated the dropout rate or the momentum parameter [1]; Transformer did not analyze the importance of the attention head dimension per parameter count [2] - such a study came years later [3]. We do consider such analysis of our hyper-parameters, both theoretically and empirically, an exciting avenue of future work.

2. The spectrum of baselines is limited. More methods must be compared against to claim the state-of-the-art performance.

We agree with the reviewer. We have added results of the more recent SOTA diffusion-based method DDOM [4], as well as of Gradient Ascent with the transformer backbone,

3. The authors miss the discussion of the computational complexity of Cliqueformer.

We understandand if the reviewer would like to see computational complexity analysis, but we do not think that lack of it is a unique flaw of our paper. Most deep learning papers do not offer such analysis. Diffusion models [5] were introduced with a U-Net architecture that added an attention layer at at least one level of resolution, but the computational complexity was omitted. COMs [6] performs 50 adversarial steps for every update (see Line 4 in their Algorithm 1) and does not provide that analysis either. Cliqueformer relies on the transformer backbone, so it's computational complexity is that of the transformer. We do not think there is anything unusual about it.

4. The authors do not provide theoretical guarantees for their method.

While we agree that, ideally, all computational methods had theoretical guarantees. However, modern deep learning - which we rely on - lacks mathematical apparatus to provide many practical methods with such guarantees together. This is why many performant works do not have them. Examples include: AlexNet [1] from computer vision, Transformer [2] from language modeling, DDOM from MBO [4], PPO [7] from reinforcement learning, StyleGAN-XL [8] from generative modeling, etc. The first of the cited by the reviewer papers, BootGen, does not provide such guarantees either.

Questions

1. How to choose the FGM configuration for a new problem?

We recommend the procedure we gave in the answer to the first critique. Additionally, we believe that, in practical applications, a method similar to how MBO research is conducted would be helpful in selecting the hyper-parameters. That is, for a given dataset, a surrogate oracle $\hat{y}(x)$ can be trained on the real-world labels $y(x)$. Then, the labels $y$ can be replaced with predictions $\hat{y}$ made by the surrogate. One can then train an MBO algorithm with respect to these new labels and see which hyper-parameters perform best.

References

[1] A. Krizhevsky et al. 2012; ImageNet Classification with Deep Convolutional Neural Networks

[2] A. Vaswani et al. 2017; Attention Is All You Need

[3] J. Kaplan et al. 2020; Scaling Laws For Neural Language Models

[4] S. Krishnamoorthy et al. 2023; Diffusion Models for Black-Box Optimization

[5] J. Ho et al. 2020; Denoising Diffusion Probabilistic Models

[6] B. Trabucco et al. 2021; Conservative Objective Models for Effective Offline Model-Based Optimization

[7] J. Schulman et al. 2017; Proximal Policy Optimization Algorithms

[8] A. Sauer et al. 2022; StyleGAN-XL: Scaling StyleGAN to Large Diverse Datasets