Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem

Brian L. Trippe; Jason Yim; Doug Tischer; David Baker; Tamara Broderick; Regina Barzilay; Tommi Jaakkola

Source-linked AI summary

Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem

Brian L. Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, Tommi Jaakkola

arXiv:2206.04119v2Jun 8, 2022q-bio.BMcs.LGstat.ML

TL;DR

Motif scaffolding requires stable structures around functional motifs, but existing machine-learning methods are limited in scaffold length or diversity. The paper combines the 3D equivariant backbone diffusion model ProtDiff with the conditional sampler SMCDiff. It reports diverse scaffolds up to 80 residues and provides a large-compute conditional-sampling guarantee, while noting limits on generalization beyond training motifs.

Problem
Motif scaffolding remains open because existing machine-learning methods are limited to short scaffolds or struggle to produce multiple diverse scaffolds.
Method
The paper trains ProtDiff, an E(3)-equivariant protein-backbone diffusion model, and uses SMCDiff to condition its samples on a desired motif.
Results
The method generates diverse motif-supporting scaffolds up to 80 residues and provides structurally self-consistent backbones evaluated with AlphaFold2.
Takeaways & Limitations
The framework demonstrates the potential of machine learning for realistic protein design settings requiring diverse, longer motif scaffolds.
Takeaways & Limitations
The reported procedure does not indicate generalization to motifs absent from the training set, and experiments were limited to a small set of proteins.

Abstract

from arXiv · show

Construction of a scaffold structure that supports a desired motif, conferring protein function, shows promise for the design of vaccines and enzymes. But a general solution to this motif-scaffolding problem remains open. Current machine-learning techniques for scaffold design are either limited to unrealistically small scaffolds (up to length 20) or struggle to produce multiple diverse scaffolds. We propose to learn a distribution over diverse and longer protein backbone structures via an E(3)-equivariant graph neural network. We develop SMCDiff to efficiently sample scaffolds from this distribution conditioned on a given motif; our algorithm is the first to theoretically guarantee conditional samples from a diffusion model in the large-compute limit. We evaluate our designed backbones by how well they align with AlphaFold2-predicted structures. We show that our method can (1) sample scaffolds up to 80 residues and (2) achieve structurally diverse scaffolds for a fixed motif.

1 INTRODUCTION

Motif scaffolding seeks stable structures that support functional motifs, but existing machine-learning methods struggle with scaffold length, diversity, and conditional generation. The paper introduces ProtDiff and SMCDiff to address these challenges.

Motif scaffolding creates stable structures that support functional motifs, with applications including vaccine and enzyme design.
Existing machine-learning approaches either limit scaffolds to about 20 residues or fail to generate multiple diverse candidates.
The framework evaluates generated backbones through ProteinMPNN sequence design and AlphaFold2 structure prediction, comparing TM-score and motif RMSD.
ProtDiff learns a distribution over realistic protein backbone structures, while SMCDiff performs conditional sampling around a specified motif.
SMCDiff provides exact conditional samples in the large-compute limit when the diffusion model matches the data distribution.
80 residues is the reported scaffold length achieved on a test case, exceeding the 20-residue boundary described for prior methods.

2 PRELIMINARIES

The paper represents proteins through amino-acid sequences and 3D C-α backbones, then uses diffusion processes to learn and sample backbone structures.

A protein is represented by an amino-acid sequence and C-α backbone coordinates, with residues divided into functional motif and scaffold sets.
Denoising diffusion probabilistic models gradually add noise to data until the final distribution is isotropic Gaussian.
The learned reverse process approximates data generation through transitions from noisy structures toward clean structures.

3 PR O TDI F F: A DIFFUSION MODEL OF PROTEIN BACKBONES IN 3D

ProtDiff uses a translation- and rotation-equivariant graph neural network to model protein backbones directly in 3D while respecting their geometric structure and sequence order.

ProtDiff implements the diffusion network as a translation- and rotation-equivariant graph neural network tailored to protein backbones.
The architecture is motivated by proteins’ geometric invariance to coordinate systems and the potential benefits of geometric inductive biases.
ProtDiff represents every residue as a node in a fully connected graph and models pairwise residue interactions through equivariant layers.
Sequence-order encodings are included so the model can learn the linear-chain constraints of protein backbones.

