PaperlayerView exact revision on arXiv

Source-linked AI summary

SE(3) diffusion model with application to protein backbone generation

Jason Yim, Brian L. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, Tommi Jaakkola

arXiv:2302.02277v3cs.LGq-bio.QMstat.ML

TL;DR

De novo protein backbone design requires models that respect geometric constraints and global rigid-motion invariance. This paper develops theoretical foundations for diffusion on multiple SE(3) frames and implements them in FrameDiff, which generates designable monomers up to length 500 without pretraining. The authors report 75% designability under decreased noise scale, while noting limitations in sampling stability and evaluation comparability.

  • Problem

    Existing protein-generation methods lacked a principled diffusion framework on multiple SE(3) frames with group invariance, while some successful methods relied on heuristic losses or pretraining.

  • Method

    The paper constructs an SE(3)-invariant diffusion process on multiple frames and trains FrameDiff with an SE(3)-equivariant score network without pretraining.

  • Results

    FrameDiff generates designable, diverse, and sometimes novel monomer backbones up to length 500; with decreased noise scale, 75% of samples were designable by scTM>0.5.

  • Takeaways & Limitations

    FrameDiff provides a non-pretrained framework for SE(3) diffusion-based protein backbone generation with reported designability across monomer lengths.

  • Takeaways & Limitations

    The authors note that final sampling steps can destabilize backbones, motivating early trajectory truncation, and they refrain from state-of-the-art claims because training and evaluation differ across methods.

Abstract

from arXiv · show

The design of novel protein structures remains a challenge in protein engineering for applications across biomedicine and chemistry. In this line of work, a diffusion model over rigid bodies in 3D (referred to as frames) has shown success in generating novel, functional protein backbones that have not been observed in nature. However, there exists no principled methodological framework for diffusion on SE(3), the space of orientation preserving rigid motions in R3, that operates on frames and confers the group invariance. We address these shortcomings by developing theoretical foundations of SE(3) invariant diffusion models on multiple frames followed by a novel framework, FrameDiff, for learning the SE(3) equivariant score over multiple frames. We apply FrameDiff on monomer backbone generation and find it can generate designable monomers up to 500 amino acids without relying on a pretrained protein structure prediction network that has been integral to previous methods. We find our samples are capable of generalizing beyond any known protein structure.

1. Introduction

The paper frames de novo protein backbone generation as a geometrically constrained design problem and develops principled, invariant diffusion on multiple SE(3) frames to address it.

  • Motivation: De novo protein design generates proteins satisfying specified structural or functional properties.This work focuses specifically on generating protein backbones.
  • Motivation: Protein backbones must satisfy physical and chemical constraints while remaining designable by at least one amino-acid sequence.The structure–sequence coupling makes computational backbone design technically challenging.
  • Geometric formulation: Backbone structures can be represented as elements of SE(3)^N, while protein function depends on relative geometries invariant to global rigid transformations.This motivates SE(3)-equivariant neural networks for structure-based methods.
  • Research gap: Existing manifold diffusion extensions do not provide the needed theoretical framework for diffusion on multiple SE(3) frames.The paper identifies this as a gap between theoretical foundations and successful deep-learning practice.
  • Contribution: The paper constructs an SE(3)^N diffusion process, characterizes Brownian motion for score matching, and separates rotational and translational components.The resulting framework is implemented as FrameDiff without pretraining.
  • Contribution: FrameDiff generates designable, diverse, and novel monomers up to length 500, with in-silico designability second only to pretrained RFdiffusion.The comparison is against RFdiffusion, which has four-fold more parameters.

2. Preliminaries and Notation

The paper represents each protein residue with a rigid frame and uses Riemannian score-based diffusion to model and reverse distributions on manifolds.

  • Backbone parameterization: An N-residue backbone is parameterized by N orientation-preserving rigid transformations, or frames, mapping fixed atom coordinates.Each frame decomposes into a rotation r_n and translation x_n.
  • Backbone parameterization: Each frame T_n=(r_n,x_n) uses r_n∈SO(3) for rotation and x_n∈R3 for translation, with an additional torsion angle ψ constructing the oxygen atom.The complete collection is T∈SE(3)^N.
  • Manifold diffusion: Riemannian score-based generation samples from p_0 by reversing a manifold-valued stochastic process that approaches an invariant density.The forward process transforms data toward noise before time reversal.
  • Score matching: The time-reversed process depends on the Riemannian score ∇log p_t, which is estimated with a denoising score-matching network.The ideal network minimizer equals the Stein score under the stated flexible-network condition.
  • Geometric setting: Lie groups such as SO(3), SE(3), and R3 combine smooth group operations with manifold structure.This structure supports geometric diffusion formulations beyond Euclidean space.

