Ambient Proteins - Training Diffusion Models on Noisy Structures

We are very pleased to see the thoughtful and very positive feedback of the Reviewer! In what follows, we do our best to address the Reviewer's questions and potential concerns. In particular:

• We present further improved motif scaffolding results.
• We clarify the ablation of Figure 5 and present additional ablations regarding reclustering.
• We better explain our current choice of the mapping fuction and we provide ablations.

Mapping function

Reviewer: "it’s not clear to me exactly how the mapping function is calculated. Is it empirically calculated separately for each of the bins and so not per-protein but per-bin? What is the reference distribution?"

Excellent question and apologies for not having it clear in the first place. The reference distribution is the distribution of "real proteins", i.e. proteins observed in nature. In the paper, we approximate this with the set of proteins for which AlphaFold is super confident (average pLDDT > 90).

Ideally, each protein would be mapped to each time separately. The mapping should rely on some proxy for the quality of the protein (e.g. pLDDT). In our paper, we did not use a continuous mapping -- instead we discretized proteins into bins based on their average pLDDT and each bin got an annotation time.

Specifically, we used the function f defined by the plddt_to_timestep mapping:

(90, 0), (80, 300), (70, 900)

Proteins with:

• pLDDT > 90 are used throughout all steps (0 to 1000).
• pLDDT in (80, 90] are used from step 300 onward.
• pLDDT in (70, 80] are used only from step 900 onward.

We will now provide ablations on this choice of mapping. We report (designability, diversity) pairs for the short-protein generation benchmark (Figure 4), as it is computationally infeasible to do long protein training for all these studies.

First, we show what happens when we slightly change the function f we used in the paper in two opposite directions.

As shown, our method is relatively robust, as small perturbations to the mapping don’t lead to a very substantial change. In fact, the inflated thresholds lead to a model that belongs to the Pareto frontier of Figure 4 in the paper. Running the same model with scale=0.65 leads to another Pareto point, achieving an (86.6, 0.882) designability-diversity pair.

Continuous mapping functions and more bins

We further present continuous generalizations. We now try:

i) a piecewise linear function. ii) a sigmoid function. iii) more bins

E.g. before all proteins with pLDDT in (80, 90] were used for times 300-1000. Under i), a protein with pLDDT x in (80, 90] will be used for times t>=300 * (90 - x) / (90 - 80), e.g. a protein with pLDDT 82 is used for times t>=240.

By optimizing the mapping function, we can even outperform the results we obtained in the paper.

Clarifications on Figure 5 and Data ablations

Reviewer: "Can you help with my interpretation of the ablations in Figure 5."

Apogologies for not having it as clear in the paper. Figure 5 compares three models:

1. Original Genie2
   Trained on short proteins with the original Genie2 dataset (pLDDT > 80 filtering, no reclustering).

2. Improved Genie2 (Baseline w/o Ambient)
   Trained on longer proteins using the re-clustered data, filtered at pLDDT > 70.

3. Ambient Diffusion
   Trained on longer proteins using the re-clustered data filtered at pLDDT > 70,
   plus the Ambient methodology for utilizing low-quality pLDDT samples.

In summary, this figure fixes the dataset (re-clustered, pLDDT > 70) and ablates the Ambient methodology.

Reviewer: "Can you present the same results on designabilty also?"

We present Designability ad Diversity numbers for Original Genie2, Improved Genie2, Proteina and our Ambient Protein Diffusion below (shown as (designability, diversity)):

As shown, our Ambient Protein Diffusion model massively outperforms all baselines.

Additional novelty numbers

We further report novelty numbers for the models of the paper (in short and long generation) to truly show that our model can generate unique proteins:

Data Ablations

Reviewer: "is it possible to infer the impact of the re-clustering of AFDB in isolation?"

We agree with the Reviewer that we need to separately ablate the need for having a re-clustered dataset. Results below:

As shown, reclustering leads to improved diversity.

