Reinforcement learning on structure-conditioned categorical diffusion for protein inverse folding

Authors' response

>The novelty of RL-DIF is limited. [...] We address the comparison to discrete diffusion/flow models below, but would like to motivate the contributions of our work here: (1) While this work does combine a number of prior methods, structural consistency was a difficult reward to optimize. In general, the methods we were building on top of did not work out of the box. To summarize the problem-specific optimizations: * Removing SASA features hampered the GradeIF architecture (see Appendix A1). We came up with a variant of the PiFold that supported discrete diffusion to mitigate this. * Our experiments showed that the D3PM hybrid loss out-performed GradeIF’s cross-entropy loss * We found the optimization to be sensitive to DDPO hyperparameters and reward standardization, requiring careful tuning * ESM3 does use structural consistency to improve an IF model, but we note some methodological differences (IRPO vs DDPO) and difficulty comparing performance (closed training set may have overlap with IF benchmark datasets). (2) We believe foldable diversity is an important contribution. IF methods have been primarily benchmarked using sequence recovery, sequence diversity, and structural consistency. Our work shows that these metrics alone do not fully capture the main purpose of an inverse folding model: sampling diverse sequences that fold into a target structure. >The motivation for fine-tuning the model with structural consistency to improve the sequence diversity is not very clear. We were motivated to fine-tune the model with structural consistency to improve structure quality, since continual training on sequence recovery plateaus in sc-TM and decreases diversity. We find DIF-Only achieves high diversity but poor sc-TM; RL-DIF then improves sc-TM significantly while maintaining competitive diversity (Table S2, S4). >The Foldable Diversity is only calculated based on 4 generated sequences for each protein structure, which can have high variance. We chose 4 to strike a balance between a precise measurement and minimizing computational resources In appendix A4, we study the impact of 4 vs 100 sequences and empirically observe a difference of at most 3%. Comparing to Table 1, we conclude that our improvements are outside of this error bar (i.e. on CATH-all, RL-DIF has a margin of 6%; DIF-Only by 9%). >Besides, it is also helpful to provide the ratio of the generated sequences that satisfy the sc-TM score constraint We have added Appendix A3, referring to the ratio you propose as “folding success”. We observe that a low folding success always leads to a low FD. However, high folding success does not guarantee high FD, as the sequences may not be diverse. It is this relationship that motivated our definition of FD. >A comparison with diffusion-based methods [1,2] is missing. Thank you for pointing out our oversight of these papers - we have added them to the Related Works section of the manuscript. However, we do not benchmark against either: * As there are many IF methods, we focus on SOTA models at time of writing; Campbell et al reports sc-RMSD greater than ProteinMPNN on CATH. * Nisinoff et al focused their stability optimization on specific proteins (instead of a full IF benchmark dataset), so we did not benchmark their model. >Hamming distance is used to measure sequence diversity. [...] In the IF setting, we are typically designing protein sequences for a very specific function (making the approximation that structure determines function). Therefore, significant functional diversity is not desirable. Sequence diversity is desirable because designs may fail evaluation for a priori unknown reasons - e.g. properties like synthesizability. Distance in ESM embedding space will be dominated by evolutionary and functional variation, not such properties. We prefer Hamming distance as it is interpretable and commonly used in prior IF work. Questions: > How are the temperature values for the pretrained models selected? ... How do they perform with a smaller T, eg. T=0.05 or T=0.01? We evaluated ProteinMPNN at a temperature of 0.05 and 0.01 on TS50; FD decreases to 5% and 1% respectively. The temperatures in the manuscript are those previously found to achieve best performance in ProteinMPNN, PiFold, and KWDesign [1]. [1] https://openreview.net/forum?id=mpqMVWgqjn; Figure 6 > For a rough comparison, how is the performance of ESM-IF trained on the same dataset as RL-DIF-100K? ESM-IF authors have not released their training code, so we are not able to reproduce their work. Instead, we trained a DIF-Only-250k on 250k structures, resulting in TS50 FD comparable to ESM-IF (see table below), with 18x fewer parameters and 48x less data. It remains for us to benchmark on CATH and TS500, and repeat RL tuning. | Model | FD | sc-TM | | --------|----| -- | | ESMIF | 43% | 0.83 | | DIF-Only-100k | 39% | 0.71 | | RL-DIF-100k | 39% | 0.80 | | DIF-Only-250k | 43% | 0.75 |

