Proteina: Scaling Flow-based Protein Structure Generative Models

Q2. Fold class classifier and its use in the new metrics.

No, there is no unfair advantage. Our fold class classifier is trained purely for the purpose of evaluation of all models and never used during inference (for sampling we use classifier-free guidance).

The experimental setup is analogous to the situation in image generation, where an Inception network-based ImageNet classifier is used universally across the literature to evaluate model performance and calculate FID and IS scores, including for class-conditional generative models trained on ImageNet. We would also like to point out that most of our evaluations (Table 1) are on unconditional generation anyway. We now split off conditional generation in the updated version of the manuscript (Table 2) and evaluated fold class-conditioned Chroma (which uses its own, separate classifier for classifier guidance during inference) as another baseline. Chroma conditions on the same labels as Proteina, thereby making the comparison fair.

It may be possible to train other classifiers or embedding networks for the FPSD metric. However, we leave this to future research. As discussed in the paper, leveraging the fold class labels to train the classifier necessarily results in a classifier sensitive to different protein structures, which is what we require to evaluate our protein structure generative model.

Q3. Do all residues of protein backbones generated with fold class conditioning use the same CAT label?

Yes, all residues use the same CAT label, this is, each backbone has one global label that defines its overall global fold class. Future work could explore generation with multiple fold classes at the same time, for instance in a multimer generation setting.

We hope that we were able to answer all your questions. We would also like to point out our additional new results that we added to the paper; please see the reply to all reviewers.

We would like to thank the reviewer for their very positive review. We are grateful for highlighting our state-of-the-art performance, our new evaluation metrics, our improved $\beta$-sheet generation thanks to the fold class conditioning, as well as our other contributions such as LoRA fine-tuning and autoguidance. Below, we address the raised questions:

1. Code and data availability for reproducibility.

Proteina will be made available publicly after publication. Also note that all data sources are publicly available under permissive licenses, and the data preparation protocol is described in detail in Appendix B. Moreover, all training, sampling and hyperparameter details are covered in Appendices F and K.

2. Training on 21 million protein structures, model size, and scaling trends.

Training and fully evaluating new models on the large D_21M dataset is beyond the scope of the rebuttal. However, we would like to point out that our motivation to train this M_21M model is to show a proof of concept that one can curate very large-scale high-quality training sets from the AFDB that result in extremely high designability, despite the fully synthetic nature of the AFDB database, while leveraging scalable architectures like ours (previous training sets are more than an order of magnitude smaller, see Sec. 3.1). Future work could build on this observation and curate different large-scale training sets from the AFDB and other sources, using different filters, tailored to different needs.

Furthermore, in the updated version we performed a scaling study with respect to the flow matching objective and model size, similar to recent seminal work in image generation [1]. See Figures 19 in Appendix M.1. We see that scaling the model size systematically reduces the loss, implying a strong scaling behavior of our architecture.

Moreover, we added the number of parameters as part of Tables 9, 10 and 11 in Appendix M.2. We list the number of parameters for all our models as well as the most relevant baselines. Our Proteina models may seem large and expensive by parameter count, but what matters in practice is not an absolute parameter count, but a model’s memory consumption and sampling speed, along with performance (which is reported in the main paper already, in Table 1). Therefore, in Figure 20 and Tables 9, 10 and 11, we also show the maximum batch size per GPU that our models and the most important baselines can achieve during inference, as well as sampling speed, both in single sample mode and in batched mode, normalized per sample. We see that while our models employ many parameters, they remain highly memory efficient and are also still as fast as or faster than the baselines. This is due to our efficient non-equivariant transformer architecture. Note that for this table, we also trained an additional smaller model with around 60M parameters, a model size more similar to the baselines – this model is particularly efficient, yet still performs very well (see Table 8). Please see the entire discussion in the newly added Appendix M.2.

Q1. Input sequence representation.

The input sequence representation (see architecture Fig. 6) only consists of the embedded backbone coordinates, as well as a sequence index and an optional self-conditioning input. We call this sequence representation, because the embedded coordinates come in sequential order, and this naming is also consistent with the transformer literature. However, we are not modeling any amino acids and our sequence representation does not correspond to an amino acid sequence, and there are also no special symbols. Proteina focuses on modeling protein backbone structures.

Thank you to the author for their response. Regarding the open-sourcing of code and data, a more proactive approach should be adopted, such as sharing via anonymous URLs. Additionally, research on scaling should include studies on data scaling, which can be easily achieved by reducing the data volume. Overall, the author’s response addresses some concerns, but I still have lingering doubts. I hope the author can address my concerns more proactively, even though I initially gave a high score.

6. Statistical significance and sensitivity in the metric.

In the newly added Appendix section M.5 (see Table 14), we repeatedly evaluate one of our models with different seeds, observing very little variation of the results, thereby demonstrating that our results can be considered statistically significant. Estimating the exact standard deviations for all models and baselines through repeated evaluation is unfortunately computationally infeasible due to the massive amount of sampling required and hence beyond the scope of the rebuttal. However, we can expect similarly accurate results for all other models and baselines as for the one for which we did the error estimation here.

We would also like to point out that all metrics vary across our different models and the baselines, thereby indicating that all metrics are sensitive to the different generated samples.

