A Regressor-Guided Graph Diffusion Model for Predicting Enzyme Mutations to Enhance Turnover Number

This work aims to develop a method for generate enzyme variants that enhance the turnover number ($k_{cat}$). We formulate it as an inverse protein folding problem with an additional optimization objective during the sampling stage with a regressor guidance signal. Given a protein structure represented by its backbone coordinates $\boldsymbol{X}^{pos}={x^{pos}_{1},\ldots,x^{pos}_{n}}$, where $n$ is the number of amino acids, our task is to predict a set of feasible amino acid sequences $\boldsymbol{X}^{aa}={x^{aa}_{1},\ldots,x^{aa}_{n}}$ that can fold into the given structure, while simultaneously identifying sequences likely to exhibit improved $k_{cat}$ values compared to the wild-type enzyme. We approach this problem by modeling the conditional probability distribution $p(\boldsymbol{X}^{aa}|\boldsymbol{X}^{pos})$, which represents the likelihood of amino acid sequences given the backbone structure. To incorporate the optimization of $k_{cat}$, we introduce a regressor function $g_{\eta}(\boldsymbol{X}^{aa},\boldsymbol{X}^{pos})\rightarrow k_{cat}$ that predicts the turnover number for a given sequence and structure. To solve it, we develop $k_{cat}$Diffuser, a regressor-guided graph diffusion model, as depicted in Fig. 1. This model combines a protein inverse folding diffusion model $p_{\theta}(\boldsymbol{X}^{aa}|\boldsymbol{X}^{pos})$ that generates diverse amino acid sequences compatible with the given backbone structure, and a regressor $g_{\eta}$ that guides the sampling process towards sequences with potentially higher $k_{cat}$ values. In our framework, we represent the protein structure as a graph $G=\{\boldsymbol{X},\boldsymbol{A},\boldsymbol{E}\}$, where $\boldsymbol{X}$ includes positional, amino acid type, and physicochemical property information, $\boldsymbol{A}$ is the adjacency matrix, and $\boldsymbol{E}$ captures edge features. This graph representation allows us to capture the complex spatial relationships within the protein structure. By integrating these components, $k_{cat}$Diffuser aims to generate enzyme variants that not only maintain the desired protein structure but also exhibit enhanced catalytic efficiency as measured by $k_{cat}$.

To implement $k_{cat}$Diffuser, we represent proteins as graphs $G=\{\boldsymbol{X},\boldsymbol{A},\boldsymbol{E}\}$ by converting PDB files into this format [7]. The node features $\boldsymbol{X}$ are defined as:

where $\boldsymbol{X}^{pos}$ denotes the 3D coordinates of the $\alpha$-carbon, $\boldsymbol{X}^{aa}$ is a one-hot encoded vector of amino acid type, and $\boldsymbol{X}^{prop}$ represents physicochemical properties. Edge attributes $\boldsymbol{E}$ capture spatial and chemical relationships between connected residues, i.e.

where $d_{ij}$ is the distance between residues $i$ and $j$, $\Delta pos_{ij}$ is their relative position, and $\phi_{ij}$ encodes dihedral angles. The graph construction utilizes a k-nearest neighbor (kNN) method, establishing connections between amino acids within a 30Åradius to preserve the protein’s tertiary structure while creating a computationally tractable representation. The corresponding adjacency matrix $\boldsymbol{A}$ is defined by

For better graph representation, we further incorporate protein backbone information, including dihedral angles ($\psi$, $\phi$) and secondary structure elements (ss), encoded as:

In the protein inverse folding diffusion model $p_{\theta}$, the diffusion process gradually adds noise to the amino acid types $\boldsymbol{X}^{aa}$ over $T$ timesteps, transforming them from the original sequence to a uniform distribution. This process is defined by a forward transition probability $q(\boldsymbol{x}_{t}|\boldsymbol{x}_{t-1})$, where $\boldsymbol{x}_{t}$ represents the noisy amino acid types at timestep $t$. The reverse denoising process, parameterized by $\theta$, aims to recover the original sequence $\boldsymbol{X}^{aa}$ through iterative refinement

where, $p_{\theta}(\boldsymbol{x}^{aa}|\boldsymbol{x}_{t})$ is predicted by an equivalent graph neural network (EGNN) as the denoising network that takes as input the noisy protein graph and additional structural information including backbone dihedral angles and secondary structure elements. To accelerate sampling, we employ the Denoising Diffusion Implicit Models (DDIM) [34], which allows for larger step sizes in the sampling process:

where $T$ controls the time step. The multi-step posterior $q(\boldsymbol{x}_{t-k}|\boldsymbol{x}_{t},\boldsymbol{x}^{aa})$ is computed using the cumulative transition matrices

where $Q_{t}$ represents the transition matrix at time $t$, and $\bar{Q}_{t}$ is the cumulative transition matrix up to time $t$ In this equation, $\text{Cat}(\cdot)$ denotes the categorical distribution, while $\odot$ represents element-wise multiplication (Hadamard product). The fraction inside $\text{Cat}(\cdot)$ represents the unnormalized probabilities for the categorical distribution, with $Q^{T}_{t}$ being the transpose of the transition matrix at time $t$.

We introduce a regressor-guided diffusion sampling scheme that combines a $k_{cat}$ regressor with the protein inverse folding diffusion model to efficiently generate enzyme variants with potentially improved $k_{cat}$ values.

The regressor employs a Transformer-based architecture, comprising 3 output layers and 3 Transformer encoding layers. It utilizes an input dimension of 20, a hidden dimension of 64, and 4 attention heads to capture diverse input features effectively. For the $k_{cat}$Diffuser, we adopted a learning rate of $0.0005$ and a dropout rate of $0.1$ to mitigate overfitting. To ensure training stability, gradient clipping was applied with a threshold of $1.0$. The denoising network in $k_{cat}$Diffuser is a EGNN with 6 layers, each with a hidden size of 128 units, and incorporates embedding layers with a dimension of 128. The training process involves a diffusion sequence length of 500 time steps. To enhance the model’s robustness and account for sequence variability, we introduced BLOSUM-based noise to the input data [7]. This noise injection simulates natural amino acid substitutions, potentially improving the model’s generalization capabilities for enzyme mutation prediction [7].

To train $k_{cat}$Diffuser, in addition to the CATH dataset, we also leverage BRENDA enzyme database, which includes EC numbers, organisms, enzyme sequences, simplified molecular-input line-entry system (SMILES) representations of substrates, and $k_{cat}$ values [26]. We focus on enzyme-substrate pairs to align with our model’s objective of optimizing $k_{cat}$. While $k_{cat}$Diffuser requires protein structures as input, BRENDA primarily provides enzyme sequences. To bridge this gap, we employ ESMFold [27] to predict 3D structures from these sequences, resulting in 15,603 enzyme structures. We divide this dataset into 12,482 enzymes for training, 1,560 for validation, and 1,561 for testing. These structures are then converted into graph representations $G=\{\boldsymbol{X},\boldsymbol{A},\boldsymbol{E}\}$ using our pre-processing pipeline. To ensure data quality and computational feasibility, we filter out empty files and structures larger than 10MB from both CATH and BRENDA datasets before pre-processing. We investigate the impact of incorporating the BRENDA dataset by training $k_{cat}$Diffuser on two configurations: CATH dataset only and combined CATH and BRENDA datasets in the experiment section. While training configuration differs, we evaluate both configurations on the BRENDA test set.

The first evaluation metric is $\Delta\log k_{cat}$, which quantifies the improvement in the enzyme’s turnover number. A higher value indicates more enhancement in catalytic efficiency. Then, we assess the model’s ability to generate sequences similar to the native protein using the recovery rate. To assess the structural quality of generated sequences, we use ESMFold to predict their 3D structures and compare them to the original crystal structures. We evaluate foldability using three metrics: pLDDT, a confidence measure for per-residue structural accuracy [28]; RMSD (Root Mean Square Deviation), which measures atomic-level differences between model and native structures [29]; and TM-score (Template Modeling score), which assesses global structural similarity, correlating strongly with overall model quality [30, 31]. These metrics provide a comprehensive evaluation of the generated sequences’ ability to maintain the desired protein structure while potentially exhibiting improved $k_{cat}$ values, aligning with the core objectives of $k_{cat}$Diffuser.