4 SMCDI F F: CONDITIONAL SAMPLING IN DIFFUSION MODELS BY PARTICLE

SMCDiff casts conditional diffusion sampling as a sequential Monte Carlo problem because direct conditional sampling is intractable. Particle filtering produces increasingly accurate scaffold samples as the particle count grows.

4.2 CONDITIONAL SAMPLING IS A SEQUENTIAL MONTE CARLO PROBLEM: SMCDiff reframes conditional sampling as sequential Monte Carlo and uses particle filtering with resampling to prioritize scaffold states consistent with the motif.
4.1 THE CHALLENGE OF CONDITIONAL SAMPLING AND THE ERROR OF THE REPLACEMENT METHOD: The conditional density involves a high-dimensional integral that is intractable to compute directly.
4.1 THE CHALLENGE OF CONDITIONAL SAMPLING AND THE ERROR OF THE REPLACEMENT METHOD: The replacement method approximately conditions diffusion by forward-diffusing the motif and replacing the motif at each reverse step.
4.1 THE CHALLENGE OF CONDITIONAL SAMPLING AND THE ERROR OF THE REPLACEMENT METHOD: The paper shows that replacement introduces irreducible approximation error that cannot be removed merely by making the diffusion model more expressive.
4.2 CONDITIONAL SAMPLING IS A SEQUENTIAL MONTE CARLO PROBLEM: The algorithm takes a motif and particle count as input, diffuses the motif forward, reverse-diffuses particles, resamples them, and returns scaffold samples.
4.2 CONDITIONAL SAMPLING IS A SEQUENTIAL MONTE CARLO PROBLEM: The number of particles K controls a trade-off between computational cost and statistical accuracy.
4.2 CONDITIONAL SAMPLING IS A SEQUENTIAL MONTE CARLO PROBLEM: As K goes to infinity, SMCDiff’s particle output converges in distribution to the target conditional distribution when the diffusion model exactly matches the forward process.

5 EXPERIMENTS

The experiments evaluate ProtDiff and SMCDiff using AlphaFold2-based self-consistency and motif RMSD, showing long, diverse motif-supporting scaffolds alongside important limitations in unconditional sampling.

5.1 In silico evaluation of designed backbones: AF2-based evaluation uses ProteinMPNN sequences and maximum TM-score, called scTM, while motif support is measured by motif RMSD.Eight sequences are sampled per scaffold with ProteinMPNN, then evaluated by AF2 without a multiple sequence alignment.
5.2 Motif-scaffolding via conditional sampling: 80 residues: SMCDiff’s lower-quartile motif RMSD stayed below 1 Å for 5trv scaffolds.The paper treats 1 Å as atomic-level resolution and concludes that scaffolding succeeds in this length range.
5.2 Motif-scaffolding via conditional sampling: 34 and 54 residues: SMCDiff generated two dissimilar, designable scaffolds with scTM > 0.5 around the same motif.The differing secondary-structure orderings illustrate structural diversity.
5.2 Motif-scaffolding via conditional sampling: Beyond 50 residues: fixed and replacement failed to generate successful 5trv scaffolds, whereas SMCDiff remained successful by lower-quartile motif RMSD.The comparison uses two naive inpainting baselines.
5.3 Unconditional sampling: 17% of short and 9% of long ProtDiff samples were designable, versus 11.8% across samples from lengths 50–128.Designable means scTM > 0.5; the designable backbones also exhibited diverse topologies.
5.3 Unconditional sampling: 45% of unconditional backbone samples contained at least one incorrect left-handed helix, and most of these had scTM < 0.5.The authors identify failure to handle chirality as one contributor to nondesignable samples.

6 DISCUSSION

The paper presents ProtDiff and SMCDiff as a generative approach for motif scaffolding, while acknowledging limited experimental scope and several modeling and data constraints.

