Source-linked AI summary
Fast non-autoregressive inverse folding with discrete diffusion
John J. Yang, Jason Yim, Regina Barzilay, Tommi Jaakkola
TL;DR
Inverse folding seeks sequences that fold into a desired 3D structure, but autoregressive decoding can be slow for large proteins. The paper fine-tunes ProteinMPNN with structure-conditioned discrete diffusion and a purity prior. On CATH, it reports a 23 times speed up over ProteinMPNN, with small drops in diversity and designability.
Problem
Inverse folding requires designing sequences that fold into a desired 3D structure, while autoregressive decoding can be prohibitively slow for large proteins.
Source: §1 · p. 1Method
The approach fine-tunes ProteinMPNN to perform structure-conditioned discrete diffusion, using a purity prior for the residue sampling order.
Source: §1 · p. 1, §2.3 · p. 3Results
23 times speed up over ProteinMPNN, with sequence recovery worse and small drops in diversity and designability.
Source: §3.1 · p. 4Takeaways & Limitations
The approach offers a way to extend ProteinMPNN with conditional diffusion for faster and controllable generation.
Source: §1 · p. 1Because the non-autoregressive approach does not edit previously generated residues, it is prone to error accumulation.
Source: §4 · p. 4
Abstract
from arXiv · showhide
Generating protein sequences that fold into a intended 3D structure is a fundamental step in de novo protein design. De facto methods utilize autoregressive generation, but this eschews higher order interactions that could be exploited to improve inference speed. We describe a non-autoregressive alternative that performs inference using a constant number of calls resulting in a 23 times speed up without a loss in performance on the CATH benchmark. Conditioned on the 3D structure, we fine-tune ProteinMPNN to perform discrete diffusion with a purity prior over the index sampling order. Our approach gives the flexibility in trading off inference speed and accuracy by modulating the diffusion speed. Code: https://github.com/johnyang101/pmpnndiff
1 Introduction
This work targets inverse protein folding, where a sequence must be designed to fold into a desired structure. It fine-tunes ProteinMPNN with discrete diffusion to reduce inference compute while maintaining foldability performance.
1 Introduction
- Inverse folding seeks sequences that fold into a desired 3D structure, the second step in de novo protein design.The preceding step is generating a structure for a desired function.Source: §1 · p. 1
- ProteinMPNN's autoregressive decoding scales linearly and can be prohibitively slow for large proteins, while its uniformly random decoding order may miss higher-order interactions.The paper identifies both as limitations of the widely used inverse-folding method.Source: §1 · p. 1
- 23 times less compute achieves equivalent foldability performance to ProteinMPNN with the diffusion model.The paper presents this as its best configuration.Source: §1 · p. 1
- The framework fine-tunes pretrained ProteinMPNN with diffusion and supports decoding the sequence in strided reverse steps.Figure 1 illustrates decoding the whole sequence in three steps.Source: §1 · p. 1, Figure 1 · p. 2
2 Method
The method treats inverse folding as learning sequences conditioned on a protein backbone, then adapts ProteinMPNN to generate them through discrete diffusion. It uses absorbing-state masking, strided sampling, and a purity-based sampling order.
2.1 Background
- Inverse folding learns a distribution of amino-acid sequences that fold into a specified backbone, while ProteinMPNN represents backbone geometry with a residue graph and message passing.Its standard inference samples a random residue order and decodes autoregressively.Source: §2.1 · p. 2
2.2 Discrete diffusion
- Discrete diffusion adds noise through a fixed forward Markov process and learns a reverse process that denoises sequences conditioned on the structure.For absorbing-state diffusion, training uses cross-entropy over masked residues.Source: §2.2 · pp. 2–3
- Absorbing-state diffusion progressively masks unmasked residues until all are masked, while reverse steps unmask residues starting from [MASK].The masking rate is controlled by βt.Source: §2.2 · p. 3
2.3 Diffusion improvements to ProteinMPNN
- Fine-tuned ProteinMPNN uses absorbing-state diffusion to mask or unmask subsets of residues, with all residues initially masked for sampling.A linear masking schedule targets approximately ⌊tL/T⌋ masked residues at time t.Source: §2.3 · p. 3
- Strided sampling decodes multiple residues per step and reduces the number of sampling steps according to an integer stride interval.At each step, the number decoded is set by the schedule and the selected decoding order.Source: §2.3 · p. 3
- The purity prior uses each residue’s relative prediction confidence to importance-sample which location to mask or unmask next.The authors chose this prior after preliminary attempts to bias the order using coupling strengths were unsuccessful.Source: §2.3 · p. 3
3 Experiments
The experiments evaluate discrete-diffusion variants on the CATH 4.2 Single Chain benchmark, measuring designability, sequence recovery, diversity, and sampling speed. The final diffusion approach is substantially faster than ProteinMPNN, with small drops in diversity and designability.
3 Experiments
- The CATH 4.2 Single Chain evaluation samples eight sequences per structure and treats ESMFold-predicted RMSD as the main designability metric.It also reports sequence recovery, pairwise Levenshtein diversity, and average wall-clock sampling time on one NVIDIA V100.Source: §3 · p. 4
- Diffusion fine-tuning uses T = 100 with ∇t = 20, requiring five model calls per sequence.ProteinMPNN is pretrained for 200 epochs, then fine-tuned on the diffusion objective for 1,000 epochs.Source: §3 · p. 4
3.1 Results
- 23 times faster than ProteinMPNN, diffusion retains on-par designability with a slight drop and has lower sequence recovery and a small drop in diversity.The comparison reports a small drop in designability and diversity alongside the speedup.Source: Table 1 · p. 4, §3.1 · p. 4
4 Discussion
The discussion identifies limitations in training, sequence refinement, and the diffusion formulation. It also reports no observed correlation between purity decoding order and structural characteristics.
4 Discussion
- Training without ProteinMPNN pretraining performs worse than fine-tuning, for reasons the authors say remain unclear.The authors also note that their approach uses discrete time and is incompatible with continuous-time SDEs.Source: §4 · p. 4
- Because generated residues are not revised, the non-autoregressive approach is prone to error accumulation.The authors suggest iterative whole-sequence refinement as a direction for future work.Source: §4 · p. 4
- The authors found no correlation between purity decoding order and structural characteristics, including whether spatially close residues should decode together.They plan further investigation of optimal decoding orders and their relation to structure.Source: §4 · pp. 4–5
A Related work
Related work spans improvements to protein structure representation and sequence generation. The section emphasizes that sequence recovery alone may not reflect designability, which has been linked to experimental validation.
A Related work
- Transformer-based inverse-folding models generate autoregressively, leading to slow inference and high variance in generation quality.Non-autoregressive alternatives have performed strongly on perplexity and sequence-recovery benchmarks.Source: A Related work · p. 7
- Sequence-recovery benchmarks can misrepresent designability: ESM-IF achieved 72% sequence recovery but had poor designability relative to methods with lower recovery [17].The passage notes that multiple works found designability correlates with experimental validation.Source: A Related work · p. 7
- The paper positions its method as a faster, non-autoregressive alternative to ProteinMPNN, a popular model with experimentally validated results.The related-work discussion also describes geometric graph architectures and Transformer-based sequence-generation approaches.Source: A Related work · p. 7