End-to-End Optimized Pipeline for Prediction of Protein Folding Kinetics

Protein folding database (PFDB) [13] was considered for protein folding kinetics (PFK) prediction. PFDB consists of experimental observations for 2S (two adjacent strands of a beta-sheet) and N2S (N adjacent strands of a beta-sheet) proteins. Structural parameters such as amino acid sequence, type of chaining, torsion angles, Ramachandran outliers, etc., were manually derived from RCSB PDB [14] for each protein. These structural parameters were considered because they significantly influence the type and nature of folding [15]. Other individually reported observations of proteins were also collected and used for testing the regressor to test its real-time capabilities.

The data is first passed to the data pre-processing module, which consists of an Outlier Remover ($\alpha$), Temperature Standardization ($\beta$), Feature Extractor ($\gamma$), and Encoder ($\delta$). As the data consists of various features, $\alpha$ ensures the removal of outliers using the percentile-based outlier detection method, specifically the interquartile range (IQR) approach. This method establishes lower and upper thresholds based on the first quartile ($Q1$) and third quartile ($Q3$) values, respectively. The lower threshold is calculated using the formula $Q1-1.5\times IQR$, while the upper threshold is determined as $Q3+1.5\times IQR$. Data points falling above the upper threshold or below the lower threshold are classified as outliers, representing values that significantly deviate from the central distribution of the data.

The Temperature Standardization ($\beta$) step is applied for the temperature feature. $\beta$ standardizes the temperature value to an optimal scale. Temperature standardization is critical when dealing with protein folding data collected at various temperatures. $\beta$ collects the rate constants of folding and unfolding ($ln(k_{f})$ and $ln(k_{u})$, respectively) at a reference temperature of 25°C using the Eyring-Kramers equation [16] which facilitates precise assessments of folding kinetics and enabling more robust conclusions in protein folding research. This correction adjusts the dataset to a standardized temperature, facilitating more accurate and reliable comparisons.

For the numerical features (Lpbd, L, pH, and temperature), $\beta$ applies Z-score normalization as the feature extraction technique. This process subtracts the mean and divides it by the standard deviation across each feature. Z-score normalization ensures that all numerical features are brought to a standardized distribution by eliminating differences in magnitude between them. Additionally, this enables more efficient and unbiased analysis.

For the categorical features, $\gamma$ ensures that all categorical data points are encoded numerically. As mentioned earlier, the M-estimate encoder [17] is used for this purpose. The encoder calculates the probabilities of each column based on their frequency in the dataset. With a $m=20$ value, the encoder effectively down-weights outliers and reduces their impact on the estimation process.

The machine learning algorithm utilized in this study is Bonsai [18], which employs a shallow tree-based approach to process data continuously and incrementally, tailored explicitly for constrained paucity systems. Data optimization and memory utilization is performed by learning the input data at a lower dimension. Bonsai’s streaming implementation allows it to support devices with limited RAM, even those incapable of storing a single vector.

These trees incorporate enhanced nodes, enabling internal and leaf nodes to make non-linear predictions. Bonsai achieves a remarkable ability to capture intricate non-linear decision boundaries by combining the predictions made by individual nodes along the path followed by a data point. In addition to enabling parameter sharing along paths, path-based prediction reduces the size of the overall model, thus increasing its efficiency. It optimizes memory allocation and maximizes prediction accuracy by learning all nodes jointly rather than node by node in a greedy manner. Bonsai offers an effective solution for machine learning tasks through its combination of sparse projection, streaming learning, enhanced nodes, and joint learning. This algorithm was specifically chosen for its low memory footprint, which would further help this pipeline to be integrated onto other fold-specific frameworks.

Feedback-driven Optimization Engine (H) ensured iterative improvement of the pipeline over time. The Performance Analyzer $\eta_{1}$ assessed the quality of predictions, which consists of various evaluation metrics as discussed in Section III-F. Sensitivity Analyzer $\eta_{2}$ facilitated a thorough observation of the impact of variations in input parameters and model architectures on the predictions, leveraging the results obtained from $\eta_{1}$, which compared the predictions against known ground truth values. Additionally, the testing data was split into different batches, and the accuracy was measured for each batch, aiding in the identification of trends and patterns in performance. With the help of $\eta_{2}$, various hyper-parameters were optimized to maximize the model’s performance. H is essentially needed for bonsai to boost its performance as the model primarily prioritizes minimizing the computational resources. Moreover, the iterative optimization process driven by H enabled continuous refinement and adaptation of the pipeline, leading to enhanced accuracy and reliability.