(Bonus): Improved motif scaffolding results

In their summary, the Reviewer mentions: "[the authors] consider the motif scaffolding benchmark from Genie2, showing reasonable performance - albeit not as clearly state-of-the-art"

We investigated this and we noticed that the scaffold length is not fixed for our baselines, and this length impacts the number of unique successful structures.

In general, for a given motif, the models produce higher numbers of unique successes for longer lengths. In the paper, the numbers reported for Ambient are using the same lengths as Genie 2. We generated samples using 1.5x the average length Genie 2 uses for each motif, and the number of sucesses icreases from approx. 2K to approx 4K (almost doubles!).

Regarding inconsistencies: For 5TPN, RFDiffusion sets the length to be between 39 and 99 residues, while Genie 2 generates samples between 50 and 75 residues. For 1QJG, Proteína sets the length range to be 95-152, while Genie 2 sets it to 53-103. The lengths used for the other motifs are not provided for either RFDiffusion or Proteína. We cannot fairly compare the performance without sampling the same lengths.

Results for 1.5x legths are shown below:

We believe that our rebuttal strongly addresses the Reviewer's questios. If so, we would be very grateful if the Reviewer considers further reinforcing their support for our work. We remain at the Reviewer's availability if there are any remaining concerns.

We are pleased that the Reviewer appreciated the novelty of our approach, the motivation, and the strong experimental performance of our method. The Reviewer has concerns regarding: i) the theoretical validity of our loss, ii) limited ablations about a) our mapping function and b) our dataset reclustering, and iii) missing evaluations on secondary structure. We provide clarifications and numerous ablations and evaluations to address these concerns.

• We provide a derivation of our loss function from first principles.
• We ablate the choice of the mapping function by providing:
  • Robustness analysis to mispecification
  • More bins
  • Continuous generalizations
• We ablate the dataset reclustering and the Ambient loss separately.
• We show improvements over the prior state-of-the-art on secondary structure evaluation.

Loss Derivation

Reviewer:: "Provide a short formal derivation of the ambient loss [...]"

Let's use $t_n$ to denote the $\epsilon$-merging time. Let also $p_0$, $q_0$ be the high-quality and low-quality distributios, respectively. Let's also use R.Vs $X_0, Y_0$ to denote samples from $p_0, q_0$. We know that $\mathrm{KL}(p_{t_n} || q_{t_n}) \leq \epsilon$. The goal is to learn the function $\mathbb E[X_0 | X_t=\cdot]$.

If we had samples from $p_{t_n}$, then it is possible to learn the function of interest, $\mathbb E[X_0 | X_t=\cdot]$, without ever seeing a sample from $p_0$. This result has been proven in the paper "Consistent Diffusion Meets Tweedie" in the context of diffusion models and relates to the Noisier2Noise method from the Signal Processing community. The idea is that we can show the following identity for all times $t \geq t_n$:

$\mathbb E[X_0 | X_t=x_t] = (1 - c_t) \mathbb E[X_{t_n} | X_t=x_t] + c_t x_t$ (for some $c_t$ that is a function of the noise level).

(see Consistent Diffusion Meets Tweedie, p14, Eq. A28.)

The proof of this identity is quite easy: we can apply Tweedie's formula twice for the R.V. $X_t$, one expressing it as a fuction of $X_0$ and one expressing it as a function of $X_{t_n}$. Using this identity and the definition of the conditional expectation as the best denoiser in l2 sense, we arrive at the objective we have in the paper.

The only issue in all this discussion is that we don't have samples from $p_{t_n}$, we have samples from $q_{t_n}$. Hence, we can only learn $\mathbb E[Y_0 | Y_t = x_t]$ through $\mathbb E[Y_{t_n} | Y_t = x_t]$ for all times $t \geq t_n$. How far is $\mathbb E[Y_{t_n} | Y_t = x_t]$ from the desired $\mathbb E[X_{t_n} | X_t = x_t]$? The answer is not much, because of the $\epsilon$-merging. We can exactly bound this by using Girsanov's theorem that expresses the KL error as an error of the conditional expectations (For a reference, see the work Sampling is as easy as learning the score, page 11, Equation 5.4 of the arXiv) Hence, since the KL error is bounded, the same holds for the errors of the conditional expectations.

