Fast non-autoregressive inverse folding
with discrete diffusion

Problem formulation. The task of inverse folding is to design a sequence $\bm{s}\in\mathcal{V}^{L}$ that folds into a given backbone structure $\bm{x}\in\mathbb{R}^{L\times 4\times 3}$ where $\mathcal{V}$ is the vocabulary of 20 amino acids plus a mask token [MASK]and $L$ is the number of residues. Superscripts $\bm{s}^{(t)}$ will be used to refer to time $t\in[1,\dots,T]$ while subscripts $s_{i}$ refer to residues $i$ (i.e. index). Here, $\bm{x}$ refers to the atomic coordinates of the 5 canonical backbone atoms $\{N,C_{\alpha},C,C_{\beta},O\}$. It has been well-documented that sequences with as low as 30% sequence similarity can fold into the same structure [10]. Therefore, the goal is to learn a distribution of possible sequences that fold into a structure, $p(\bm{s}|\bm{x})$.

ProteinMPNN. First described in Dauparas et al. [3], we provide a brief summary. Each residue is represented as a node in an attributed graph $\mathcal{G}(\bm{x})=(\mathcal{V}(\bm{x}),\mathcal{E}(\bm{x}))$. $\mathcal{V}(\bm{x})$ encodes geometric features such as orientation and sequence index of each residue; $\mathcal{E}(\bm{x})$ constructs a $k$-nearest neighbor graph based on $C_{\alpha}$ distance and encodes relative geometric features. The neural network uses message passing to learn embeddings of protein geometry through a masked language modeling objective that was first described in Ingraham et al. [8]. Inference is performed by sampling a uniformly random decoding order then autoregressively decoding residues one by one.

We provide a overview of discrete diffusion models (D3PM) described in Austin et al. [1]. Discrete diffusion models are a class of latent variable generative models that are defined by a fixed forward Markov process $q(\bm{s}^{(1:T)}|\bm{s}^{(0)})=\prod_{t=1}^{T}q(\bm{s}^{(t)}|\bm{s}^{(t-1)})q(\bm{s}^{(0)})$ from a starting sequence $\bm{s}^{(0)}$ and sequence of increasingly noisy latent variables $\bm{s}^{(1:T)}=(\bm{s}^{(1)},\dots,\bm{s}^{(T)})$. The goal is to learn a model that parameterizes the reverse process $p_{\theta}(\bm{s}^{(0:T)},\bm{x})=\prod_{t=1}^{T}p_{\theta}(\bm{s}^{(t-1)}|\bm{s}^{(t)})$ where,

$\displaystyle p_{\theta}(\bm{s}^{(t-1)}|\bm{s}^{(t)})=\sum_{\bm{s}^{(0)}}q(\bm{s}^{(t-1)}|\bm{s}^{(t)},\bm{s}^{(0)})p_{\theta}(\bm{s}^{(0)}|\bm{s}^{(t)}).$

(1)

Thus, a model with weights $\theta$ learns how to denoise by predicting $p_{\theta}(\bm{s}^{(0)}|\bm{s}^{(t)})$. During training, we optimize the evidence-weighted lower bound (ELBO) between $q(\bm{s}^{(t-1)}|\bm{s}^{(t)},\bm{s}^{(0)})$ and $p(\bm{s}_{t-1}|\bm{s}_{t})$. For absorbing state diffusion, the ELBO reduces to the cross-entropy loss over masked residues,

$\displaystyle\mathcal{L}=\mathbb{E}_{q(\bm{s}^{(0)})}\left[\sum_{t=1}^{T}\frac{T-t+1}{T}\mathbb{E}_{q(\bm{s}^{(t)}|\bm{s}^{(0)})}\left[\sum_{i=1}^{L}\mathds{1}\{s^{(t)}_{i}=\texttt{[MASK]}\}\log{p_{\theta}\left(\hat{s}^{(0)}_{i}|\bm{s}^{(t)}\right)}\right]\right].$

(2)

Following Bond-Taylor et al. [2], each inner expectation is re-weighted to take into account the difficulty of denoising steps. Discrete diffusion depends on the noising process using a categorical distribution over each residue $q(s^{(t)}_{i}|s^{(t-1)}_{i})=\text{Cat}(s^{(t)}_{i};\bm{p}=s^{(t-1)}_{i}\bm{Q}_{t})$ that is parameterized by the probabilities $\bm{p}$ of each category where $\bm{Q}_{t}$ is a doubly stochastic transition matrix with $s_{i}^{(t-1)}$ represented as a one-hot encoding. Each residue is noised independently. We can derive a closed form for sampling the $t$-step marginal and reverse step,

	$\displaystyle q(s_{i}^{(t)}|s_{i}^{(0)})$	$\displaystyle=\text{Cat}\left(s^{(t)}_{i};\bm{p}=s^{(0)}_{i}\overline{\bm{Q}}_{t}\right),\quad\text{where}\quad\overline{\bm{Q}}_{j}=\prod_{j=1}^{t}\bm{Q}_{j}$		(3)
	$\displaystyle q(s^{(t-1)}_{i}|s^{(t)}_{i},s_{i}^{(0)})$	$\displaystyle=\text{Cat}\left(s^{(t-1)}_{i};\bm{p}=\frac{s^{(t)}_{i}\bm{Q}_{t}^{\top}\odot s^{(0)}_{i}\overline{\bm{Q}}_{t-1}}{s^{(0)}_{i}\overline{\bm{Q}}_{t}(s^{(t)}_{i})^{\top}}\right).$		(4)

We choose to use the absorbing state transition matrix,

$\displaystyle\bm{Q}_{t}=(1-\beta_{t})\mathds{I}+\beta_{t}\mathds{1}e_{m}^{\top}$

(5)

where $\beta_{t}$ controls the masking rate, the vocabulary is extended with a mask token, and $e_{m}$ is a vector with 1 on the index of the mask token and 0 elsewhere. On each forward step, a percentage of the unmasked tokens transition to mask and stay masked until the final step, where all tokens are masked. The reverse process performs the opposite with unmasking starting with residue as [MASK].

ProteinMPNN’s autoregressive sampling can be seen as an absorbing state diffusion that unmasks one residue on each reverse step with a uniformly random decoding order. Given this perspective, we sought to extend ProteinMPNN to be more efficient with discrete diffusion. In the context of inverse folding, diffusion is conditioned on the structure, $p_{\theta}(\bm{s}_{t-1}|\bm{s}_{t},\bm{x})$. Borrowing techniques from Bond-Taylor et al. [2], we first fine-tune ProteinMPNN with non-autoregressive diffusion training and strided sampling. Second, we utilize more informed sampling orders based on a purity prior [14].

Our results are presented in Table 1. Using the same model, we evaluate three different sampling variants based on whether purity order and strided sampling $(\nabla t>1)$ are used. Our baseline does not include purity or strided sampling. Using purity to determine the decoding order results in improvement over recovery and designability but a drop in diversity. Including strided sampling gives a 10 times improvement in speed. In comparison to ProteinMPNN, we find sequence recovery is worse with diffusion but designability remains on par with a slight drop in diversity. The advantage with diffusion is clear: a 23 times speed up with a small drop in diversity and designability.