3. Diffusion models on SE(3)

The paper defines a principled diffusion on SE(3)^N by choosing a metric and reference measure, then enforces global invariance through centered translations and rotational equivariance.

  • Forward diffusion: No canonical forward diffusion exists on SE(3)^N, so the construction requires choosing an inner product and a reference measure.These choices determine the Brownian motion and invariant density.
  • Forward diffusion: The chosen metric identifies SE(3) with SO(3)×R3 from a Riemannian perspective, yielding a well-defined Brownian motion.The rotational and translational components can then be treated independently.
  • Rotational diffusion: The rotation process uses Brownian motion on SO(3), whose transition density is expressed through the heat kernel and rotation angle.For compact Lie groups, the heat kernel is characterized using irreducible characters.
  • Translational diffusion: The translation process uses an Ornstein–Uhlenbeck process converging geometrically to N(0, Id), with an analytically computable conditional score.The SE(3) conditional score factors into independently computed rotational and translational terms.
  • SE(3) invariance: An SE(3)-invariant probability measure cannot exist on SE(3)^N because no translation-invariant probability measure exists on R3N.The paper therefore centers the diffusion by removing the center of mass.

4. Protein backbone diffusion model

FrameDiff models protein backbones as multiple SE(3) frames and learns an SE(3)-equivariant score with a neural architecture combining geometric and sequence-based updates. Its training adds auxiliary structural losses, while sampling uses centered translation and rotation diffusion with truncated, noise-scaled reverse trajectories.

  • Frame representation: FrameDiff models protein backbones as frames T in SE(3)^N, with each frame combining a rotation r_n and translation x_n.The frame representation describes the N−Cα−C atoms of each residue; backbone oxygen placement additionally uses a torsion angle ψ.
  • FramePred architecture: Each network layer updates node embeddings with IPA and a Transformer, then uses the updated embeddings to update edges and predict frame updates.Initial and post-IPA node embeddings are concatenated before the Transformer, with a skip connection after it.
  • Score and torsion prediction: FrameDiff predicts denoised frames and torsion angles, converting frame predictions into separate rotation and translation score predictions for each residue.The score is trained with denoising score matching on the SE(3) diffusion process.
  • Training objectives: Auxiliary backbone-position and local pairwise-distance losses discourage chain breaks and steric clashes during low-noise training steps.These losses are applied when t < T_F/4, when fine-grained characteristics emerge.
  • Sampling: Sampling uses Euler–Maruyama geodesic random walks with Gaussian translation and rotation tangent-space noise, removes center-of-mass motion, and returns the prediction at time ϵ.The reverse trajectories may use noise downscaled by ζ, and sampling is truncated early because backbones commonly destabilized in final steps.

5. Experiments

Experiments evaluate unconditional monomer backbone generation by designability, diversity, and novelty, comparing FrameDiff with prior methods under differing training and evaluation setups. FrameDiff generates designable samples without pretraining, with designability–diversity trade-offs controlled by sampling parameters and reported speed advantages at reduced step counts.

  • Experimental setup: FrameDiff was trained on 20,312 filtered PDB backbones using 17.4 million parameters for one week on two A100 GPUs.The experiment uses a four-layer model for monomer backbone generation.
  • Evaluation: Designability is measured by self-consistency between ProteinMPNN-designed sequences and ESMFold-predicted structures, while diversity is the proportion of unique MaxCluster structural clusters.Novelty is measured by each sample’s highest FoldSeek TM-score to any PDB chain, reported as pdbTM.
  • Sampling parameters: Changing the noise scale to ζ = 0.5 improved designability at the cost of lower diversity, while increasing Nseq improved designability with significant compute cost.With Nsteps = 100, a 100-amino-acid backbone took 4.4 seconds on an A100 GPU, reported as more than an order-of-magnitude faster than RFdiffusion.
  • Ablations: The best ablation model included self-conditioning, auxiliary losses, and the DSM form of the SO(3) loss.Using the squared Frobenius norm loss instead caused a slight reduction in designability, and hyperparameter searches were left for future work.
  • Designability and comparison: FrameDiff generates designable samples without pretraining, whereas RFdiffusion demonstrated designable sequences only when initialized with pretrained weights.The authors caution that training data, model size, and evaluation procedures differ across methods.
  • Results: 75% of samples across a range of lengths were designable by scTM>0.5 with decreased noise-scale, while prior non-pretrained methods reporting this metric were below 55%.The authors refrain from state-of-the-art claims because training and evaluation differ across methods.