We hope this proof sketch makes sense. We will add a detailed derivation in the Camera Ready version.

Secondary Structure Evaluation

Reviewer: "Quantify secondary-structure composition—e.g., helix/strand/loop percentages—mirroring Genie 2’s Figure 2."

Following the Reviewer’s recommendation, we systematically evaluate the Secondary Structure Proportions of our generated proteins. Specifically, we first use biotite (as in Proteina) to assign secondary structure labels to all the generated proteins. Then, we look at the designable PDB files, we calculate the percentages of alpha/beta/loops, and then we average those over all designable PDBs. The results are shown below:

As shown, Ambient Protein Diffusion generates proteins with more beta sheets and loops compared to Proteina and Genie2, illustrating progress towards designing more natural proteins.

SS conditional model

Over the past two months (since the submission), we further trained a model that can be conditioned on a per-residue SS label. As we increase the Classifier Free Guidance strength towards beta sheets, this model generates proteins with more interesting SS structure:

We plan to make these models available together with the code release upon acceptance of our work. We will further add all the Ablations and the Secondary Structure Evaluations in the paper.

More Dataset & Ambient Ablations

Reviewer: "The paper lacks an ablation table to separate their individual contributions."

Figure 5 in the paper directly ablates the effect of using Ambient Diffusion, as it shows the effect of having it and not having it on the same exact dataset (reclustered filtered at pLDDT > 70).

We agree with the Reviewer that we need to separately ablate the need for having a re-clustered dataset. Results below:

As shown, reclustering leads to improved diversity. Combining the reclustered dataset with Ambient gives the best results.

Mapping function ablations

Reviewer: "Add ablations for: alternative f(pLDDT) mappings."

Robustness

In the paper, we used the function f defined by the plddt_to_timestep mapping:

(90, 0), (80, 300), (70, 900)

Proteins with:

• pLDDT > 90 are used throughout all steps (0 to 1000).
• pLDDT in (80, 90] are used from step 300 onward.
• pLDDT in (70, 80] are used only from step 900 onward.

We will now show what happens when we slightly change the function f we used in the paper in two opposite directions.

As shown, our method is relatively robust, as small perturbations to the mapping don’t lead to a very substantial change. In fact, the inflated thresholds lead to a model that belongs to the Pareto frontier of Figure 4 in the paper. Running the same model with scale=0.65 leads to another Pareto point, achieving an (86.6, 0.882) designability-diversity pair.

Continuous mapping functions and more bins

We further present continuous generalizations of our original choice. In the paper, we chose a very rough categorization of the proteins into “high-quality”, medium-quality, and “low-quality” for didactic purposes (L191-200). We now try:

i) a piecewise linear continuous generalization. ii) a sigmoid. iii) more bins

E.g. before all proteins with pLDDT in (80, 90] were used for times 300-1000. Under i), a protein with pLDDT x in (80, 90] will be used for times t>=300 * (90 - x) / (90 - 80), e.g. a protein with pLDDT 82 is used for times t>=240.

As shown, by optimizing the mapping function, we can even outperform the results we obtained in the paper. That said, even a simplistic choice, as the ones we opted for in the paper, can yield significant boosts.

Limitations of pLDDT and AlphaFoldDB

Reviewer: "pLDDT is (sometimes) poorly calibrated".

This is true -- pLDDT is known to be unreliable sometimes. We will acknowledge this limitation more strongly in the paper. That said, our method can use any other metric of quality (other than pLDDT). For the purposes of the paper, we chose to use pLDDT because existing pipelines (Genie2/Proteina/RFDiffusion) were filtering proteins based on that score.