Response to authors

Thank the authors for the further clarifications about the contribution and motivation of their work and for attempting diffusion-based models. In light of this, I increase my score. I would like the authors to add these clarifications in the revised manuscript. Regarding the diffusion-based models, I understand the time limitation of the rebuttal period. I would expect an inverse folding model to work well if the authors follow the steps in Section F.4.1 in [2] and encourage the authors to have a try.

Clarification question

Thank you for your review! We have a quick clarification question as we work on producing answers to your questions: when you refer to "structural similarity", do you wish to see the similarity amongst generated structures (e.g. a form of structural diversity) or between the generated structures and the target structure (i.e. what we call "structural consistency" in the paper)?

Authors' response

> The paper feels close to the acceptance threshold for me. However,I felt the authors could do a bit more. The increase in seq diversity generated by RL-DIF is modest in some cases. The authors could provide (A) more information on the trade off between structural similarity and sequence diversity along the training trajectory (b) More information on how hyper parameter can be used to relax structural similarity constant to increase diversity (c) exploration of regularization strategies like entropy based regularization to increase the diversity of sequences generated about the gains shown. (d) validation of structural similarity of generated sequences via alpha fold. A:We have added Appendix A3 to study this, both along the training trajectory and across different models/sampling strategies. We compare foldable diversity (FD) with two other metrics: folding success (fraction of sequences that meet the sc-TM threshold) and folded diversity (FdD; sequence diversity among sequences meeting the threshold). Table S4 shows that folding success improves over the first 1k RL steps, while FD is maintained. B: Essentially the only training hyperparameters that can affect diversity are (1) sc-TM threshold and (2) number of RL steps. Profiling (1) is expensive (requiring multiple full retrainings), and lowering the threshold significantly would reward biologically inconsistent structures. In Table S4 we study the effect of (2), which does suggest that diversity increases initially, but then collapses with too many RL optimization steps. C: We did not explore adding an entropy regularization term in this work, but agree that could be a more direct way to boost diversity. D: We appreciate your suggestion to use AlphaFold for structural validation. While AlphaFold is generally recognized for its higher accuracy, we opted for ESMFold primarily due to its lower computational resource requirements. In our study, where we evaluated thousands of sequences, using AlphaFold would have been prohibitively expensive in terms of computational cost and time. We also have some confidence that evaluating with AlphaFold would not significantly change the deltas between benchmarked IF methods. TM-score between ESMFold and AlphaFold are correlated, especially if we are binarizing with a sc-TM threshold of 0.7 [1]. Therefore, the relative ranking of models in terms of foldable diversity would likely remain unchanged. Nonetheless, we agree with the reviewer that repeating this work with AlphaFold could be an important next step. [1] https://www.sciencedirect.com/science/article/pii/S0022283624001888#f0005 (Figure 1) Questions: > Could the authors could provide (A) more information (plots) on the trade off between structural similarity and sequence diversity along the RL training trajectory (b) More information on how hyper parameter can be used to relax structural similarity constant to increase diversity (c) exploration of regularization strategies like entropy based regularization to increase the diversity of sequences generated about the gains shown. (d) validation of structural similarity of generated sequences via alpha fold. We believe our response to the questions posed in Weaknesses addresses these questions.

Clarification question

Thank you for your review! We are working on thorough responses to all your questions, but would like to quickly clarify the question re: Table 1 - we believe Figure 2 shows what you are requesting (curves of diversity as scTM threshold is varied, across all benchmarked methods). Could you clarify if there is something else you'd like to see?

