All-atom inverse protein folding through discrete flow matching

1.What is more impactful, the architecture introduced or the discrete flow matching? A lot of attention to the architecture is provided and it clearly is beneficial but deeper ablations as to what degree does the underlying generative framework and architecture play a role would strengthen the contribution.

Thank you for your comment. Both the multi-scale GNN architecture and the discrete flow matching framework play important and complementary roles in ADFLIP. The flow matching framework provides a flexible generative backbone that allows for integrating diverse sources of information—such as multiple structural states and external guidance signals. As shown in Table 4, this enables improved generation performance under dynamic structural contexts. On the other hand, the multi-scale architecture dynamically captures both atom-level and residue-level information, enabling the model to reason over partial side-chain context during sampling. To assess its effect, we conducted an ablation study:

While the numerical gains in recovery rate are moderate, we believe the use of side-chain information has significant implications for downstream applications of inverse folding, such as enzyme design or protein–ligand interaction modeling—because the chemical interactions with small-molecule ligands or substrates are typically through the protein side chains, and where even minor torsional differences in side chains can result in substantial functional changes. We will point out this advantage of our all-atom approach more explicitly in the revised version.

What are the mean/std or distributions of the sequence recoveries? We have computed the mean and standard deviation of the sequence recovery rates for LigandMPNN and ADFLIP as below. We will also include the distribution in the revised manuscript.

What happens when you do not use purity sampling wrt the benchmarks? We performed an ablation study to evaluate the impact of purity sampling on performance across benchmarks. Specifically, we varied the purity threshold and compared it to a fixed-step denoising strategy. Purity Sampling (variable threshold):

Fixed-Step Sampling:

We will add these new results as supplementary data to the revised manuscript. [Will we add these to the revised manuscript?] Kai: Yes

How is classifier guidance implemented for the discrete flow models? A algorithm would be useful here.

Thank you for this suggestion. We have included the training-free classifier guidance sampling algorithm below, which we will incorporate into the revised manuscript for clarity.

Algorithm: Training-Free Classifier Guidance Sampling

Input:

• $N$ protein and non-protein structures $\{x_1, ..., x_N\}$
• Initial sequence $s_0 = (`[MASK]`, ..., `[MASK]`)$
• Initial sidechains $\chi_0 = \emptyset$
• Time step $\Delta t$
• Denoising network $f_\theta$, sidechain packing network $g_\eta$, regressor network $h_\phi$
• Target property value $y$

Procedure:

1. Initialize $t = 0$
2. While $t < 1$:
   • For $n = 1$ to $N$:
     • Compute $p^n(\hat{s}_1) = f(s_t, x_n, \chi_t)$
   • Average over structures:
     $p(\hat{s}_1) = \frac{1}{N} \sum_n p^n(\hat{s}_1)$
   • Compute predicted property:
     $\hat{y} = h_\phi(p(\hat{s}_1))$
   • Compute guidance respect to $s_1$:
     $p(y \mid \hat{s}_1) \approx -\nabla||y - \hat{y} ||^2$
   • Sample $s_1 \sim p(\hat{s}_1) \cdot p(y | \hat{s}_1)$
   • For $n = 1$ to $N$:
     • Compute sidechains:
       $\chi^n_1 = g_\eta(s_1, x_n)$
   • Compute reward:
     $R_t(s_t, j) = E_{p_{1|t}(s_1|s_t)}[R_t(s_t, j|s_1)]$
   • Sample $s_{t+\Delta t}$ using Eq.1
   • Update time: $t \leftarrow t + \Delta t$
3. Return: Final sequence $s_1$

1.Related Work

Thank you for pointing this out. We mentioned the paper by Yi et al. [1] in the introduction, but agree it should also be included in the related work. We will rewrite the related work section and discuss the discrete diffusion method for inverse protein folding from Yi et al and also another diffusion model [2][3]. This will help clarify how our work connects to previous studies.

[1]Yi, Kai, et al. "Graph denoising diffusion for inverse protein folding." Advances in Neural Information Processing Systems 36 (2023): 10238-10257.

[2]Wang, Xinyou, et al. "Diffusion language models are versatile protein learners." arXiv preprint arXiv:2402.18567 (2024).

[3]Wang, Xinyou, et al. "Dplm-2: A multimodal diffusion protein language model." arXiv preprint arXiv:2410.13782 (2024).

2.Methodology Explanation for Flow Matching

Thank you for this suggestion. It is correct that, in our framework, discrete flow matching generates the distribution $p(s_t​)$ at each timestep. However, unlike continuous flow models, where the denoiser predicts noise, the denoiser of discrete flow matching directly predicts data (e.g., sequence tokens). In our implementation, a trained denoiser first estimates $p(s_1)$, and then $p(s_t)$ is computed by taking an expectation over this estimate, as shown in Equation (line 137) and Algorithm 1.

In the revised manuscript, we will add clearer descriptions, expanded mathematical formulations, and improved illustrations to explain the conditional discrete flow matching process more intuitively.

3.Classifier Guidance