Reviewer: "generator might learn AlphaFold’s statistical prior rather than experimental structure distributions"

If the quality estimation oracle (in this case, pLDDT) is assumed to be perfect, then that's not going to happen because AlphaFold's bad predictions and hallucinations will be treated as nosiy data. In fact, that's the problem that previous methods are suffering from and our method fixes, at least partially. That said, pLDDT is not perfect (see above) and some biases might remain.

We believe that our rebuttal strongly addresses the Reviewer’s concerns, and we would be grateful if the Reviewer considers upgrading their score to reflect this.

We thank the Reviewer for their time and thoughtful feedback. The Reviewer appreciated the novelty of our method, our strong experimental results and the potential relevance of our approach for scientific domains with noisy data. The only concerns seem to be about i) missing ablation studies about the mapping function and ii) missing secondary structure evaluation.

In what follows, we address these raised issues by:

• providing numerous ablations about the mapping function
• evaluating comprehensivel our secondary structure performance
• Providing novelty numbers
• Presenting further improved Secondary Structure results using a conditional model.

For all our ablations on the mapping function, we report (designability, diversity) pairs for the short-protein generation benchmark (Figure 4 of the paper), as it is computationally infeasible to do long protein training for all these studies.

Mapping function ablations

Robustness

Reviewer: "there is insufficient exploration of the sensitivity/robustness".

In the paper, we used the function f defined by the plddt_to_timestep mapping:

(90, 0), (80, 300), (70, 900)

Proteins with:

• pLDDT > 90 are used throughout all steps (0 to 1000).
• pLDDT in (80, 90] are used from step 300 onward.
• pLDDT in (70, 80] are used only from step 900 onward.

We will now show what happens when we slightly change the function f we used in the paper in two opposite directions.

As shown, our method is relatively robust, as small perturbations to the mapping don’t lead to a very substantial change. In fact, the inflated thresholds lead to a model that belongs to the Pareto frontier of Figure 4 in the paper. Running the same model with scale=0.65 leads to another Pareto point, achieving an (86.6, 0.882) designability-diversity pair.

Continuous mapping functions and more bins

We further present continuous generalizations of our original choice. In the paper, we chose a very rough categorization of the proteins into “high-quality”, medium-quality, and “low-quality” for didactic purposes (L191-200). We now try:

i) a piecewise linear continuous generalization. ii) a sigmoid. iii) more bins

E.g. before all proteins with pLDDT in (80, 90] were used for times 300-1000. Under i), a protein with pLDDT x in (80, 90] will be used for times t>=300 * (90 - x) / (90 - 80), e.g. a protein with pLDDT 82 is used for times t>=240.

As shown, by optimizing the mapping function, we can even outperform the results we obtained in the paper. That said, even a simplistic choice, as the ones we opted for in the paper, can yield significant boosts.

Reviewer: "Is there a principled way to optimize binning or the function?"

Excellent question. Indeed, there is a principled way to find the mapping function, but implementing it might add some additional engineering steps. Essentially, the mapping function predicts the minimum amount of noise we need so that some distance (e.g., TV or KL) between the two distributions becomes less than epsilon. We can train a discriminator network to estimate this distance by seeing how much noise we need to add to confuse the discriminator. In our paper, we did not opt for this more principled solution because we were afraid it might hinder the wider adoption of the method, but we can certainly have this discussion for the Camera Ready.

Secondary Structure Evals

Reviewer: "have the authors examined generated structure topologies in more detail?"

Following the Reviewer’s recommendation, we systematically evaluate the Secondary Structure Proportions of our generated proteins. Specifically, we first use biotite (as in Proteina) to assign secondary structure labels to all the generated proteins. Then, we look at the designable PDB files, we calculate the percentages of alpha/beta/loops, and then we average those over all designable PDBs. The results are shown below:

As shown, Ambient Protein Diffusion generates proteins with more beta sheets and loops compared to Proteina and Genie2, illustrating progress towards designing more natural proteins.

Reviewer: "What steps could mitigate the bias towards alpha helices?"

Overall, with unconditional generation is really hard to fully capture the underlying distribution. This problem is known for image generative models too, e.g. in the recent NVIDIA paper “Guiding a diffusion model with a bad version of itself” the authors mention:

“"Furthermore, the unconditional case tends to work so poorly that the corresponding quantitative numbers are hardly ever reported. The EDM2-S model trained with ImageNet-512, for example, yields a FID of 2.56 in the class-conditional setting and 11.67 in the unconditional setting."”

The obvious solution is thus conditioning. Over the past two months (since the submission), we trained a model that can be conditioned on a per-residue SS label. As we increase the Classifier Free Guidance strength towards beta sheets, this model generates proteins with more interesting SS structure:

We plan to make these models available together with the code release upon acceptance of our work. We will further add all the Ablations and the Secondary Structure Evaluations in the paper.

Reviewer: "This could limit the diversity of discovered proteins and undercuts one of the paper's main claims of achieving unprecedented diversity."

Besides the SS evaluation results shown above, we further report novelty numbers for the models of the paper (in short and long generation) to truly show that our model can generate unique proteins:

We believe that our rebuttal strongly addresses the Reviewer’s concerns, and we would be grateful if the Reviewer considers upgrading their score to reflect this.

Dear Reviewer,

The end of the rebuttal is fast approaching. The authors would like to know if their new responses satisfy you further. I encourage you to please respond to their message at your earliest convenience.

We thank the Reviewer for their time and feedback. In what follows, we present numerous ablations. We report (designability, diversity) pairs for the short-protein generation benchmark (Figure 4), as it is computationally infeasible to do long protein training for all these studies.

Mapping Function Ablations

Reviewer: "How does the mapping function influence the results?"

Great question – we thank the Reviewer for raising this.

Robustness

In the paper, we used the function f defined by the plddt_to_timestep mapping:

(90, 0), (80, 300), (70, 900)

Proteins with:

• pLDDT > 90 are used throughout all steps (0 to 1000).
• pLDDT in (80, 90] are used from step 300 onward.
• pLDDT in (70, 80] are used only from step 900 onward.

We will now show what happens when we slightly change the function f we used in the paper in two opposite directions.

As shown, our method is relatively robust, as small perturbations to the mapping don’t lead to a very substantial change. In fact, the inflated thresholds lead to a model that belongs to the Pareto frontier of Figure 4 in the paper. Running the same model with scale=0.65 leads to another Pareto point, achieving an (86.6, 0.882) designability-diversity pair.

Continuous mapping functions and more bins

We further present continuous generalizations of our original choice. In the paper, we chose a very rough categorization of the proteins into “high-quality”, medium-quality, and “low-quality” for didactic purposes (L191-200). We now try:

i) a piecewise linear continuous generalization. ii) a sigmoid. iii) more bins

E.g. before all proteins with pLDDT in (80, 90] were used for times 300-1000. Under i), a protein with pLDDT x in (80, 90] will be used for times t>=300 * (90 - x) / (90 - 80), e.g. a protein with pLDDT 82 is used for times t>=240.

As shown, by optimizing the mapping function, we can even outperform the results we obtained in the paper. That said, even a simplistic choice, as the ones we opted for in the paper, can yield significant boosts.

Weight Ablations

Reviewer: "Can you present diffirent choices of the weighting factor w(t)?"

We ablate a constant weighting function (w=1) as used in the Genie2 paper and the square root of the weighting used in the paper.

As shown, the weighting we picked in the paper yields the optimal results. The reason this reweighting is needed is due to the factor multiplying the network prediction in the loss of Eq. (2). The reweighting we proposed corrects for this multiplication, following the derivations in the paper:
Elucidating the Design Space of Diffusion Models (EDM)
— Appendix, p. 27, Section B6.