7. The code is not provided.

Proteina will be made available publicly after publication. Also note that all data sources are publicly available under permissive licenses, and the data preparation protocol is described in detail in Appendix B. Moreover, all training, sampling and hyperparameter details are covered in Appendices F and K.

Q1. The quality of a fold class predictor can be evaluated through classifier-guidance.

It is not clear to us what exactly is meant with this question. However, we would like to point out that our fold class predictor is validated by evaluation on a held-out test set (see Appendix E.3). Further, whether our fold class-conditioned Proteina models accurately generate proteins from the correct fold classes is then evaluated by re-classifying the generated backbones with the fold class predictor and seeing if the predictions match the input conditioning (see Appendix H). We would like to kindly ask the reviewer to clarify their question, in case our reply is not what they were looking for.

Q2. GearNet-based classifier details.

We trained the GearNet fold class predictor from scratch, without contrastive pretraining. We leverage the GearNet architecture, but trained the model in a fully supervised fashion using the CAT labels. Note that the original GearNet architecture processes both sequence and structure data; we modify GearNet to focus solely on predicting fold classes based on structure. Please see Appendix E.3, where the fold class predictor training is explained in detail.

Q3. Use of inception score and exponentiation.

Our definition of the fold Score (fS) follows the standard definition of the inception score in computer vision [2], which includes the exponentiation of the Kullback-Leibler term, see definition in Sec. 3.5. This exponentiation ensures that the scale of the score is in ~[1, N], where N is the number of classes, making the metric more interpretable. Although indeed less popular, the inception score nonetheless remains a relevant metric to quantify a generative model’s coverage of a class distribution. This is particularly interesting when quantifying a class-conditional model like ours, which builds on classifier-free guidance. In fact, the inception score, together with the FID, was used in the original classifier-free guidance paper to analyze the results [3]. Since we similarly use classifier-free guidance (for the first time) for protein generation, we believe it is similarly a good idea to study our fold Score (fS), our protein backbone version of the inception score.

We hope we were able to address your questions and concerns. In that case, we would like to kindly ask you to consider raising your score accordingly. Otherwise, please let us know if there is anything else we can clarify. Finally, please see the reply to all reviewers for an overview over all additional new results that we added to the paper.

[1] Ingraham et al., Illuminating protein space with a programmable generative model. Nature 2023.

[2] Salimans et al., Improved Techniques for Training GANs. NeurIPS, 2016.

[3] Ho and Salimans, Classifier-Free Diffusion Guidance. arXiv, 2022.

3. Use of different training datasets and reference datasets in new metrics.

In fact, not only RFDiffusion, ESM3 and Genie2 were trained on different datasets, but all baselines use different training sets (those works training on the PDB all use differently filtered PDB subsets). In the protein structure generation literature, the choice of the dataset is treated as a hyperparameter itself. This is because from the protein designer perspective all that matters is that the model can produce diverse, novel and designable proteins, while a protein designer is not interested in the training dataset (hence, a better TM-score-based diversity is generally desirable, irrespective of the training data). Considering that all baselines trained on different data, it is not feasible to re-train all models on the same data.

That being said, our M_FS and M_FS^no-tri models are trained on exactly the same AFDB subset as Genie2, which is the strongest among the baselines. As pointed out above, we do outperform Genie2 in designability and diversity. Regarding the question whether PDB/AFDB data was part of training, please see Sec. 3. In short, our model employs two different AFDB subsets. PDB training data is only leveraged in the LoRA fine-tuning experiment.

Generally, the PDB (natural proteins) and the AFDB (synthetic proteins predicted from AlphaFold2) represent the two main protein structure databases, from which different works and models construct their training datasets. Therefore, we chose a representative subset from the PDB as well as a representative subset from the AFDB as our reference distributions in the FPSD and fJSD metrics. From the protein design perspective, this allows one to study how close in distribution a model’s generated proteins are to either the synthetic proteins from the AFDB or the natural proteins from the PDB, which provides valuable insights. We do agree with the reviewer that the value of the metrics always needs to be considered with a model’s training data in mind. That is another reason why we now removed all bold/underline highlighting for the new metrics (see above). We kept the “higher/lower is better” arrow indications, to indicate general trends; this is, a lower FPSD/fJSD indicates a closer match to the reference distribution and a higher fS indicates more diverse samples.

4. Competitors for fold-class conditional generation.

RFDiffusion can be conditioned on pairwise distance maps between residues, which can be used to also specify certain folds in a fine-grained manner, but RFDiffusion cannot be conditioned on the broad and universal CAT fold class labels that we leverage and that provide a semantic and hierarchical way to control the backbone generation at different levels of detail. Therefore, a meaningful comparison with RFDiffusion is not possible. However, Chroma [1] can indeed be used with CAT fold class conditioning, leveraging classifier guidance (in contrast to our classifier-free guidance). Hence, in the updated version we evaluated fold class-conditioned Chroma as a new baseline, which we outperform (discussed already above in point 1).

5. TM-score-based diversity and violin plots.