Contributions: ProtDiff models protein backbones in 3D, while SMCDiff generates scaffolds conditioned on a motif.The approach targets diverse motif-supporting scaffolds and is presented as a generative modeling solution to motif scaffolding.
Results: The method is reported as the first to generate diverse scaffolds longer than 20 residues, with computation time reliably on the order of minutes or less.The authors describe these results as demonstrating potential for realistic protein-design settings.
Scope: SMCDiff is generic in principle, but extending it beyond motif scaffolding is outside the present work’s scope.The paper reports an advantage over a replacement method on toy MNIST inpainting, without claiming superiority over state-of-the-art conditional diffusion models elsewhere.
Limitations: The experiments were limited to a small set of proteins, and the procedure is not shown to generalize to motifs absent from the training set.The authors suggest improved protein modeling could provide better inductive biases for generalization.
Limitations: The training set is small, longer-protein scaling is challenging, and standard motif-scaffolding benchmarks do not yet exist.The authors identify collecting biologically important motifs and developing benchmarks as directions for future work.

B PROBLEM ASSUMPTIONS AND MODELING HEURISTICS

The paper states simplifying assumptions for motif scaffolding and describes a coarse-grained, EGNN-based diffusion model with coordinate scaling and conditional-sampling procedures.

Problem assumptions and modeling heuristics: The formulation relies on simplifying assumptions and heuristics that the paper identifies as targets for future methodological development.These assumptions concern the problem definition, representation, scaffold specification, and modeling choices.
Motif sequence and side-chains: Motifs are represented by C-α coordinates, although biochemical function can depend on residue identities and side-chain atom placement.The paper notes that a more complete motif representation would include amino-acid sequence and side-chain coordinates.
Scaffold length and motif placement: Scaffold size and motif-residue indices are assumed known a priori, despite practical cases allowing multiple lengths and placements.Previous methods used brute-force sampling and post-hoc filtering, while future ML methods could generalize this assumption.
Sequence and side-chain modeling: ProtDiff models only backbone coordinates and delegates sequence design to a separately trained ProteinMPNN model.The paper describes joint sequence-and-structure modeling as a more complete but currently unrealized representation.
Backbone model: Each graph node represents a residue with 3D coordinates and features, and the EGNN predicts the diffusion noise through EGCL layers.Edges carry features, while the network updates coordinates and node features across multiple equivariant graph-convolution layers.
Coordinate scaling: Coordinates are divided by 10 to convert Angstroms to nanometers, matching the spatial scale of the stationary unit-variance Gaussian reference distribution.The final diffusion step is reported to be indistinguishable from an isotropic Gaussian.
Conditional sampling: The replacement sampler forward-diffuses the motif, reverse-diffuses the scaffold, replaces motif coordinates at each step, and returns the scaffold trajectory.The appendix separately documents SMCDiff, including its proof and residual-resampling step.

D.1 THE REPLACEMENT METHOD AND ITS ERROR

The replacement method provides approximate conditional samples but introduces irreducible error from the forward process. SMCDiff instead uses sequential Monte Carlo to approximate conditional sampling, with convergence guarantees as particle count grows.

The replacement method’s output is treated as a putative inpainting solution or approximate conditional sample.
The replacement method introduces approximation error intrinsic to the forward process, which greater reverse-model expressiveness cannot eliminate.
SMCDiff decomposes conditional sampling into three stages: exact motif diffusion, approximate reverse-process sampling, and final scaffold sampling.
The non-trivial reverse-process step requires approximating an intractable, high-dimensional integral with weighted particle approximations.
As K goes to infinity, the particle approximations converge weakly to the target conditional distributions, yielding conditional samples in distribution.
Particles are exchangeable but not independent because they share the diffused motif trajectory, so independent approximate samples require multiple algorithm runs.

D.3 PROOFS AND LEMMAS

The proofs formalize SMCDiff’s convergence using particle-filter asymptotics and show when its final sampling step is exact under the proposition’s assumptions. They also establish the relevant weak-convergence result for residual resampling.

Standard particle-filter asymptotics imply that P(1)_K converges weakly to the target distribution as K increases.
The convergence proof represents SMCDiff as a Feynman–Kac particle-filtering model with transition kernels and potential functions.
The convergence result also applies to residual resampling, although the cited proof directly treats multinomial resampling.
The proof relies on agreement between the forward and approximated reverse processes and standard particle-filter convergence results.
The replacement-method error can be non-trivial, with a lower bound that decreases monotonically in σ2.

