Generating All-Atom Protein Structure from Sequence-Only Training Data

We’d like to thank the Reviewer for their insightful comments, and for highlighting the importance of being able to train on larger sequence-only datasets. We also are glad to see the Reviewer highlight the unique difficulties of all-atom generation.

To address some weaknesses:

“The CHEAP autoencoder used in this work does not seem like a variational autoencoder. Therefore, the latent space may not be very smooth and can make diffusion model training difficult.”

• This is a great observation. In initial experiments, we struggled a lot with the ‘roughness’ of the inherent latent space of ESMFold; this is actually why we use the CHEAP autoencoder. CHEAP applies a channel normalization to combat massive activations, and though it is not an VAE, the authors show that the latent space is smooth and Gaussian-like after smoothing & compression. If interested, Reference [1] discusses this at length.
• We've added a visualization of embedding values after compression in Appendix Figure 7 to address this comment.
• We also updated Figure 5 to show how noising in the latent space maps back to sequence and structure space before and after using the CHEAP encoder.

“The unconditional generation performance of the model after training on Pfam is somewhat lackluster compared with MultiFlow.”

• It's important to note that Multiflow is not an all-atom generative model, and only generates backbones and the residue identity. Please see our General Response for more details on this point.
• When examining performance by length (Updated Figure 5), we also find that PLAID better balances sample quality and diversity at higher sequence lengths.

“I also believe that an experiment showing how the performance of the proposed model scale with more data can be insightful.”

• This is a wonderful suggestion. It was difficult for models to train for long enough to say something concrete about this before the end of the discussion period, but if accepted, we will add this to the camera-ready version.

“The latent diffusion choice of the framework can be really beneficial in terms of speed when generating very long sequences (>600 aa).”

• We thank the Reviewer for this suggestion. We've added a comparison of sampling speeds for $L=600$ to Table 4, both for the batched and unbatched setting.

Re: CHEAP model details:

• We’ve included more detail on how the CHEAP [1] module works. We've also added experiments to show diffusion without CHEAP compression in updated Table 8.

Re: conditional generation details:

• We’ve included a detailed panel (Updated Figure 2D) to describe how conditioning was added via AdaLN in DiT blocks. We also included a longer discussion of classifier-free guidance.

Conclusion. We’ve aimed to address the Reviewer’s concerns around writing/figure clarity, performance, and included additional experiments for sampling speed. We hope the Reviewer will consider reviewing our updated manuscript and improve their score. We also welcome any additional comments.

[1] Tokenized and Continuous Embedding Compressions of Protein Sequence and Structure. (Lu et al., 2024)

To address questions:

“If we focus less on the structures obtained from experiments and instead accept those predicted by models, the available structural data is still substantial.”

• Synthetic data has been known to create errors such as mode collapse [2]. (Interestingly, this is precisely what we observe in Multiflow for higher sequence lengths, which does use distillation.)
• Using latent diffusion in a compressed space allows easier exploitation of hardware-aware attention kernels. Structure inputs often require architectures and data representations less amenable to scaling.
• Our approach democratizes all-atom generation. Synthetic data is computationally expensive to generate, but PLAID can be used as long as model weights are available from a pretrained foundation model.

“Does designing based on structure prediction models introduce an unnatural bias toward certain protein structures? For example, AF2 has been found to predict overly "clean" structures.”

• We agree that the bias towards overly “clean” structures is a pervasive issue, and this actually is precisely what motivated our work. The inherent bias towards “clean” and crystallizable structures in the PDB is severe; by incorporating data from sequence databases, we can actually sample from those regions, which would otherwise not be captured in existing protein structure generation methods.

“Is there a more detailed explanation or experimental validation available for using the CHEAP module?”

• Thanks for this suggestion; we’ve added results from learning on the latent space without using the CHEAP module (Updated Appendix Figure 8), and also plot & visualize how noise in the latent space maps back to corruptions in sequence and structure space (Updated Figure 4). We also visualize the latent space in Updated Appendix Figure 7.