As discussed above, from a protein designer’s perspective a more diverse model is generally desirable, implying that a lower TM-score-based diversity value is generally better. That said, we agree with the reviewer that it is helpful to study diversity in the context of the diversity observed in the PDB and AFDB, from which the models construct their different training sets. In Table 13 in Appendix M.4, we now report all metrics when evaluated on our PDB and AFDB reference sets used. We added a discussion in the appendix, but would like to highlight the TM-score results here; we find that the data generally has a higher diversity (lower TM-score) than all models and baselines, which is expected and implying that existing models can still potentially be optimized further for diversity.

We also agree that it can be helpful to better analyze the distribution over TM-scores. Following the reviewers’ suggestion, we added violin plots for the TM-score for different lengths for our main models as well as the most relevant baselines (Genie 2, RFDiffusion, FrameFlow), see Figure 21 in Appendix M.3. We find that all models show reasonable TM-score distributions without mode collapse.

We would like to thank the reviewer for their comprehensive review. We appreciate that the reviewer considers our work very well-written and easy to follow, and we thank the reviewer for highlighting Proteina’s successful large-scale non-equivariant flow-based generation of protein structures. Below, we address the raised questions:

1. Performance of Proteina compared to previous work and additional fold class-conditioned baseline.

We would like to point out that we do indeed outperform all baselines in terms of designability and diversity in the unconditional generation benchmark in Table 1. Different models and sampling parameters (noise scale gamma) show different tradeoffs, but for instance the M_FS (gamma=0.45) model achieves a designability of 96.4% as well as a cluster-based diversity of 0.63 (305) and a TM-score diversity of 0.36. No baseline scores as well. FrameFlow has similar TM-score diversity, but much worse designability and cluster-based diversity. FoldFlow can achieve slightly higher designability, but has much worse diversity, and both our M_FS (gamma=0.35) model as well as our M_21M (gamma=0.3) model outperform FoldFlow not only in designability, but also diversity. Increasing the noise scale to gamma=0.5 for the M_FS model further boosts diversity, far beyond all baselines, while maintaining above 90% designability. Meanwhile, our M_21M (gamma=0.3) model achieves 99% designability, also far beyond all baselines. In novelty, we are only slightly behind Genie2, which we outperform in designability and diversity.

Moreover, in the long chain generation experiments (see Sec. 4.2), we outperform all baselines by a large margin.

Regarding fold class-conditional generation, there is only one previous work, Chroma [1], which qualitatively demonstrated fold class conditioning. We now also systematically sampled Chroma with fold-class conditioning and compared with our results in the updated version of the paper. We find that we significantly outperform Chroma in designability and TM-score diversity, and we also generate more designable clusters (see Table 2). Further, we also outperform Chroma in re-classification accuracy, which validates whether specific classes are generated correctly, see Table 6 in Appendix I.1. Chroma fails to reliably generate protein structures according to input query fold class labels. We would like to stress that our work is the first that quantitatively studies fold class conditioning, opening the door towards more advanced control in protein backbone generation and new applications. Note that we have split Table 1 into Table 1 and Table 2 in the updated version to better emphasize that the conditional and unconditional settings are not directly comparable.

Considering the findings in these different experiments, we do believe that we overall indeed outperform relevant previous works, achieving state-of-the-art performance. To be more precise, we slightly rephrased our claim in the last paragraph of the introduction to “(v) We achieve state-of-the-art designable and diverse protein backbone generation performance …”.

2. Role of the new metrics (FPSD, FS, fJSD).

In practice, a protein designer is interested in diverse, novel and designable protein backbones. Hence, these are the primary metrics that should be optimized. With this in mind, our new metrics (FPSD, fS, fJSD) are not meant to represent “yet another benchmark metric to beat”. Instead, they can provide novel and detailed insights, help analyze results, and guide model development. They are therefore complementary to the other, established metrics. For instance, as discussed in section 4.1, they have been useful to better understand the results on full distribution and fold class-conditional generation. To emphasize that, although the new metrics can bring new insights, they should not be the target of optimization and used to rank models, we have now removed all bold and underline highlighting from the tables for these metrics. We have also added a brief discussion on this at the end of section 3. Our claims regarding state-of-the-art designable and diverse protein structure generation performance are not related to these new metrics, but based on the results discussed above.

Regarding the scale of the metrics, we would like to emphasize that we are primarily interested in relative comparisons between the methods, not absolute values. We could easily re-scale the metrics to vary in any desired interval, but this would not lead to additional insights.

3. Use of different training datasets and reference datasets in new metrics.

In fact, all models and baselines in Table 2 are trained on different datasets. Some models use subsets of the AFDB, like Genie 2 and Proteina, while other models leverage subsets of the PDB. Those works using PDB proteins all use different filtering strategies, resulting in different training sets and results with different designability, diversity and novelty tradeoffs.

The crucial point is that in the protein structure generation literature, the choice of the dataset is treated as a hyperparameter itself. This is because from a protein designer’s perspective all that matters is that the model can produce diverse, novel and designable proteins, while a protein designer is less interested in the employed training dataset. Considering that all baselines trained on different data, it is not feasible to re-train all models on the same data.

That being said, our M_FS and M_FS^no-tri models are trained on exactly the same AFDB subset as Genie2, which is the strongest among the baselines. As pointed out above, we do outperform Genie2 in designability and diversity.