F TRAINING DETAILS

ProtDiff was trained on filtered single-chain protein structures using an E(3)-equivariant graph architecture and fixed optimization settings.

ProtDiff uses 4 equivariant graph convolutional layers with 256-dimensional node and edge embeddings.
The training set contains 4,269 single-chain PDB proteins of lengths 40–128 residues after filtering structures with resolution worse than 5Å.
Training used Adam for 1,000,000 steps with batch size 16 and learning rate 1e-4.

G ADDITIONAL METRIC DETAILS

The additional metric and experiment details describe self-consistency evaluation, chirality detection, test-case structures, and a harder motif-scaffolding failure case.

Additional metric details: Self-consistency samples eight ProteinMPNN sequences, predicts structures with AlphaFold2, and records the maximum TM-score against the designed backbone.
Additional metric details: Chirality is classified from C-α dihedral angles: 0.6–1.2 radians indicates right-handed helices, while −1.2 to −0.6 indicates left-handed helices.
Additional metric details: Left-handed helices and mixed chirality are failure modes because natural proteins contain right-handed alpha helices and mixed chirality cannot be fixed by simple reflection.
Additional experimental results: Figure 7 shows native proteins 6exz and 5trv in grey with their extracted motifs in orange for motif-scaffolding tests.
Additional experimental results: The additional experiments analyze chain breaks, steric clashes, and incorrect chirality in ProtDiff backbone samples.
Additional motif-scaffolding results: For the RSV motif target, all inpainting methods failed to recapitulate the motif within a motif RMSD of 1 Å in a 62-residue scaffold.

H.2 QUALITATIVE ANALYSIS OF S CTM IN DIFFERENT RANGES

Qualitative examples show that high-scTM backbones closely match their AF2 predictions, whereas intermediate-scTM designs retain only coarse shape similarity and low-scTM designs often fail structurally. The section also describes interpolation and clustering analyses of generated samples.

scTM > 0.9: scTM > 0.9 designs closely match their AlphaFold2 predictions and often have a related PDB example.Figure 10 compares each backbone with its highest-scTM AF2 prediction and the closest PDB chain.
scTM ∼0.5: scTM ∼0.5 designs preserve coarse shape but may differ in secondary-structure ordering and composition.The length 69 example has a closest PDB-chain TM-score of 0.65, while the length 100 example has no matching monomer in PDB.
scTM < 0.25: scTM < 0.25 cases contain many disordered regions, little structural similarity to the original backbone, and no similar PDB chains.The authors identify these as failure cases and expect improved backbone generators to avoid them.
Latent interpolation: 89- and 63-residue backbone samples were used for additional latent interpolations between model samples.These interpolations are shown in Figures 11 and 12.
Clustering: All 92 samples with scTM > 0.5 were compared sequence-independently and hierarchically clustered using 1 - TM-score as the distance metric.Clustering used average linkage and a TM-score threshold of 0.5; Figure 13 displays cluster centroids and selected remaining designs.

I APPLICABILITY OF SMCDIFF BEYOND PROTEINS: MNIST INPAINTING

The authors test SMCDiff beyond proteins by applying it to MNIST inpainting and comparing it with a replacement method. SMCDiff generally produces continuations that complement the conditioning information, but incurs additional computation.

Experimental setup: SMCDiff was evaluated on MNIST by sampling the missing half of occluded digits and comparing against a replacement method.A DDPM was trained on MNIST before the inpainting comparison.
Experimental setup: For each occluded image, 16 samples were drawn from the replacement method and from SMCDiff using 16 or 64 particles.The methods used a fixed forward trajectory for each occluded image.
Results: SMCDiff tended to generate digits that complemented the conditioning information, while replacement more often produced incoherent continuations.Replacement sometimes generated coherent samples, but this occurred less often than incoherent outputs.
Limitations: SMCDiff adds computation overhead proportional to particle count and can be more expensive than replacement, despite producing higher-quality samples.The authors identify harder datasets and improved architectures as future research directions.