“Why is a design result that finds more homologous sequences considered more novel in the context of defining Hit metrics?”

• We've removed Hit % as it was a confusing choice of metrics. In Updated Table 3, we use sequence identity % to the closest mmseqs neighbor and structure TMScore to closest foldseek neighbor, to assess sequence and structure novelty, respectively.

Conclusion. We thank the Reviewer for their detailed examination and suggestions, which has improved this work. We hope that our response make our reasoning around choice of conditioning clearer, and that our updated results reflect how PLAID circumvents mode collapse observed at longer lengths in other approaches. We would ask the reviewer to consider increasing their score, and are happy to address any additional comments or suggestions during the discussion period.

[1] Tokenized and Continuous Embedding Compressions of Protein Sequence and Structure. (Lu et al., 2024)

[2] Synthetic Data, Real Errors: How (Not) to Publish and Use Synthetic Data. (Breugel et al., 2023)

We want to thank the Reviewer for acknowledging the innovative and refreshing aspects of our approach. We're also grateful for the suggestions, which have helped us make many improvements. We'd appreciate it if you could take a new look at the PDF file and see if your concerns were addressed.

To address the stated weaknesses:

Re: performance: “My main concern with the paper lies in its underwhelming performance…”

• We've greatly expanded the depth with which we've been examining the methods. When we expanded sample lengths to 512 and examined performance by length, we found that PLAID is better at balancing quality and diversity for longer sequences (Updated Figure 6 & Appendix Figure 11).
• It should be emphasized the Multiflow is not an all-atom method. Amongst all-atom generation methods, PLAID achieves SOTA on consistency and naturalness metrics, both of which are important metrics that have been neglected by co-generation literature. Please see our General Rebuttal for more on this point.

“The lack of support for motif scaffolding raises concerns about the controllability of the current model. Although function and organism terms can be used as conditioning, this form of conditioning may still fall short in achieving precise control.”

• We were deliberate in choosing capabilities that would enlarge the surface area of how generative models can help drug development. Organism conditioning can be key to many developability and humanization bottlenecks. GO terms act as proof-of-concept for the eventual aim of scaling to more complex annotations, such as natural language, which affords precise control that is relevant to many domains. Please see our General Response for additional detail on this point.
• In our updated manuscript, to highlight precise control, we’ve investigated function-conditioned sampling, and found that PLAID preserves precise catalytic motifs (Updated Figure 7, Appendix Figure 12).

“The paper…overlooks hallucination models that also manipulate structure prediction models in similar ways.”

• We thank the reviewer for feedback on our Related Works section and have expanded the discussion to address how our approach differs substantially from hallucination. Anishchenko et al., focus on sequence design with MCMC in sequence space. It doesn’t propose new structures, and instead looks for a sequence that folds into that model. In contrast, we are doing multimodal / all-atom design in latent space. This problem setting encompasses that examined in Anishchenko et al., but is vastly more difficult, since we also search through structure space.

“The related model design strategies require further clarification...”

• Thank you for this feedback. We've added an ablations table (Updated Table 1) to highlight this. We've also rewritten the Methods section. Many design decisions were driven by scalability and generalizability to new methods, and being able to use maximally efficient attention kernels, such that it would be easier to build upon it (either the method or the weights itself).

“There is no explanation provided on how the model integrates function and organism conditioning…”

• We added a detailed panel (Updated Figure 3D) which describes the AdaLN operation used in DiT blocks to incorporate conditioning information.

“The design of the sequence decoder…also requires further elaboration…”

• As noted in Methods, the sequence decoder is taken from [1]. However, we've incorporated the feedback and restructured our manuscript to make this information easier to find.

“The paper does not provide an explanation of how FIDs…is adapted for application in protein design.”