Generally, the PDB (natural proteins) and the AFDB (synthetic proteins predicted from AlphaFold2) represent the two main protein structure databases, from which different works and models construct their training datasets. Therefore, we chose a representative subset from the PDB as well as a representative subset from the AFDB as our reference distributions in the FPSD, fJSD and novelty metrics. For novelty, this practice simply follows Genie2, with which we wanted to be consistent. From the protein design perspective, this allows one to study how close in distribution a model’s generated proteins are to either the synthetic proteins from the AFDB or the natural proteins from the PDB, which provides valuable insights. We do agree with the reviewer that the value of the metrics always needs to be considered with a model’s training data in mind. That is another reason why we now removed all bold/underline highlighting for the new metrics (see above). We kept the “higher/lower is better” arrow indications, to indicate general trends; this is, a lower FPSD/fJSD indicates a closer match to the reference distribution and a higher fS indicates more diverse samples.

That aside, the training data for our M_FS and M_FS^no-tri models and for Genie2 is in fact the same as the AFDB reference data used in the different metrics (hence, lower FPSD and fJSD scores indeed indicate a better match of the training data in those cases).

4. TM-score-based diversity, violin plots, and different clustering thresholds for cluster-based diversity.

As discussed above, from a protein designer’s perspective a more diverse model is generally desirable, implying that a lower TM-score-based diversity value is generally better. That said, we agree with the reviewer that it is helpful to study diversity in the context of the diversity observed in the PDB and AFDB, from which the models construct their different training sets. In Table 13 in Appendix M.4, we now report all metrics when evaluated on our PDB and AFDB reference sets used. We added a discussion in the appendix, but would like to highlight the TM-score results here; we find that the data generally has a higher diversity (lower TM-score) than all models and baselines, which is expected and implies that existing models can potentially still be optimized further for diversity.

We also agree that it can be helpful to better analyze the distribution over TM-scores. Following the reviewers’ suggestion, we added violin plots for the TM-score for different lengths for our main models as well as the most relevant baselines (Genie 2, RFDiffusion, FrameFlow), see Figure 21 in Appendix M.3. We find that all models show reasonable TM-score distributions without mode collapse.

We further follow the reviewer’s suggestion and report cluster-based diversity with different thresholds for three of our “best” models as well as the most relevant Genie2, RFDiffusion and FrameFlow baselines (see Table 12 in Appendix M.3). We find that for looser thresholds our M_FS (gamma=0.45) model generally “wins”, while when considering very strict clustering thresholds, all models and baselines produce highly diverse results with many different clusters.

We would like to thank the reviewer for their comprehensive review. We appreciate that the reviewer considers our work exceptionally well-written, we thank the reviewer for pointing out that our work pushes the frontier in scaling protein structure generation, and we are grateful for highlighting our extensive experiments and ablations. Below, we address the raised questions:

1. Performance of Proteina compared to previous work, experimental design, task scope and additional fold class-conditioned baseline.

We would like to point out that we do indeed outperform all baselines in terms of designability and diversity in the unconditional generation benchmark in Table 1. Different models and sampling parameters (noise scale gamma) show different tradeoffs, but for instance the M_FS (gamma=0.45) model achieves a designability of 96.4% as well as a cluster-based diversity of 0.63 (305) and a TM-score diversity of 0.36. No baseline scores as well. FrameFlow has similar TM-score diversity, but much worse designability and cluster-based diversity. FoldFlow can achieve slightly higher designability, but has much worse diversity, and both our M_FS (gamma=0.35) model as well as our M_21M (gamma=0.3) model outperform FoldFlow not only in designability, but also diversity. Increasing the noise scale to gamma=0.5 for the M_FS model further boosts diversity, far beyond all baselines, while maintaining above 90% designability. Meanwhile, our M_21M (gamma=0.3) model achieves 99% designability, also far beyond all baselines. In novelty, we are only slightly behind Genie2, which we outperform in designability and diversity.

Moreover, in long chain generation (see Sec. 4.2), we outperform all baselines by a large margin.

Regarding fold class-conditional generation, there is only one previous work, Chroma [1], which qualitatively demonstrated fold class conditioning. We now also systematically sampled Chroma with fold-class conditioning and compared with our results in the updated version of the paper. We find that we significantly outperform Chroma in designability and TM-score diversity, and we also generate more designable clusters (see Table 2). Further, we also outperform Chroma in re-classification accuracy, which validates whether specific classes are generated correctly, see Table 6 in Appendix I.1. Chroma fails to reliably generate protein structures according to input query fold class labels. We would like to stress that our work is the first that quantitatively studies fold class conditioning, opening the door towards more advanced control in protein backbone generation and new applications. Note that we have split Table 1 into Table 1 and Table 2 in the updated version to better emphasize that the conditional and unconditional settings are not directly comparable.

Considering the findings in these different experiments, we do believe that we overall indeed outperform relevant previous works, achieving state-of-the-art performance. To be more precise, we slightly rephrased our claim in the last paragraph of the introduction to “(v) We achieve state-of-the-art designable and diverse protein backbone generation performance …”.

2. Role of the new metrics (FPSD, FS, fJSD).