6. Related work

Related work spans protein diffusion models using different structural representations and general diffusion frameworks on manifolds. FrameDiff differs by developing a principled diffusion model on SE(3) that recovers the exact backward process at optimality.

  • Diffusion models on proteins: Prior protein diffusion models used representations including backbone atoms, torsion angles, and other structure-based parameterizations, with Chroma reporting the highest designability among cited non-pretrained methods.FrameDiff instead develops a principled SE(3) diffusion approach for frame-based protein modeling.
  • Diffusion models on manifolds: General manifold diffusion frameworks extend score-based or maximum-likelihood methods to Riemannian manifolds, while FrameDiff specializes the construction to the SE(3) Lie group.The paper identifies exact backward-process recovery at optimality as a major theoretical and methodological distinction.

7. Discussion

FrameDiff provides an SE(3)-invariant diffusion framework with an equivariant score network that generates designable and diverse monomer backbones without pretraining. Samples can generalize beyond PDB under stringent filters, but novelty claims require experimental characterization and comparisons remain confounded by differing setups.

  • Contributions: FrameDiff combines theoretical foundations for SE(3)^N diffusion with an SE(3)-equivariant, non-pretrained score network for protein backbone generation.The framework is applied to monomer backbones and generates designable and diverse samples.
  • Novelty and scope: Under stringent filters, some samples generalize beyond PDB, but claims of generating novel proteins require experimental characterization.The conclusion limits the interpretation of computational novelty evidence.
  • Comparison and future work: FrameDiff results are competitive with Chroma and RFdiffusion, although differences in training and evaluation confound rigorous comparisons.The paper proposes extending FrameDiff to conditional generation and scaling training data and optimization in future work.

Supplementary to:

The supplementary material develops Lie-group and representation-theoretic tools used to analyze SO(3) and SE(3), including irreducible representations and Lie-algebra structure.

  • Representation theory: A group representation maps group elements into invertible linear transformations, and irreducibility restricts invariant subspaces to zero or the full space.For compact groups, unitary representations decompose into finite-dimensional irreducible components.
  • Lie-group foundations: Lie groups combine smooth manifold and group structures, with Lie algebras given by tangent spaces at the identity.SO(3), SE(3), and R3 are examples of Lie groups.
  • Lie-algebra operations: The Lie bracket measures the noncommutativity of vector fields or Lie-algebra elements and reduces to the matrix commutator for matrix Lie groups.The bracket is [X,Y] = XY − YX when the group is represented by matrices.
  • Representation theory: Irreducible representations of sl(2,C) are constructed on homogeneous polynomials and shown to be equivalent to the representations indexed by m.The basis actions of H, E, and F establish irreducibility and the stated equivalence.
  • Representation theory: The supplementary construction derives explicit actions of H, E, and F on polynomial bases, including the rescaled basis used to identify equivalent representations.The resulting actions are given by ρ(H)(wj) = (−m + 2j)wj, ρ(E)(wj) = (m − j)wj+1, and ρ(F)(wj) = jwj−1.

D.1. Canonical metric on SO(3)

This section constructs the canonical bi-invariant metric on SO(3), characterizes its Lie-algebra operations, and introduces the SE(3) group structure and adjoint representation.

  • Canonical metric on SO(3): The Killing form defines a canonical metric on compact Lie groups, and on SO(3) it is B(X,Y) = Tr(XY).In the basis (Y1,Y2,Y3), B = −2 Id, so −B/2 makes the basis orthonormal.
  • Adjoint representations: The SO(3) Lie brackets satisfy [Y1,Y2] = Y3, [Y2,Y3] = Y1, and [Y3,Y1] = Y2.These relations determine the adjoint and infinitesimal adjoint actions used in the metric analysis.
  • Canonical metric on SO(3): The metric induced by −B/2 on SO(3) is bi-invariant, and its Laplace–Beltrami operator is expressed through the basis of so(3).The associated Riemannian exponential agrees with the matrix exponential in this setting.
  • SE(3) structure: SE(3) consists of orientation-preserving rigid motions and has composition (R,x) ∗ (R′,x′) = (RR′, x + Rx′).It is the semidirect product SO(3) ⋊ R3, not the direct product SO(3) × R3.
  • SE(3) structure: SE(3) is a 6-dimensional Lie group whose Lie algebra has basis {Y1,Y2,Y3,e1,e2,e3}, enabling an explicit adjoint representation.The adjoint representation is computed in this rotational and translational basis.