As the number of features directly influences the efficiency of the model, finding the optimal number of features is essential to maximize the performance of the regressor while downsizing the number of feature samples and memory footprint. This would also reduce the computational cost for $\Delta$ and the overall pipeline.
Pearson’s correlation coefficient [19] was used to determine the base relationship between the target variable ($ln(kf)$) and other feature sets. The correlation value (R) ranges from -1 to +1, where -1 (negative correlation) and +1 (positive correlation) signify strong correlation, while values closer to 0 signify weak or no correlation. R was computed, as shown in Equation  5, where $X_{i}$, $Y_{i}$ are individual datapoints and $X_{m}$, $Y_{m}$ is the mean of the X and Y dataset respectively, the absolute was taken to reside in the range of 0 to 1 for more uncomplicated depiction and understanding of the magnitude of correlation vectors as shown in Fig  4.

PFDB served as the primary dataset for training the regressor model (bonsai). PFDB contains a diverse collection of PFK information on both 2S and N2S proteins. A comprehensive train test strategy was implemented to assess the adaptability and generalizability of the trained model. The dataset was split into train and test as this approach ensured that the models were evaluated on unseen data, allowing for a robust assessment of their predictive capabilities. To optimize the model, MAE was iteratively reduced by tuning the hyperparameters. The process of testing the mini-batches and updating the parameters occurs concurrently by multi-threading. This approach reduced the time to compute results and significantly decreased MAE.

The regressor was trained on different mini-batches, each comprising a unique feature subset. The training was conducted on $F_{n}=\{n=2,4,5,6,7,8,9\}$, where $n$ represents the feature subset derived from Pearson’s correlation coefficient analysis. Specifically, the batches $D^{\text{AA}}$, $D^{\text{BB}}$, $D^{\text{AB}}$, and $D^{\text{BA}}$ were trained on the prominent features $F_{\text{best}}^{\text{A}}$, $F_{\text{best}}^{\text{B}}$, $F_{\text{best}}^{\text{A}}\cup F_{\text{best}}^{\text{B}}$, and $F_{\text{best}}^{\text{A}}\cap F_{\text{best}}^{\text{B}}$, respectively. The features are chosen from $F_{n}$ as formulated in Equation  4 and  5.

The metrics actively used to evaluate the bonsai throughout this work are formulated below:

$y_{i}$ = Actual kinetic value of the protein
$\hat{y}_{i}$ = Predicted kinetic value of the protein
$\bar{y}$ = Mean of the actual kinetic values, for $i=1,2,\ldots,n$.
Other evaluation metrics actively used in this work are inference time and model size. These metrics enable informed decisions regarding feature subset selection and model optimization.

As discussed above, an optimal feature subset heavily influences the performance of the bonsai model. The basis behind choosing an optimal feature subset is based on the performance, memory footprint, and inference time while availing the feature subsets. The detailed reasons for this are discussed in the upcoming sections.

Several regressor models, such as Decision Tree (DT), Random Forest (RF), XGBoost (XGB), and LightGBM (LGBM), were trained in the same environment as the bonsai model to compare and contrast the difference in performance among them. Since bonsai employs tree-based regression, five tree-based regressors were chosen for comparison. These models were trained on the most prominent feature set $F_{9}^{B}$, as bonsai is designed to minimize computational resources, taking the best balance model into account would not be an appropriate comparison. Bonsai’s ability to achieve such high-performance gains while maintaining its compactness and accuracy sets it apart. The best regressor model (LGBM), in terms of performance, was evidently outperformed by the bonsai model, which consumed only 9% of the memory consumed by LGBM while being 6% more accurate when evaluated using R2 score. Bonsai was also 156x faster than DT, the model with the lowest inference time (8.736 sec).

Even though the bonsai model’s performance and memory footprint is carefully analyzed and reported, analysis of the pipeline is also fairly significant. As depicted in Fig  5, $\Delta$ consumes the most time (71%), out of which normalization and standardization ($\beta$) take up to 37% and feature extractor consumes 34% individually. The regressor consumes 21% of the pipeline inference time, while H takes up 4%-9% of inference time, which is comparatively insignificant. This pipeline was constructed primarily to consume minimal computational resources while being highly accurate and fast (24.7 ms). This makes the pipeline highly compatible for future integration onto fold-specific prediction frameworks and systems.

The above chart was plotted for the best performing model which was trained on $F_{9}^{B}$. Between the models trained on $F_{9}^{B}$ and $F_{8}^{A}$, $F_{9}^{B}$ was chosen so as to understand and illustrate the computational needs of the best regressor. $F_{8}^{A}$ was not considered as it does not exhibit the best performance but only the best balance between the considered factors, model size, performance, and inference time.