In practice, a protein designer is interested in diverse, novel and designable protein backbones. Hence, these are the primary metrics that should be optimized. With this in mind, our new metrics (FPSD, fS, fJSD), which also involve our fold class classifier, are not meant to represent “yet another benchmark metric to beat”. Instead, they can provide novel and detailed insights, help analyze results, and guide model development. They are therefore complementary to the other, established metrics. For instance, as discussed in section 4.1, they have been useful to better understand the results on full distribution and fold class-conditional generation. To emphasize that, while the new metrics can bring new insights, they should not be the target of optimization and used to rank models, we have now removed all bold and underline highlighting from the tables for these metrics. We have also added a brief discussion on this at the end of section 3. Our claims regarding state-of-the-art designable and diverse protein structure generation performance are not related to these new metrics, but based on the results discussed above on the conventional metrics. Also, since we are the first to propose such new distribution-focused metrics, we naturally had to train our own classifier, which we rigorously validated and which does not represent any unfair advantage, see Appendix E.

5. Fold class classifier and noisy data.

Indeed, if a model generates noisy data, then the classifier may not perform well. This will be reflected in the classifier’s prediction and in its internal feature maps, thereby negatively influencing the FPSD, fJSD and fS metrics that rely on this classifier. However, this is desirable, because in that way the metrics can immediately catch noisy generated structures. This is what is validated in the experiments in Appendix E.4. This approach follows the practice from the image generation literature, where the inception-based classifier widely used in the FID [1] and IS [2] metrics is also not specifically trained for noise robustness, thereby leading to metrics sensitive to noisy data.

Finally, we agree that a classifier trained on noisy data could be used for classifier guidance. However, Proteina leverages classifier-free guidance for fold class conditioning and the classifier is only used for evaluation purposes, never during sample generation itself, which ensures that all evaluations are fair.

Q1. GearNet-based classifier details.

We trained the GearNet fold class predictor from scratch, without contrastive pretraining. We leverage the GearNet architecture, but trained the model in a fully supervised fashion using the CAT labels. Note that the original GearNet architecture processes both sequence and structure data; we modify GearNet to focus solely on predicting fold classes based on structure. Please see Appendix E.3, where the fold class predictor training is explained in detail.

We hope we were able to address your questions and concerns. In that case, we would like to kindly ask you to consider raising your score accordingly. Otherwise, please let us know if there is anything else we can clarify. Finally, please see the reply to all reviewers for an overview over all additional new results that we added to the paper.

[1] Heusel et al., GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium. NeurIPS, 2017.

[2] Salimans et al., Improved Techniques for Training GANs. NeurIPS, 2016.

I appreciate that authors have made several improvements to the manuscript. Adding Chroma to fold class conditioning comparison provides some context for conditional generation evaluation. I appreciate the plots of TM-score distributions and TM-score clustering at different thresholds. These experiments support the concern that mean pairwise TM-score is a rather non-informative metric. The provided clustering and visualization of the TM-score distributions reveal much more details about the generated samples, e.g. a small fraction of duplicate structures. Finally, I found the provided additional details on the training of the GearNet-based classifier important for reproducibility. It would be great to update the manuscript accordingly.

Still, several major methodological concerns remain:

1. Self-evaluation bias - using their own encoder/classifier and metrics without external validation raises concerns about objectivity. I acknowledge that many established protein structure encoders like ProteinMPNN or GVP-Transformer cannot be directly applied to CA-only backbones and use 4- and 3-atom-per-residue backbones. Nonetheless, the authors could validate their approach by comparing their encoder/classifier representations of protein backbones with those from established methods on tasks where both are applicable, or by demonstrating performance on independent classification tasks using a separate validation dataset.

2. The authors' argument that "dataset is a hyperparameter" sidesteps fundamental scientific reproducibility concerns. Furthermore, comparing distributional distances across models trained on different data remains problematic from a methodological perspective, as it confounds model improvements with dataset effects.

3. The proposed metrics (FPSD/fJSD) may inadvertently favor models trained on data more similar to reference sets. Without controlled comparisons, it becomes difficult to determine whether better scores come from improved modeling capabilities or simply from dataset similarity.

4. The apparent correlation between FPSD and fJSD metrics raises questions about their redundancy. This correlation could stem from either inherent metric similarity, similar performance patterns across evaluated models, or biases introduced by training and using the classifier on related datasets.

To address these concerns, I suggest the following improvements:

1.1. Use multiple established structure encoders alongside the self-trained encoder/classifier to provide independent validation

1.2. Through independent tasks demonstrate what unique structural information is captured by the self-trained encoder/classifier compared to established encoders

1.3. Add ablation studies using different encoders for the distributional metrics, carefully selecting the data to alleviate potential data leakage.

2. Add controlled experiments comparing models trained on the same datasets

3. Provide correlation analysis between the proposed metrics

4. (Minor) Include reference measurements between random splits of training data to establish scales for distributional metrics

Without solving problems with FPSD/fJSD/fS (need for independently validated encoder and inconclusive comparisons between models trained on different datasets), the generation benchmark boils down to one comparison in unconditional generation (with Genie2) and one in conditional (with Chroma) evaluated on Designability, Diversity and Novelty. Given the fair criticism of these metrics from the authors (Section 3.5), with which I completely agree, it becomes crucial to establish more rigorous evaluation protocols. The paper presents important technical advances in scaling protein structure generation, but without addressing these methodological concerns, the comparative results should be interpreted with appropriate caution. I encourage the authors to focus on strengthening the evaluation framework to match the impressive technical achievements demonstrated in the work.