• This is a good point, and in retrospect, “Frechet Distance” would have been a better description. In our Updated Figure 8, we instead use Sinkhorn Distance between the generated latents and the corresponding embeddings from a distribution of real proteins. This also allows us to use smaller sample sets, so we are able to compare more classes.

(Continued below...)

The author's response has addressed some of my concerns, and I have accordingly raised my score. However, I still feel that the author’s perspective on their work is too narrowly defined. Specifically, whether one designs structure and sequence simultaneously, designs structure first and then sequence, or explores the sequence space before folding the structure, all fall within the realm of protein design. A comprehensive study should include comparisons of different approaches.

We thank the reviewer for these helpful comments, and for acknowledging the method’s strength in being able to scale across larger datasets.

In Weaknesses, the Reviewer makes an astute observation that the performance of ESMFold limits this model. We think an extension of this paradigm to include a finetuning of the ESMFold model, or developing new frontier biological foundation models for PLAID, would be extremely useful – hence why we’d like to share this work with the research community.

Regarding questions:

• “The statement in line 178, p(x) = \phi(s) = \phi(\omega), I am not sure if this is correct…”

$\phi(\cdot)$ is a mapping that transforms different distributions. We’ve updated the writing to make this clearer.

• “how the ESMFold first component that maps the sequence to the latent captures evolutionary prior?”

ESMFold uses ESM2 in the first component, which was trained with a masked language modeling objective on UniRef. This is also empirically justified by the fact that replacing the explicit MSA construction in AlphaFold with ESM2 yields comparable results. We’ve updated the writing to clarify this.

• “The equation at the end of the line 187, isn't the structure module takes actually x and map it to the strucuture \omega, the equation written as if it takes the structure as input?”

We use the notation for inverting a function to say that an inversion of the structure module would provide a mapping from structure to $x$.

• “I am not sure if I understood the part the sequence decoder need to be separately trained”.

Yes; as noted in the Methods, the sequence decoder is separately trained, and provided in [1]; it reaches a validation accuracy of 99.7%. We’ve reorganized the sections to make this easier to find.