D.3. Choice of metric and Laplacian derivation

The authors choose a left-invariant metric on SE(3) that separates rotational and translational geometry, yielding factorized gradients, Laplacians, and Brownian motion.

  • Metric choice: The chosen metric uses the identity form on se(3) and combines the SO(3) metric with the Euclidean inner product on translations.This metric is left-invariant but not right-invariant.
  • Metric and Laplacian: Under this metric, the SE(3) Laplacian factorizes as Δ_SE(3)f(T) = Δ_SO(3)f(r,x) + Δ_R3f(r,x).The corresponding gradient also separates into rotational and translational components.
  • Brownian motion: The resulting Brownian motion on SE(3) can be represented using rotational and translational Brownian components.The construction supports independent SO(3) and R3 forward processes conditional on the initial frame.
  • SO(3) heat kernel: SO(3) Brownian-motion densities are represented through heat-kernel expansions and approximated in practice by truncating the series at N = 2000 terms.The conditional score is obtained by differentiating the resulting density with respect to the rotation.
  • Caveats: The metric choice affects Brownian-motion speed and the normalization of the conditional score, while confusing heat-kernel and Brownian densities introduces a factor 1/2.The heat kernel satisfies ∂tqt = Δqt, whereas the Brownian density satisfies ∂tpt = 1/2 Δpt.

E.4. Pytorch implementation of IGSO3, and simulation of forward and reverse process on a toy example

The implementation section provides SO(3) parameterization and simulation primitives, then verifies on a toy example that forward and reverse geodesic random walks have closely agreeing marginals.

  • Toy experiment: The toy experiment uses a discrete target on SO(3), samples uniformly distributed rotations, and simulates processes over T = 4 with 200 time points.The target is formed from three randomly selected atoms on SO(3).
  • Toy experiment: The forward and reverse processes have closely agreeing marginal distributions at each simulated time t.This comparison tests the implemented geodesic random-walk approximation on the toy target.
  • SO(3) implementation: The implementation supplies primitives for logarithmic and exponential maps, rotation angles, tangent-space Gaussians, and Riemannian gradients on SO(3).These operations move between matrix, Lie-algebra, and rotation-vector parameterizations.
  • Density and score: The IGSO3 density is computed from a truncated power series and used to define noised target densities and their Stein scores.The score is obtained as the Riemannian gradient of the log density.
  • Process simulation: Forward and reverse Brownian processes are simulated with a geodesic random walk that combines drift updates with tangent-space Gaussian noise.Reverse simulation uses negative time increments in the same procedure.
  • Scaling rules: Scaling the inner product by α rescales gradients, Laplacians, orthonormal bases, and tangent-space Gaussian samples according to explicit rules.The exponential map remains unchanged under this scaling.

F. Invariant diffusion processes

The paper establishes conditions under which diffusion processes on Lie groups preserve invariance, then uses measure disintegration to characterize invariant measures on SE(3)^N and SO(3).

  • Invariant diffusion processes: H-invariant initial distributions remain H-invariant under the SDE when the diffusion covariance and drift satisfy the stated equivariance conditions.The proof proceeds by showing the transformed density satisfies the same Fokker–Planck equation and using uniqueness of the SDE solution.
  • Invariant diffusion processes: An H-equivariant vector field has H-invariant divergence, supporting invariance of the associated Fokker–Planck operator.The result is established through equivariance of covariant derivatives or, alternatively, the divergence theorem.
  • Invariant diffusion processes: The anisotropic operator ∆Σ preserves the required transformation relation when Σ and the gradient transform compatibly under the group action.When Σ = Id, the result reduces to equivariance of the standard Laplacian.
  • Invariant measures: An SE(3)-invariant measure on SE(3)^N disintegrates into a translation component proportional to Lebesgue measure and an SO(3)-invariant conditional measure.The construction uses translation shifts of the conditional kernel and proves rotational invariance of the centered measure.
  • Invariant measures: Combining an SO(3)-invariant probability measure with Lebesgue measure on R3 yields an SE(3)-invariant measure on SE(3)^N.This characterizes the measure construction used to handle global translations and rotations.