Note that while developing the classifier, we did explore different architectures. In fact, we also experimented with a GVP architecture [1]. Both GearNet and GVP led to very similar, strong performance. Because GearNet is more efficient in this case, we chose it over GVP.

We would also like to stress that the classifier is purely used for evaluation purposes. We never use the classifier during generative model training or inference. Therefore, the use of the classifier is fair among all baselines and there is no “self-evaluation bias”. We would also like to stress that our practice exactly follows what is being done in the large literature on image generative modeling. The Frechet Inception Distance (FID) and the Inception Score (IS) for evaluation of image generators [2,3] are based on an inception network classifier, trained on ImageNet. FID and IS have been used in hundreds of works to evaluate generative models trained both on ImageNet itself, and other datasets. The way in which this inception classifier is used and how it is sensitive to global and local image and texture variations is exactly analogous to how our GearNet classifier is used and how it is sensitive to complex protein structure variations. We are following standard and well-established practices from the large image generation literature.

2. Models and baselines trained on different datasets. Reproducibility.

We respectfully disagree with the reviewer on this point. Training models on different datasets has nothing to do with reproducibility. All datasets used in our work are publicly and freely available to both academic and commercial entities and all data processing is explained in detail in the Appendix. Moreover, we now made a version of the code available to the reviewers (for reviewing purposes only), and we will publicly release code upon publication. We believe that we extensively address scientific reproducibility.

Another question is to which degree different models can be compared. We acknowledge that the training dataset does of course play an important role in a model’s performance. However, all works in protein structure generation do indeed use different datasets. In this literature, the dataset is indeed treated as a form of hyperparameter. This is because one looks at the problem from the perspective of a protein designer, who does not care about training datasets, but only about the ability of the models to generate designable, diverse and novel proteins — hence, models are compared on these aspects, irrespective of data and architecture details.

To be clear, let us give some examples:

• FoldFlow trained on 22,248 monomeric samples from the PDB, using samples with <5A resolution, <50% loops, sampling uniformly over clusters with 30% sequence similarity.
• For FrameFlow, there exist two versions, a stronger one trained on a somewhat similar PDB subset as FoldFlow, and another one trained on a smaller subset of just 3,938 backbones (SCOPe data).
• Proteus trained on 50,773 structures from the PDB, also using oligomers split into individual chains.
• Genie 2 trained on 588,570 structures from the AFDB, the same as our D_FS.
• Chroma is trained on 28,819 protein structures. These contain non-membrane, X-ray structures with resolution <2.8A from the PDB, clustered at 50% sequence identity, taking one sample from each cluster. The data is further enriched with 1726 non-redundant antibodies clustered at 90% sequence identity.
• RFDiffusion first trains RoseTTAFold, a folding model, additionally leveraging sequence information, in contrast to all other models. Only then it is fine-tuned for generation without sequences. It uses 25,600 samples from the PDB, and the exact filtering and clustering details are unclear.

The chain lengths the models are trained on vary, too. We see that all previous works use different datasets with different preprocessing, filtering, clustering or pre-training steps, all of which matter. And they also all use different architectures. Nonetheless, it is common practice in the literature to compare these models with each other, as has been done in all previous papers in the area (which we cite in our related work section). We would like to acknowledge that we sympathize with the reviewer’s point, though. It could indeed be helpful if the field of protein structure generation established more well-defined benchmarks. However, this is beyond the scope of our paper, and we follow the current practices of the field in our work. In fact, it is one of our key contributions to demonstrate successful training on large-scale protein backbone datasets using efficient and scalable transformer architectures.

We would like to thank the reviewer for their comprehensive review. We appreciate that the reviewer considers our work well-written and enjoyed reading it. We would also like to thank the reviewer for highlighting our various interesting explorations and evaluations, pointing out the impact on future work. Below, we address the raised questions:

1. Scaling Laws.

In our case of flow matching with Gaussian conditional probability paths, the flow matching loss can be seen as a rescaled version of the evidence lower bound (see [1,2,3] for details), thereby being closely related to the likelihood. Similar to recent seminal work in image generation [4], which analyzed the scaling of the flow matching objective, we therefore also directly analyze the flow matching loss for differently-sized Proteina transformer models. In the newly added Fig. 19 in Appendix M.1, we see that scaling model size indeed systematically reduces the loss, implying a strong scaling behavior of our architecture.

2. Models in Table 1, and number of parameters.

We trained three separate models in our paper, M_FS, M_FS^no-tri, and M_21M, see first paragraph of Sec. 4. The “FS/21M” subscript denotes which data the model was trained on (see Sec. 3). “No-tri” indicates that no triangular multiplicate layers were used (see discussion in Sec. 3.3). Each model is trained such that it can be sampled both fold-class conditioned, or in an unconditional manner (see Sec. 3.2). Moreover, varying the noise scale $\gamma$ during inference also leads to different results. These aspects result in different evaluation setups, presented in the rows of Tables 1 and 2.