The results, summarized in Table I, demonstrate the effectiveness of our approach across multiple metrics. Specifically, $k_{cat}$Diffuser achieved the highest improvement in enzyme turnover number, with a $\Delta\log k_{cat}$ of 0.209. This represents an enhancement over ProteinMPNN (0.117) and PiFold (0.087), while GraDe-IF showed a slight decrease (-0.057), demonstrating the effectiveness of our regressor-guided diffusion approach in optimizing enzyme activity. Our $k_{cat}$Diffuser also outperformed all baselines in in terms of recovery rate, achieving 0.716. This is higher than PiFold (0.473), GraDe-IF (0.406), and ProteinMPNN (0.342). The high recovery rate indicates that $k_{cat}$Diffuser generates sequences that closely resemble the native protein while still introducing beneficial mutations. In terms of structural quality, $k_{cat}$Diffuser maintained high fidelity while improving enzyme activity. Our model achieved a pLDDT score of 92.515, slightly lower than PiFold (92.968) but higher than ProteinMPNN (92.038) and GraDe-IF (89.165), indicating high confidence in the local structural accuracy of our generated sequences. $k_{cat}$Diffuser achieved the lowest RMSD of 3.764, better than all baselines, suggesting that our generated structures closely align with the native structures at the atomic level. Furthermore, our model attained the highest TM-score of 0.934, indicating excellent global structural similarity to the native proteins. These results demonstrate that $k_{cat}$Diffuser can balance enzyme activity improvement with structural integrity. To further illustrate the performance of $k_{cat}$Diffuser, we conducted a case study across five diverse enzyme classes (Fig. 2). The visual comparison and accompanying metrics demonstrate the superiority of our approach, particularly when using regressor guidance. Across all cases, $k_{cat}$Diffuser consistently achieved higher $\Delta\log k_{cat}$ values, indicating greater improvements in enzyme activity. For instance, in the EC 2.5.1.31 case, $k_{cat}$Diffuser with regressor guidance achieved a $\Delta\log k_{cat}$ of 0.486, outperforming other methods. Importantly, these activity improvements were achieved while maintaining high structural fidelity, as evidenced by the consistently low RMSD values and high TM-scores. The generated structures (cyan) closely align with the original proteins (green), demonstrating $k_{cat}$Diffuser’s ability to optimize enzyme activity without compromising structural integrity.

We conducted a comprehensive analysis of model complexity, comparing $k_{cat}$Diffuser with ProteinMPNN, PiFold, and GraDe-IF. Table II summarizes the results in terms of number of parameters, memory usage, and inference time. $k_{cat}$Diffuser has the highest number of parameters (8.85M) among the compared models. This increased complexity allows $k_{cat}$Diffuser to capture relationships between protein structure and enzyme activity. In terms of memory usage, $k_{cat}$Diffuser (170.0 MB) sits between the memory-efficient PiFold (108.0 MB) and the more memory-intensive ProteinMPNN (237.1 MB). This moderate memory footprint makes KcatDiffuser suitable for deployment on a wide range of hardware configurations, balancing performance with resource requirements. The inference time of $k_{cat}$Diffuser (5.23s) is notably higher than the other models, which range from $0.26$s to $0.60$s. This increased computational cost is primarily due to the iterative nature of the diffusion process and the additional computations required for the regressor-guided sampling. However, this trade-off in speed enables $k_{cat}$Diffuser to perform multi-site mutations and optimize for enzyme activity, capabilities not present in the faster models.

To explore the impact of regressor guidance on $k_{cat}$Diffuser’s performance, we conducted an ablation study by varying the regressor guidance strength parameter $\lambda$. Table III presents the results of this study, demonstrating the trade-off between enzyme activity improvement and structural integrity. For lower values of $\lambda$ (0.1 to 1.0), we observe relatively stable performance across all metrics. The model maintains high structural fidelity, as evidenced by consistent pLDDT scores around 92.4, low RMSD values (3.64), and high TM-scores (0.937). The recovery rates are also highest in this range (0.730-0.733), indicating that the generated sequences closely resemble the native proteins. As $\lambda$ increases to 5.0, we see an improvement in $\Delta\log k_{cat}$ (0.209), suggesting enhanced enzyme activity. This comes with a slight decrease in recovery rate (0.716) and marginal changes in structural metrics, indicating a good balance between activity improvement and structural preservation. However, further increases in $\lambda$ (10.0 and 20.0) lead to a decline in performance across all metrics. The $\Delta\log k_{cat}$ decreases, and we observe a marked deterioration in structural integrity, particularly at $\lambda=20.0$ (RMSD of 16.390 and TM-score of 0.613). This ablation study reveals that moderate regressor guidance ($\lambda=5.0$) yields the best results, optimizing enzyme activity while maintaining structural stability.