• “the line 199 ( the inference ... is shown in Figure 2B, is it typo?”

Thank you for noticing this. We’ve entirely updated Figure 2 to also include a schematic of how conditioning information is incorporated into each DiT block.

• “if I summarize the training of this model, …” Yes, this is the correct summary. “...not sure how x_norm is obtained…”

This comes from the CHEAP [1] model, where massive activations are removed via a post-hoc channel norm operation. We've added a longer discussion of details of the CHEAP model.

• “Any intuition behind the why the compressed latent first dimension set L/2?”

This reduces memory usage, since we use a transformer architecture for scalability, and memory increases quadratically with length. We added this clarification to the updated manuscript.

• “I am wondering what happens if one apply the diffusion model directly on the latent of the ESM2 model…”

We’ve added Appendix Figure 8 to show performance before and after compression, and include a longer discussion of high-resolution image synthesis via latent diffusion in related works.

• “The figure 4D has beed resolution, and color of the round dots are a bit misleading.”

We’ve incorporated this feedback and updated this figure with expanded results in Updated Figure 9.

• “what is the red vertical line in the figure 4 B represents?”

This is the threshold for which “designability” is defined. In the Updated Figure 6, which examines performance by length, we’ve made sure to make this more clear.

• “I am wondering what happens if one trains the full Model in figure 2C instead of freezing some parts?” This is a great suggestion and opportunity for follow-up work. It should presumably improve performance, but retraining/finetuning ESMFold in addition to the models trained for this paper was not feasible with our resource constraints.

Conclusion. We're very grateful for the comments from Reviewer n5hb, and hope that the updates can address these concerns. It would be great if the Reviewer could improve their score, and we’re happy to address any further comments before the end of the Discussion Period.

As the discussion period is coming to a close, we'd appreciate it if Reviewer n4hb can find a chance to review the updated PDF. We've made significant changes to address the Reviewers' comments, including rewriting the Methods section to reflect Reviewer n4hb's suggestions, experiments that analyze when and how PLAID outperform baselines, expanded discussion of why we used embedding compression, and better highlighting how PLAID correctly positions sidechains at active sites when prompted by function.

If we've been able to address the Reviewer's concerns on the impact of this paper, we'd appreciate it if you can improve your score accordingly. Thanks for taking time to engage with our paper.

3. Re: similarities to ESM3: “In newer work like ESM3, the language model is directly used for structure generation via explicit structure token embeddings. How does this compare…”

• This is an interesting insight. This work was developed concurrently with ESM3; we’d also considered a unified tokenized approach at first, but chose diffusion because existing literature shows that diffusion models at similar parameter counts can achieve better performance than autoregressive ones. [4]
• It’s quite straightforward to extend PLAID to use tokenized representations, since CHEAP [2] embeddings also include a tokenized version.
• We've updated the related works based on this suggesiton.

4. Re: naturalness: “…should naturalness in terms of aromaticity etc be a target for a "good" protein generative model?”

Previous literature [1] shows that distribution conformity to natural proteins is very correlated with expression. (Please see our General Response for further discussion.)

“...For many applications these models are used for non-natural properties like higher melting temperatures etc”

This is a fascinating suggestion. In PLAID, we can do this by conditioning by a thermophilic organism. We'll add this experiment before the end of the discussion period.

Re: conditional generation evaluation: “...how well do these case studies work across the board beyond the two examples shown in the paper?”

We have significantly expanded our analyses, showing that function-conditioned generations conserve active site residues, global structural similarity, and hydrophobicity patterns of membrane proteins.

Re: PLAID compatibility to other conditioning capabilities:

“In protein generation it is often useful to employ structure constraints such as motifs, specific secondary structure or binding interfaces for generating samples. Is this possible in a latent framework such as the one proposed here?”

• To guide PLAID by per-token secondary structure, one can use a text-conditioned approach. In an early iteration of the model, we’d conditioned the model by secondary structure content, but are focusing on GO term and organism control for better applicability to drug development and introducing new capabilities for all-atom generative models.
• For binding interfaces, one can take an in-painting approach. Since at inference time, the user specifies the length, one can provide a sequence and leave “additional space” for the model to in-fill the binder.
• For motif scaffolding, one can keep the input motif fixed during inference.

Conclusion. The Reviewer’s suggestions have been helpful for improving the clarity of our paper and the strength of our claims. In light of the additional results we provide, we’d appreciate it if the Reviewer would consider reassessing this work and improving their score.

[1] Protein Discovery with Discrete Walk-Jump Sampling. (Frey et al., 2023)

[2] Tokenized and Continuous Embedding Compressions of Protein Sequence and Structure. (Lu et al., 2024)

[3] GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models. (Nichol et al, 2021)

[4] Denoising Diffusion Probabilistic Models. (Ho et al., 2021)

We thank Reviewer bWWC for the insightful comments that have helped us improve this work, and for acknowledging the novelty and scalability of our approach.

To address the stated weaknesses:

1. Designability/diversity performance:

• We'd like to note that Multiflow does not produce side chain positions; if we assume each side chain to have 0 to 4 rotamers, this induces $4^L$ more degrees of freedom. Multiflow also cannot mirror PLAID’s ability to capture sidechain placements and precisely model ion binding and active sites (Updated Figure 6).
• Designability metrics are extremely fragile, and have strong structure bias. Since ProteinMPNN, etc. are trained on the PDB, these metrics essentially assess if new samples are in distribution with the PDB. Our goal is explicitly to remove over-emphasis on structure. They are also fallible; the reported sequence recovery rate of ProteinMPNN is 52.4% [2].
• Natural proteins themselves have low self-consistency performance. This can be seen in Tables 2 & 3. This is also why we emphasize cross-consistency; as seen in Updated Table 2, natural proteins actually achieve near perfect performance here.
• Novelty performance: in our Updated Table 3, we’ve removed names like “Hit %” that have caused confusion.

2. Control experiment using ProteinMPNN:

• We’ve added this experiment by using ProteinMPNN to predict sequence from a given structure, and compare its performance to the original structure (i.e. the scTM and scRMSD metrics) in Updated Figure 5, Table 2, 3. Consistency here is lower than using the sequences directly produced by PLAID. This reaffirms our original thesis and the simplicity, robustness, and novelty of end-to-end training and co-generation without alternating model calls.

3. Naturalness of sequences:

• We’re grateful for the Reviewer’s careful consideration and have updated the table to reflect length normalized molecular weights.
• Re: stability metric, this is a heuristic calculated from the prevalence of dipeptides, and should not be interpreted as a gold standard. * To illustrate this, we’ve included a stability distogram of real, experimentally stable proteins in Updated Appendix Figure 10E.

4. Ambiguous usage of FID term:

• We apologize for confusion here; “Frechet Distance” would have been more suitable. We updated this experiment using Sinkhorn Distance. This is also more resilient to small sample sizes, which enabled us to look at more function classes.
• As to why we run this experiment: this assesses the distance in latent space between the real and generated proteins, to assess “proteinness” separately from decoder performance.

Regarding questions:

1. Conditioning capabilities: In Table 1, the authors imply that their model can perform two GO term and organism conditioning, while the other methods according to the table cannot do conditioning…

We’ve removed Table 1 given that it’s impossible to do a full comparison.

“...for example RFDiffusion can do a lot of conditioning constraints that PLAID cannot (symmetric oligmers, binders, ...) which for practical applications have arguable more relevance.”:

We’d respectfully disagree that function & organism conditioning are less relevant; they are in fact crucial to developability and humanization bottlenecks in drug discovery, and introducing new capabilities should be considered a strength. Please see our General Rebuttal for more on this point.

2. Re: domain generation: “Since the model is trained on protein domains in Pfam only, does this limit the generation capabilities to single-domain proteins?”

We discuss our rationale for using Pfam in the “data” section. In brief: (1) memory scales quadratically with length, so we wanted to maximize content in short sequences; (2) It’s difficult to perform in silico comparisons on full proteins containing disordered regions; (3) Most of our evals are around sequence-structure and motif conservation, so nothing is lost by focusing on single domains.

(continued below)

Thanks for these comments and encouragement on the novelty of our approach.

1. I think we are in agreement that for unconditional design, the best validation is to see if it can be manufactured in the wet-lab; and in the absence of that, we need to look at in silico metric that seem to correlate with it. Distributional conformity has been shown in literature to correlate with real-world expressibility. WJS [1] achieved 70% expressibility using distributional conformity scores, which at the time was "the highest reported binding rate of any antibody design method applied to trastuzumab CDR H3 redesign." We therefore think that this is a valid metric to at least consider, if not prioritize. Designability metrics have heavy bias towards samples in-distribution with what ProteinMPNN/etc. are trained on; biophysical parameter do not have this same bias. We hope this better describes our rationale behind this metric and can put us in better agreement. Furthermore, even if we examine only designability, PLAID achieves the best results of all other all-atom methods. If there are metrics that the Reviewer think is explicitly missing, please let us know.

2. As mentioned, having to model side chain positions adds up to $4^L$ additional degrees of freedom, which is non-trivial. If we only wanted to solve the Multiflow setting of backbone atoms + residue identity, that would give much more leeway to method design. Re: side chain capabilities: Updated Figures 9 and 12 shows several cases of precise ion coordination, such as learning the tetrahedral geometry of cysteine residues coordinating the iron ion, learning the DHDH motif for zinc binding, and more. All of these side chain positions were directly produced by PLAID. This capability actually directly mirrors the recent preprint referenced by the reviewer [2]. Are there any other specific experiments that the Reviewer might want to see on this front?

3. Re: ProteinMPNN fallibility: one of the major motivations for our work is to move away from reliance on these tools, since it compounds errors, increases installation dependencies, is slower, etc. Our point is not to say that ProteinMPNN is not a good tool -- it's of course been successfully and widely used -- but as a field, there's benefit to not putting all of our eggs in one basket. It keeps the field more nimble for progress. The point we wanted to convey is that many decisions made here were very intentionally, including why we do not use ProteinMPNN for inverse folding.

4. Re: Frechet and Sinkhorn Distances: the Frechet Distance and Sinkhorn Distance both provide a means of characterizing distance between two high-dimensional quantities. Typically this is used to assess distance between embeddings, for e.g. Inception embeddings, but since in our work, we are directly generating latent embeddings, we just directly calculate the distance in this space, i.e. in the CHEAP latent space. We will update this in the manuscript to make it more clear, along with the other fixes suggested.

5. On motif scaffolding: we maintain that not centering our method on motif scaffolding is a strength rather than weakness, as it increases the surface area of what we can examine as a field. We're happy to implement motif scaffolding in PLAID since that has been important to the Reviewer, but we want to re-highlight the importance of the conditioning tasks we choose:

• Organism expressibility is very important; not being able to express your antigen or binder in a given system is an important bottleneck in scientific discovery
• The goal of motif scaffolding is ultimately to preserve function, which GO terms also provide definition for (and describe along a different axis)
• To transfer a function from one organism to another, motif scaffolding could fail if the motif is different in the target organism; our compositional conditioning approach tackles this.

We appreciate you engaging with our paper. If this discussion has helped us reach better agreement on our contributions, it would be great if you can improve your score.

[1] Protein Discovery with Discrete Walk-Jump Sampling. https://arxiv.org/abs/2306.12360

[2] Computational Design of Metallohydrolases. https://www.biorxiv.org/content/10.1101/2024.11.13.623507v1

I appreciate the effort the authors put into the rebuttal by fixing the mistakes and missing figures in the manuscript, clarifying some of my questions regarding naming and metrics in the paper and conducting additional experiments and increase my score accordingly. Some mistakes such as missing reference links (see "Table ??" at the end of page 8) and should be fixed.

I think the paper has an interesting novel approach as mentioned before and can be turned into great work. However, as I mentioned in my initial review, a few things such as some of the experiments now acknowledged as interesting by the authors as well as proper metrics validating the approach should be added in order to make the work more convincing and stronger. Some more detailed responses below.

• Although I agree that designability metrics etc can be biased for structure, it is an experimentally verified metric for lab success of designed proteins. It is not necessarily the aim of protein design to exactly mimic all natural protein properties (e.g. better expressibility, thermostability etc). If the authors think that their approach would benefit from a different metric, they should compare on that metric and thoroughly verify that this is a meaningful metric in a practical sense.
• Re the point about the MultiFlow comparison: it is not fully clear to me yet why the all-atom capability in your use cases is necessarily improving upon just modelling backbones + sequence as in Multiflow; experiments that show the advantage of this capability via side-chain conditioning as in other recent work [1] or other leverage of the all-atom capacity for design would strengthen that point. As long as such things are not possible, the clear advantage of producing an all-atom structure in the end is not obvious to me.
• What is meant regarding ProteinMPNN being fallible with a 52% sequence recovery? Of course, the model is not perfect, but why is the 52% sequence recovery indicative of that? I at least could not point to the "best" value for sequence recovery since many sequences can fold into the same structure.

2. FID vs Sinkhorn: I still do not understand what model was used for the Frechet distances; this needs to be computed via some classifier now which embedding space the distance is calculated, and I see no classifier mentioned anywhere nor justified why it is useful/appropriate for that task.

[`1] Kim, D., Woodbury, S. M., Ahern, W., Kalvet, I., Hanikel, N., Salike, S., ... & Baker, D. (2024). Computational Design of Metallohydrolases. bioRxiv, 2024-11