We will detail the derivation in the Camera Ready version of our work.

Figure 5 clarifications

Reviewer: "In line 318, did you remove the low-quality data?"

Figure 5 of our paper compares three models:

1. Original Genie2
   Trained on short proteins with the original Genie2 dataset (pLDDT > 80 filtering, no reclustering).

2. Improved Genie2 (Baseline w/o Ambient)
   Trained on longer proteins using the re-clustered data, filtered at pLDDT > 70.

3. Ambient Diffusion
   Trained on longer proteins using the re-clustered data filtered at pLDDT > 70,
   plus the Ambient methodology for utilizing low-quality pLDDT samples.

In summary, this figure fixes the dataset (re-clustered, pLDDT > 70) and ablates the Ambient methodology.

Reviewer: "In Figure 5, the ablation plot lacks a direct comparison with Proteína, which performs well on long proteins."

The comparison (designability, diversity) is shown below:

As shown, our Ambient Protein Diffusion model massively outperforms Proteina in long protein generation.

More Dataset & Ambient Ablations

Reviewer: "No systematic ablations separate the effects of the dataset from the low-pLDDT utilization method"

Figure 5 in the paper directly ablates the effect of using Ambient Diffusion, as it shows the effect of having it and not having it on the same exact dataset (reclustered filtered at pLDDT > 70).

We agree with the Reviewer that we need to separately ablate the need for having a re-clustered dataset. Results below:

Reclustering leads to improved diversity.

Other clarifications

Reviewer: "why do you filter structures with pLDDT <70?"

It is true that we could use proteins with lower pLDDTs too. That said, proteins with pLDDT ~70, are already used for only 10% of the diffusion training (steps from 900-1000). Hence, we don’t expect a dramatic improvement by using proteins with pLDDT<70 with our method, and we skipped those to avoid the whole AFDB dataset (estimated size ~28 terabytes).

Reviewer: " Training distinct models for short and long proteins [...] could be a bit strange"

There seems to be a misconception here. Our Genie2 baseline is only trained (and optimized) for short protein generation. Proteina also uses different models to produce their short-generation and long-generation results. In fact, in Proteina there are four different models trained for short generation (Table 1 from https://arxiv.org/pdf/2503.00710). Proteina finetunes one of the four short generation models (the MFS^{no-tri}) for long generation and reports results in Figure 8 of their paper. Similarly, we fine-tune the short-protein generation model for long generation, and we report results in Figure 3 of our submission.

Reviewer: "[...] you do not need to change the loss function [...]"

Excellent question – let us clarify it here.

The Reviewer proposes predicting x0. For low pLDDT structures, x0 cannot be trusted. Instead, our loss involves predicting x_{t_n}, which can be trusted more than x0 because the noise has erased some of the AlphaFold prediction mistakes. So in terms of loss, the prediction of x_{t_n} instead of x_0 is actually critical. For completeness and to support our argument, we did one more ablation during the rebuttal where we trained a model with the same setup as the short-protein generation in the paper, but using x0 prediction loss instead of the proposed loss. Results:

Moving on to the Reviewer’s point about the sampling of times: is it necessary to first noise to $x_{t_n}$ and then to $x_t$ instead of noising all the way directly?

The answer is surprisingly yes (as we hint in the paper L208-209 and in our pseudocode Algorithm 1, Appendix page 6). The reason has to do with the information leakage about x0 that happens as we do multiple epochs. When the corruption to noise level x_t_n happens once, as we do in our paper, x0 cannot be fully recovered – some part of the signal has been destroyed by noise. However, if we are allowed to get multiple noise versions of it at the same noise level then we can average them out and approximately recover the signal. That’s why in a multiple epochs setting the two ways of getting x_t are not equivalent.

We believe that our rebuttal strongly addresses the Reviewer’s concerns, and we would be grateful if the Reviewer considers upgrading their score to reflect this.