To not make these tables even larger, we added the number of parameters as part of the Tables 9, 10 and 11 in Appendix M.2. We list the number of parameters for all our models as well as the most relevant baselines. Our Proteina models may seem large and expensive by parameter count, but what matters in practice is not an absolute parameter count, but a model’s memory consumption and sampling speed, along with performance (which is reported in the main paper already, in Table 1). Therefore, in Figure 20 and Tables 9, 10 and 11, we also show the maximum batch size per GPU that our models and the most important baselines can achieve during inference, as well as sampling speed, both in single sample mode and in batched mode, normalized per sample. We see that while our models employ many parameters, they remain highly memory efficient and are also still as fast as or faster than the baselines. This is due to our efficient non-equivariant transformer architecture. Note that for this table, we also trained an additional smaller model with around 60M parameters, a model size more similar to the baselines – this model is particularly efficient, yet still performs very well (see Table 8). Please see the entire discussion in the newly added Appendix M.2.

3. Performance of the M_21M model.

The reviewer also pointed to the performance of the M_21M model. The fact that this model exhibits a different performance tradeoff compared to our other models and several baselines is not because of the model architecture. It is because it was trained on the D_21M dataset, which is subject to careful filtering to only contain high-quality well-structured backbones (see Sec. 3.1). Hence, this leads to slightly lower diversity, but state-of-the-art 99% designability (see Tab. 1). Our motivation to train this model is to show a proof of concept that one can curate such very large-scale high-quality training sets from the AFDB that result in extremely high designability, despite the fully synthetic nature of the AFDB database, while leveraging scalable architectures like ours (previous training sets are more than an order of magnitude smaller, see Sec. 3.1). Future work could build on this observation and curate different large-scale training sets from the AFDB and other sources, using different filters, tailored to different needs.

We hope we were able to address your questions and concerns. In that case, we would like to kindly ask you to consider raising your score accordingly. Otherwise, please let us know if there is anything else we can clarify.

Finally, we would also like to point out our other additional new results that we added to the paper; please see the reply to all reviewers.

[1] Ho et al., Denoising Diffusion Probabilistic Models. NeurIPS, 2020.

[2] Albergo et al., Stochastic Interpolants: A Unifying Framework for Flows and Diffusions. arXiv, 2023.

[3] Kingma and Gao, Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation. NeurIPS, 2023.

[4] Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis. ICML 2024.

We are glad that we were able to address your concerns and are happy to also provide some insights regarding your follow-up question.

I am curious that considering the performance only, would flow matching-based approaches outperform the diffusion-based approaches over protein backbone generation? If so, it is there any intuitions or insights behind that?

To answer this question, we would like to first point out that we are using flow matching to couple the training data distribution (the protein backbones for training) with a Gaussian noise distribution, from which the generation process is initialized when sampling new protein backbones after training. In this case, i.e. when coupling with a Gaussian distribution, flow matching models and diffusion models can actually be shown to be equivalent, up to reparametrizations. This is because diffusion models generally use a Gaussian diffusion process, thereby also defining Gaussian conditional probability paths, similar to the Gaussian conditional probability paths in flow matching with a Gaussian noise distribution.

For instance, when using a Gaussian noise distribution, one can rewrite the velocity prediction objective used in flow matching as a noise prediction objective, which is frequently encountered in diffusion models. Different noise schedules in diffusion models can be related to different time variable reparametrizations in flow models. Most importantly, for Gaussian flow matching, we can derive a relationship between the score of the interpolated distributions $\nabla_{\mathbf{x}_t}\log p_t(\mathbf{x}_t)$ and the flow’s velocity (see Eq. (3) in our paper as well as Appendix F.1). The score is the key quantity in score-based diffusion models. Using this relation, diffusion-like stochastic samplers for flow models can be derived, as well as flow-like deterministic ODE samplers for diffusion models (see below). In conclusion, we could in theory equally consider our flow models score-based diffusion models. With that in mind, to answer your question, from a pure performance perspective when coupling with a Gaussian noise distribution flow matching-based approaches and diffusion-based approaches should in principle perform similarly well. In practice, it all boils down to choosing the best training objective formulation, the best time sampling distribution to give appropriate relative weight to the objective (see Section 3.2), etc. – these aspects dictate model performance, independently of whether one approaches the problem from a diffusion model or a flow matching perspective.

In fact, we directly leverage the connections between diffusion and flow models when developing our stochastic samplers (see Appendix F.1) and guidance schemes. Both classifier-free guidance [1] and autoguidance [2] were proposed for diffusion models, but due to the relations between score and velocity, we can also apply them to our flow models (to the best of our knowledge, our work is the first to demonstrate classifier-free guidance and autoguidance for flow matching of protein backbone generation). Please see our Appendix F.2, which has all technical details.

Considering these relations, why did we then overall opt for the flow matching formulation and perspective? (i) Flow matching is somewhat simpler to implement and explain, as it is based on simple interpolations between data and noise samples. No complex stochastic diffusion processes need to be considered. (ii) Flow matching offers the flexibility to be directly extended to more complex interpolations, beyond Gaussians and diffusion-like methods. For instance, we may consider optimal transport couplings [3,4] to obtain straighter and faster generation paths or we could explore other more complex non-Gaussian noise distributions. We plan to further improve Proteina in the future and flow matching offers more flexibility in that regard. At the same time, when using Gaussian noise, all tricks from the diffusion literature still remain applicable. (iii) The popular and state-of-the-art large-scale image generation system Stable Diffusion 3 is similarly based on flow matching [5]. This work unambiguously demonstrated that flow matching can be scaled to large-scale generative modeling problems.