Our classifier-guidance approach shares a similar objective with Chroma and RFdiffusion in that we aim to condition generation on an external property $y$, estimated via $p(y∣s_t)$. However, there is a key difference in how this is achieved. Chroma and RFdiffusion use continuous diffusion models and typically require training a separate classifier or regressor (y = f(s_t)) directly on noisy intermediate states, $s_t$. In contrast, our method uses a training-free guidance strategy: we leverage a pre-trained classifier or regressor that operates on clean data $s_1​$, such as AlphaFold. We estimate $p(y∣s_1)$ from this clean output and then derive $p(y∣s_t)$ by taking the expectation over s1∼p(s_1∣s_t), as described in Equation (line 307). This approach avoids the need to retrain an external model and enables seamless integration of existing structure-based predictors.

4.Methodology Clarification: Does removing sidechains at masked positions bias decoding toward early samples, and is there an ablation study on this effect?

Thank you for this insightful comment. We conducted an ablation study on the parameter $\tau$ for purity sampling, which controls how much side-chain information is preserved for masked positions during the denoising process. A smaller $\tau$ value retains more side-chain atoms, providing richer structural context to the denoiser. Conversely, a higher $\tau$ results in fewer sampled residues and less side-chain information.

5. Need for Ablation Studies for sidechain

We conducted an ablation study to assess the impact of incorporating predicted side-chain atoms into the denoiser's input. The results are summarized below:

While the inclusion of side-chain information leads to moderate improvements in sequence recovery, we believe its primary value lies in enhancing biological fidelity rather than optimizing this metric alone. In many downstream tasks—such as enzyme design or protein–ligand interaction modeling—minor torsional differences in side chains can lead to substantial functional changes (e.g., in binding affinity or catalytic activity). Therefore, even small gains in sequence recovery reflect a more meaningful structural signal that allows the model to better capture fine-grained biophysical nuances. We will discuss this in more detail in the revised manuscript.

6.(Conditional Flow Matching Formulation) Additional information about the conditional and the formulation from Campbell et al would be helpful for readers.

We agree, and will add more detail about Conditional Flow Matching in the revised manuscript.

How does ADFlip perform on protein only dataset such as CATH

Thank you for the suggestion. We retrained ADFlip on the CATH 4.2 dataset and evaluated its performance based on sequence recovery rate. The results are summarized below:

How does other inverse folding model perform on all-atom dataset

Thank you for the suggestion. Due to time constraints during the first review period, we are still working on additional evaluations. We aim to include results for other inverse folding models on the all-atom dataset in the second-round discussion and the revised manuscript.

Clarification on "All-Atom" Structure in ADFLIP

Perhaps we had not explained clearly enough what we mean by ‘all-atom’. This interpretation of “all-atom” was introduced previously with RoseTTAFold-all atom[1]. In this context, all-atom refers not only to including the full set of protein atoms (backbone and side chains) but also to incorporating non-protein biological components such as ligands, ions, and nucleotides. Most inverse folding methods—including PiFold, ProteinMPNN, and others—limit their input representation to backbone atoms (typically N, C, O).

By contrast, our method, ADFLIP, is designed to operate on full biological assemblies that may include a broader set of atoms from diverse molecules beyond proteins. This enables modeling protein-ligand, protein-RNA, or protein-ion interactions in a generalizable manner. Therefore, our all-atom terminology emphasizes the inclusion of all relevant atomic context in biomolecular complexes—not just the geometric features derived from protein atoms—enabling applications in more complex and realistic biological environments. We will revise the introduction and related work to more clearly explain this distinction.

Novelty and Technical Contributions of ADFLIP

While previous works such as UniIF have explored inverse folding in all-atom settings, the objectives differ. UniIF is designed as a general framework for modeling all biological structures, including proteins, RNAs, and small molecules, rather than focusing specifically on protein design. In contrast, ADFlip is specific for protein design tasks, particularly for proteins that interact with ligands, nucleotides, and other biomolecules. Our work introduces several novel contributions:

1. Generative Modeling with Discrete Flow Matching: ADFLIP is, to our knowledge, the first inverse folding framework that applies discrete flow matching to protein sequence design, taking into account protein complex context. This is particularly advantageous when handling input structures with high uncertainty, which are common when flexible ligands or RNA components are present, where traditional autoregressive models struggle. Our experiments demonstrate that ADFLIP outperforms autoregressive baselines on both single-structure and multi-structure inputs.
2. Training-Free Classifier Guidance for Conditional Generation: We propose a novel, training-free classifier-guidance approach for conditional generation with flow matching. Existing guidance techniques often require retraining regressors or classifiers to accept noisy intermediate inputs $s_t$. Instead, we use a pre-trained classifier/regressor (e.g., AlphaFold) that operates on clean inputs $s_1$, and approximate $p(y | s_t)$ by taking the expectation over $p(s_1 | s_t)$, as described in Eq. (line 308). This allows us to plug in any pretrained predictor without retraining, enabling flexible and efficient conditional design.
3. Support for Multiple Structural States: ADFLIP supports input ensembles with multiple structural conformations, capturing dynamic aspects of protein complexes. This is crucial for modeling biological systems where a single static structure may be insufficient. In summary, ADFLIP introduces a general-purpose, generative approach to inverse folding in full biological assemblies by (1) leveraging a broader all-atom context, (2) introducing discrete flow matching for improved sample quality and uncertainty modeling, and (3) enabling flexible, training-free guidance using pretrained models. We hope this response clarifies the distinctiveness of ADFLIP and addresses your concerns. We thank you for your valuable feedback, which will help in improving our manuscript.

[1]Krishna, Rohith, et al. "Generalized biomolecular modeling and design with RoseTTAFold All-Atom." Science 384.6693 (2024): eadl2528.