Authors' response

> The proposed method is largely a direct combination of available techniques. ... Of course, it's fine to mix and match methods with adjustments to produce a well-performing technique but then the paper rests on its empirical results. We thank the reviewer for their comment and would like to take this opportunity to clarify our contributions: (1) While this work does combine a number of prior methods, we found structural consistency to be a difficult reward to optimize. In general, the methods we were building on top of did not work out of the box. To summarize the problem-specific optimizations: * Removing SASA features significantly hampered the GradeIF architecture (see Appendix A1). We came up with a variant of the PiFold that supported discrete diffusion to mitigate this regression. * Our experiments showed that the D3PM hybrid loss out-performed GradeIF’s cross-entropy loss * We found the optimization to be extremely sensitive to DDPO hyperparameters and reward standardization, requiring careful tuning It is a generally open problem of how to design diverse protein sequences meeting complex quality criteria. In this paper, we aim to demonstrate that a careful application of these prior methods results in a model optimized for one of the most fundamental of those criteria (structural consistency) (2) We believe an important contribution of our work is the introduction of foldable diversity. IF methods have been primarily benchmarked using sequence recovery, sequence diversity, and structural consistency. Our work shows that these metrics alone do not fully capture the main purpose of an inverse folding model: sampling diverse sequences that fold into a target structure. We demonstrate an inherent tradeoff between the foldability of model-generated sequences and the diversity of those sequences. Foldable diversity is a metric designed to measure this; models that are less sensitive to the tradeoff have a better foldable diversity. (3) We explore scaling properties of inverse folding models in an open-source way. In particular, we report that a careful scaling of pretraining data, model size, and on-policy training enables RL-DIF to achieve competitive performance with much larger models trained on far more data. To enable future work and direct comparison by the research community, we feel it is important to release the IDs of the proteins used in training our models. > I think it would be better to replace eq (8) with a measure calculated only within foldable sequences so as to separate foldability from diversity (so they can be assessed separately and together). Diversity was taken as the key motivation for the paper but the two notions (diversity and foldability) are now mixed together in the proposed metric. E.g., the downward slope in Figure 2 is in part just due to the fact that higher folding threshold automatically decreases the value of the current metric. Figure 2 can also give a wrong impression that RL-diff increases diversity since the metric can be improved by increasing foldability instead. In fact, the reward used for fine-tuning the diffusion model is sc-TM so tailored to increase foldability (not control diversity). We thank the reviewer for pointing out this important point. This is indeed something we considered as we constructed the foldable diversity (FD) metric, and have added Appendix A3 to study your proposed metric (which we call folded diversity (FdD)) in detail. Our observations demonstrate the expected trade-off (FdD can be increased at the expense of folding success). But we consider FD to be a more grounded metric because it measures diversity in the realistic scenario of a fixed sampling budget. FdD measures the diversity among good sequences, after rejecting bad samples. This can be maximized by sampling from a uniform distribution over sequences and simply rejecting those that fail to fold (some settings in Table S2 have a 90+% fail rate). > since the metric can be improved by increasing foldability instead In general, we observe IF models increase foldability at the expense of diversity, which does not improve FD alone. RL-DIF is less sensitive to this tradeoff and consistently increases the folding success over DIF-Only from anywhere from 2% to 19% while maintaining a competitive foldable diversity. Questions: > The expectation after eq (3) should really be a sum > q(S_t|S_{t-1},v) should be q(S_{t-1}|S_t,v) in eq (9) > In the model architecture, line 183, do you mean p(s_0|s_t) rather than p(s_{t+1}|s_t) as currently written? Thank you for catching these typos - fixed in the manuscript. > Table 1 would be easier to read if summarized in terms of curves. For example, calculate diversity within all generated samples from a method that pass a sc-TM threshold and then vary this threshold. We believe Figure 2 shows the requested curves, but are happy to iterate on the figure if we have misunderstood this point.