(continued below)

Q2. Clarifying the specific focus of each of the new metrics (FPSD, fS, fJSD), and possibility to merge into a more unified measure.

The Frechet Protein Structure Distance (FPSD) is inspired by the Frechet Inception Distance [5] widely used in image generation. It embeds sets of generated and reference protein backbones into a feature space and measures a Wasserstein distance between the embedding distributions. Therefore, FPSD measures distributional similarity between generated and reference backbones at the structure level, including detailed structure variations.

The fold Jensen-Shannon Divergence (fJSD) is a somewhat similar measure, but only operates at the coarser fold class level. Specifically, sets of generated and reference backbones are classified with our fold class classifier, and then the predicted categorical fold class label distributions are compared, using the Jensen-Shannon Divergence. In contrast to the FPSD, fJSD is less sensitive to local structure variations, but instead measures fold class-level similarity between generated and reference data sets.

The fold Score (fS) is inspired by the Inception Score in computer vision. This metric does not rely on a reference distribution and instead measures how well a set of generated protein backbones covers all possible fold classes. Hence, we can interpret the fS as a novel fold class-based diversity measure.

Both fJSD and fS can be evaluated at different levels of the CAT fold class hierarchy, thereby operating at coarser or more fine-grained fold class levels (we aggregate the fJSD for all different levels of the CAT hierarchy into one score). We would like to point the reviewer to Appendix E, where the metrics are discussed, introduced and validated in detail.

In theory, one could probably merge all new metrics into a single score. However, we believe that this would not be intuitive nor helpful, and goes against our intentions with these metrics. As discussed, the different metrics capture different aspects, and these would not be discernable anymore if combined into one “average” unified score. We would like to emphasize that we did not introduce these metrics just to have “yet another benchmark metric to beat”, but rather to provide novel and detailed insights (see discussions in Sec. 4.1). To this end, it is beneficial to be able to study the metrics individually, enabling different analyses.

We hope we were able to answer your questions. Please let us know if there is anything else we can clarify. We would also like to point out our additional new results that we added to the paper; please see the reply to all reviewers.

[1] Jumper et al., Highly accurate protein structure prediction with Alphafold. Nature, 2021.

[2] Lin et al., Out of Many, One: Designing and Scaffolding Proteins at the Scale of the Structural Universe with Genie 2. arXiv, 2024.

[3] Huguet et al., Sequence-Augmented SE(3)-Flow Matching For Conditional Protein Backbone Generation. arXiv, 2024.

[4] Duval et al., A Hitchhiker's Guide to Geometric GNNs for 3D Atomic Systems. arXiv, 2023.

[5] Heusel et al., GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium. NeurIPS, 2017.

[6] Salimans et al., Improved Techniques for Training GANs. NeurIPS, 2016.

We would like to thank the reviewer for their positive evaluation of our work. Below, we address the raised questions:

Continuity equation and differential symbols.

We thank the reviewer for pointing this out. We have fixed the continuity equation in the updated version of the manuscript and also modified all differential symbols to be upright.

Q1. Choice of non-equivariant transformer, data augmentations, and performance.

We chose a non-equivariant transformer as the main backbone architecture in Proteina primarily for practical and performance reasons. Building strict equivariance into networks often comes with severe architecture constraints, and the most powerful equivariant architectures often scale less gracefully than non-equivariant models when training on large and complex data. For instance, in protein structure generation, triangle attention and triangle multiplicative layers are common network components employed in high-performance equivariant architectures [1,2,3]. However, their memory consumption scales cubically with sequence length, making them highly expensive when applied to large biomolecular systems. Other equivariant networks involve very computationally expensive and memory consuming higher-order tensor products [4]. In contrast, non-equivariant transformers have been shown to scale to large-scale generative modeling tasks in the language and image modeling literature. This is the main reason why we primarily rely on non-equivariant transformer architectures also in our work. In fact, this scalability is what allowed us for the first time to scale protein backbone generation to unprecedented chain lengths of up to 800 residues (see Sec. 4.2 in the paper).

Note that in some models we do in fact also use triangle layers to boost performance, but our main driving backbone architecture is a transformer that performs extremely well also on its own, as manifested for instance in the long chain generation experiments.

Since the group over 3D rotations is “small”, we can nonetheless achieve approximate equivariance even with our non-equivariant architectures through data augmentation, improving the model’s generalization (also see the related analysis in Appendix G). It is an interesting question to explore whether re-introducing strict equivariance into our network architecture itself would further boost performance. However, as discussed, strictly equivariant architectures often do not scale well. Hence, it is questionable whether this would lead to actual benefits, when relying on current equivariant network layers. Future work could explore developing novel, highly expressive, yet computationally efficient and scalable equivariant architectures for protein structure generation.

We are glad that we were able to address the questions. Thank you for the helpful review and constructive feedback.