I. Additional method details

The additional details specify how frames and coordinates are parameterized, how FrameDiff updates geometric and sequence features, and how translation and rotation diffusion schedules are selected.

  • Frame and atom parameterization: Backbone coordinates are represented with residue frames, while idealized atom positions are reconstructed through frame2atom and torsion-angle transformations.The frame construction uses backbone atoms and a rigidFrom3Point-style procedure; torsion ψ_n rotates idealized oxygen around the C−Cα bond.
  • FramePred architecture: FrameDiff initializes residue and pair features from residue indices, timestep, and relative sequence distances, then iteratively updates nodes, edges, and frames.Node updates use IPA and a Transformer, edge updates use an MLP, and backbone updates follow AF2’s BackboneUpdate algorithm.
  • Frame and score prediction: After the final network layer, the predicted frame is converted into a score for each residue, with torsion prediction following AF2.The model can also use self-conditioning based on predicted Cα displacements, set to zero half the time.
  • Diffusion schedules: Translation diffusion uses separate drift and diffusion coefficients from rotation diffusion, allowing their rates to be decoupled while approaching the invariant distribution by t = 1.The translation process uses a beta schedule and has a time-rescaled Ornstein–Uhlenbeck interpretation.
  • Diffusion schedules: A logarithmic rotation schedule makes variance decay slower than a linear schedule and was found to produce slightly improved samples.The selected hyperparameters are σ_min = 0.1 and σ_max = 1.5 for rotations.
  • Training objective: The denoising score matching loss is preferred to the Frobenius rotation loss because it is compatible with the theory of score-based generative models and slightly improves designability in ablation.The comparison concerns learning different rotation-related objects in practice.

J.1. Training details

Training uses filtered monomer structures, a memory-conscious batching strategy, and a two-stage sampling procedure that reconstructs atom coordinates from generated frames.

  • Training data: 20,312 monomers remained after filtering 23,913 PDB proteins for length 60–512, resolution below 5 Å, and no more than 50% loops.The authors report that removing high-loop proteins improved training and sample quality.
  • Batched training: FrameDiff batches different diffusion timesteps of the same backbone because fully connected graphs require memory that scales quadratically with backbone length.The implementation uses a simple batching strategy rather than PyTorch Geometric’s efficient batching methods.
  • Optimization: Training uses Adam with learning rate 0.0001 and runs for two weeks on two A100 GPUs.The optimizer parameters are β1 = 0.9 and β2 = 0.999.
  • Sampling: Sampling initializes frames from the invariant reference distribution, reverses diffusion with discretized Langevin dynamics, and re-centers intermediate frames at every step.A final network pass at t = 0 produces frames that are converted into idealized backbone atoms with frame2atom.
  • Designability: Designability evaluation designs multiple sequences with ProteinMPNN, folds them with ESMFold, and compares predicted and sampled backbones.Agreement is measured using scRMSD or scTM.

J.4. Additional results

Additional analyses show that FrameDiff generates designable backbones up to length 500, while stricter scRMSD designability becomes difficult beyond length 400. Samples retain high diversity across lower noise scales, exhibit varied secondary structures with longer backbones tending toward helical composition, and outperform FoldingDiff in direct comparison.

  • Designability: FrameDiff generates backbones up to length 500 meeting the scTM> 0.5 designability criterion.
  • Designability: Reliably achieving scRMSD< 2 designability past length 400 remains difficult and is identified as future work.
  • Secondary structure: Samples span a wide range of secondary-structure compositions, with a preference for helical backbones and mostly helical structures beyond length 400.
  • Diversity: FrameDiff maintains high diversity at lower noise scales, whereas RFdiffusion reports decreased diversity under those conditions.
  • Comparison to FoldingDiff: FrameDiff greatly outperforms FoldingDiff when both are evaluated on proteins up to length 500 using the same dataset.The comparison uses FoldingDiff’s evaluation code and FrameDiff noise scale ζ = 1.0, Nsteps = 500, and Nseq = 8.
Loading 2302.02277v3…