Enhancing Multimodal Protein Function Prediction Through Dual-Branch Dynamic Selection with Reconstructive Pre-Training

1 Introduction

Protein function prediction has become a key challenge in biology, with the rapid development of bioinformatics [Hasselgren and Oprea, 2024]. The Gene Ontology (GO) framework [Ma et al., 2025] standardizes protein functions into three categories: biological process (BPO), molecular function (MFO), and cellular component (CCO) [Aleksander et al., 2023]. In recent decades, numerous deep learning methods [You et al., 2021; Zhang et al., 2023] have been developed to predict protein functions. However, using singlemodal features often faces data limitations [Kulmanov and Hoehndorf, 2020]. Many studies [Fan et al., 2020] have shown that using protein sequence information significantly improves the accuracy of MFO. Still, many proteins share functional similarities but have dissimilar sequences [Lin et al., 2024]. As a result, for proteins with low sequence similarity, the accuracy of predictions may be compromised. Moreover, structure-based methods usually perform better, but the high complexity of protein structures and data acquisition costs limit their application [Paysan-Lafosse et al., 2023]. Furthermore, the noise introduced during the generation of protein-protein interaction (PPI) networks [Wang et al., 2022] through high-throughput techniques poses risks to the accuracy of predictions [Chen and Luo, 2024].

∗Co-corresponding authors: Xiaopeng Jin, Jie Wen.

Therefore, integrating these different types of protein data and taking advantage of their complementary advantages in functional prediction is an important way [Zhao et al., 2024] to improve the performance of protein function prediction. These methods mainly adopt two strategies: graph neural networks (GNNs) [You et al., 2021] and autoencoders [Gligorijevi´c et al., 2018; Fan et al., 2020; Pan et al., 2023]. Graph2GO [Fan et al., 2020] integrates sequence similarity and PPI networks using GNNs, treating protein sequences and structures as node features. However, those using GNNs [Zhou et al., 2019] may amplify noise and face issues with over-smoothing. To address these limitations, CFAGO [Wu et al., 2023] introduces Transformer-based fusion within autoencoders to enhance multimodal feature integration.

However, current multimodal approaches mainly fuse information without exploring the potential complementarity between different modalities. To address this issue, we propose a multimodal method for protein function prediction that efficiently mines the complex internal relationships among spatial structure features, such as PPI networks, subcellular locations, and protein domains, as well as sequence features, specifically the amino acid sequence. Furthermore, due to the complexity of protein information, existing models tend to ignore the detailed features inside the information, such as PPI local network topology, connection strength, amino acid frequency distribution, and key sequence fragments. We add a reconstruction pre-training step to obtain more low-semantic and fine-grained features from protein information of multiple modes. By learning these basic features, the model provides a richer representational basis for downstream tasks.

In addition, large language models play an important role in improving protein function prediction. Inspired by large language models, the protein sequence information in our method is extracted using the pre-trained ProtT5 [Elnaggar et al., 2021]. In this work, to better learn multimodal information, our proposed DSRPGO model includes a shared and an interactive learning branch. In the shared learning branch, we concatenate features from different modalities and perform joint analysis in a unified representation space. Moreover, we introduce the Bidirectional Interaction Module (BInM), where each modality both influences and receives information from others, enhancing overall understanding.

Besides, faced with thousands of protein functions, accurately predicting the protein function of a sample remains a challenging issue. Protein function prediction is essentially a complex hierarchical multi-label classification problem. In this situation, we propose the Dynamic Selection Module (DSM) to dynamically select the optimal feature combination for fitting more diverse protein functions. The code and supplementary materials have been open-sourced 1 . Our main contributions can be summarized as follows:

• We propose a multimodal feature-based approach for protein function prediction that overcomes the limitations of single-modality methods, effectively representing protein functional characteristics.
• A reconstructive pre-training phase is designed to make the model capable of learning more low-semantic finegrained features to assist the model in understanding protein function.
• Our proposed BInM incorporates a bidirectional interaction mechanism to promote efficient fusion and information exchange between sequence and spatial features, enhancing the model's ability to capture strong protein information between different modes.
• We construct the DSM that enables the model to adaptively select channel features most relevant to specific functional labels, resulting in enhanced performance.

2 Methodology

Our proposed method efficiently captures multimodal information about proteins through a strategy for two-step training. In the pre-training stage, we use the encoder-decoder model to learn and inject multimodal knowledge. For spatial features including PPI, subcellular location, and protein domains, a Protein Spatial Structured Information (PSSI) encoder-decoder model using the BiMamba blocks is introduced in this stage. To mine sequence features including protein sequences, we design a Protein Sequence Information (PSeI) encoder-decoder model based on the Transformer blocks for pre-training. Then, during our DSRPGO model training phase, we integrate and learn features from multimodal information. The proposed model is primarily divided into two major branches: one is the multimodal shared learning branch (MSL-Branch), and the other is the multimodal interactive learning branch (MIL-Branch). Protein data are processed through these branches to generate several sets of features, which serve as inputs for DSM. Finally, the model dynamic selects the optimal features for the current protein, to enhance performance in protein function prediction. An illustration of our proposed method can be seen in Figure 1.

1 https://github.com/kioedru/DSRPGO

2.1 Reconstructive Pre-training

In the reconstructive pre-training stage, to obtain feature extractors that are good at mining fine-grained features from multi-modal protein information, we utilize the PSSI and PSeI encoder-decoder model for feature reconstruction.

PSSI Encoder-Decoder Learning

The PPI network gets an N × N adjacency matrix by matrix conversion as input to the encoder. Moreover, another input to the encoder is obtained by concatenating the bag-of-words encodings of subcellular location and Protein Domain.

Mamba Preliminaries. Mamba [Gu and Dao, 2023] extends the capabilities of the State-Space Models (SSMs) [Gu et al., 2023] by enabling the transformation of a continuous 1D input xt ∈ R to yt ∈ R via a learnable hidden state ht ∈ R ˆN with discrete parameters ¯ A ∈ R ˆN× ˆN , ¯B ∈ R1× ˆN , and ¯ C ∈ R1× ˆN as follows:

¯ A and ¯B are continuous A and B converted to discrete evolution parameters using a timescale parameter ∆. To process discrete-time sequences sampled at intervals of ∆, SSMs can be calculated using the recurrence formula. ¯ C represents the projection parameters. In addition, the models compute output through a global convolution as follows:

where ˆN is the length of x, and ¯K is a convolutional kernel.

BiMamba Block. Inspired by the selective scan mechanism in Vision Mamba [Zhu et al., 2024], BiMamba Block introduces a novel bidirectional selective scanning mechanism designed for protein data, capturing both the start and end of spatial structure features for enhanced detail and context. Multi-dimensional features are first converted into onedimensional vectors. Features xsp fromPPI, subcellular location, and protein domains are then passed through BiMamba blocks, interleaved with linear layers and residual operations. As shown in Figure 2, forward (FSScan) and backward selective scans (BSScan) extract bidirectional matrix features via positional transformations and reconstructions. Transformed tokens are scanned using Equation 1 to produce new features, with BiMamba's output ˜xsp expressed as:

where the operation ⊙ denotes the Hadamard product.

PSSI Encoder. In this section, we propose a PSSI encoder architecture designed to effectively map high-dimensional input data into a low-dimensional latent space. The PSSI encoder consists of multilayer perceptrons (MLPs), BiMamba block, Linear and Norm layers, which work in concert to extract features from the input data and generate a compact latent representation. Assume that the input feature xh(k) i ∈ RHk i is a high-dimensional vector of the i-th protein, where Hk i represents the feature dimension of the k-th input source. This feature is reconstructed using MLP to output a low-dimensional representation xd(k) i ∈ RD , where D denotes the size of the MLP hidden layer. PSSI Decoder. The architecture of the PSSI decoder is a counterpart to that of the encoder. The PSSI decoder rebuilds the given protein spatial structure information based on the hidden representations output by the encoder. This process involves BiMamba computation and residual operations, optimizing the crossentropy loss function to enhance the performance. After taking the output xd(k) i of the PSSI encoder and passing through the BiMamba block, alternating Linear and Norm layers, we obtain the recovered high-dimensional features ¯xh(k) i ∈ RHk i . The overarching objective of the encoder-decoder architecture is to minimize the sample wise binary cross-entropy loss between the original and reconstructed source features, thereby enhancing the model's predictive accuracy and fidelity in representing protein data. The loss function of PSSI encoder-decoder is:

where N is the number of total proteins, K is the number of input sources, xh(k) ij and ¯xh(k) ij denotes the j-th dimension vector of xh(k) i and ¯xh(k) i .

PSeI Encoder-Decoder Learning

In PSeI encoder-decoder, the transformer block with multihead self-attention (MSA) mechanism [Dosovitskiy et al., 2021] extracts long-distance features from protein sequences.

Then, to further leverage these features, we use the pretrained ProtT5 [Elnaggar et al., 2021] model to parse the protein sequences. To achieve this, we froze the parameters of ProtT5 and connected it to the PSeI encoder for further pretaining.

PSeI Encoder. The PSeI encoder consists of an MLP block and 6 self-attention blocks. The self-attention block includes an MSA computation layer, as well as alternating linear and norm layers, connected through a residual structure. Assuming the input of the self-attention block is ˜ sdi = MLP(shi ), the output feature is ˆ sdi ∈ RD :

where shi ∈ RHi is the i-th input sequence feature of encoder, and Hi is the dimension of input feature. L(x) denotes the fuction of Linear layer, and N(x) denotes the Norm layer.

PSeI Decoder. The PSeI decoder takes the hidden states from the encoder as input, which contains compressed information about the input sequence. To obtain the final protein sequence encoding, we designed the PSeI decoder using a combination of 6 self-attention blocks and one MLP block. Then, the output feature of the PSeI decoder is ˆ shi ∈ RHi . Like the PSSI encoder-decoder, the loss function Lse for the PSeI encoder-decoder also adopts the form of cross-entropy:

where i denotes the sequence input of the i-th protein, j is the j-th dimension vector of the feature map.

2.2 Bidirectional Interaction and Dynamic Selection for Protein Function Prediction

In this section, we apply the encoders sensitive to low semantic features obtained in the pre-training stage to high semantic tasks. Specifically, to improve the performance of protein function prediction, BInM and DSM modules are proposed to capture deep interaction information between multimodal features and dynamically screen the features most suitable for the current task.

Bidirectional InteractionModule

The proposed BInM enhances the model's ability to learn complex patterns by integrating information across modalities. Using cross-attention, it compares query (Q) vectors

Algorithm 1 Dynamic Selection Moudle Procedure

Input: Protein vector Xdsm , Threshold t

Output: Fusion feature after DSM

• 1: Initialize expert weights W ← 0N.
• 2: Compute expert confidence coefficients ˆ p ← Softmax(MLP(Xdsm)).
• 3: Select active experts S ← {Ei|ˆ pi ≥ t}.
• 4: for each experts Ei in S do
• 5: Normalize ˆ p to obtain weights Wi ← ˆ pi ∑Ej∈S ˆ pj .
• 6: end for
• 7: return DSM(Xdsm) ←Concat(Wi · Ei(Xdsm))

with key (K) vectors from the opposite branch, enabling bidirectional interaction. This approach captures interdependencies between branches more effectively, similar to MSA but focused on cross-branch connections.

Therefore, we assume that the features transformed by PPI are represented as x(1) i , and the features obtained from the encoding of subcellular location and protein domains are concatenated to form x(2) i , while the features extracted through the ProtT foundation model for protein sequences are denoted as x(3) i . Subsequently, x(1) i and x(2) i get features with the same dimension after the MLP reconstruction features, and their concatenated feature map ˜ xBi is used as the input of the first branch of BInM. Similarly, xBi , the input to the second branch of BInM, is derived from x(3) i after its transformation through the MLP. In BInM, the input embedded patches F 1 a ∈ RLa×Da and F 2 a ∈ RLa×Da are initially and randomly divided into multiple heads vectors F 1 b ∈ RLa×Db×Hb and F 2 b ∈ RLa×Db×Hb , where Hb is the number of multiple heads.

As shown in Figure 1, F 1 b and F 2 b are converted into queries Q1 (F 1 b ) and Q2 (F 2 b ). The key K1 and value V 1 of F 1 b , and the key K2 and value V 2 of F 2 b are obtained using three generators Q, K, and V. Then, F 1 c ∈ RLa×Db×Hb obtained by cross-attention is defined as:

where the operation T means matrix transpose, the operation ⊗ represents matrix multiplication, and the goal of softmax function is to normalize the F 1 c . Finally, the cross-attention output feature F 1 d ∈ RLa×Da of the first branch is obtained by feature mapping. Similarly, we can get the cross-attention output F 2 d ∈ RLa×Da of the second branch. In this way, the model takes into account not only the meaning of each branch itself but also the relationships with other branch features, resulting in a more complete representation of multimodal data.

Dynamic Selection Module

In the final feature selection stage, we introduce DSM to enhance key features and mitigate the impact of conflicting ones. As illustrated in Algorithm 1 and Figure 1, this module employs an improved Mixture-of-Experts (MoE) strategy based on Masoudnia et al [Masoudnia and Ebrahimpour, 2014]. The MSL-Branch and MIL-Branch each output a single vector with three channels, where the three channels represent PPI, sequence, and subcellular localization combined with domain features, respectively. All six-channel feature maps serve as the input Xdsm = (x1dsm, x2dsm, · · · , xVdsm) for the DSM. The function of DSM is:

where Ej is the experts belonging to the selected expert group S, ˆ pi denotes the confidence coefficient of expert Ei.

Loss Functions

In this work, protein function prediction is modeled as the multi-label classification task. The predictor, constructed from fully connected layers, takes the output features of the DSM as input and produces an M-dimensional score vector of GO terms: Pi = (p1i , p2i , · · · , pMi )). In the context of protein function prediction using GO terms, there are significantly more negative proteins than positive ones in the training set. Consequently, we employ an asymmetric loss [Wu et al., 2023] as the prediction loss L.

where ym i represents the ground truth label for the i-th protein, while pmi denotes the predicted score. The symbols {y+} and {y-} refer to the positive and negative focusing parameters respectively.

3 Experiments

In this section, we present the experimental setup, including the datasets, baseline models, training details, and evaluation metrics. Then we provide an analysis of the experimental results, supported by ablation studies and Davies-Bouldin scores to validate the effectiveness of the model.

Further experiments on the model components, structures, and parameters can be seen in Appendix Sections 1, 2, and 5.

3.1 Experimental Setup

Dataset Settings. We construct our dataset based on CFAGO [Wu et al., 2023]. PPI data comes from the STRING [Szklarczyk et al., 2023] database (v11.5), and protein sequences, subcellular localization, and domain data are from the UniProt [Consortium, 2022] database (v3.5.175). A total of 19,385 proteins are used for pretraining. For fine-tuning, we collect protein function annotations from the Gene Ontology [Aleksander et al., 2023] database (v2022-01-13). The fine-tuning datasets for each GO branch, split by two-time points, including BPO: 3,197 training, 304 validation, 182 testing proteins (45 GO terms), MFO: 2,747 training, 503 validation, 719 testing proteins (38 GO terms), and CCO: 5,263 training, 577 validation, 119 testing proteins (35 GO terms).

More details about sequence similarity and model performance are in Appendix Sections 3 and 6.

Implementation Details. We conduct all experiments on NVIDIA GTX 4090. We set the dropout rate to 0.1 during pre-training, and the model trains for 5000 epochs, with a learning rate of 1e-5 for the first 2500 epochs and 1e-6 for the remaining 2500 epochs. During fine-tuning, we use a dropout rate of 0.3 and train for 100 epochs with the AdamW optimizer. The learning rate is set to 1e-3 for the first 50 epochs and reduced to 1e-4 for the remaining 50 epochs.

Compared Methods. We compare DSRPGO with nine methods, which are categorized into two groups based on their data utilization strategies. Unimodal-based methods: Naive [Radivojac et al., 2013], BLAST[Altschul et al., 1990], GeneMANIA[Mostafavi et al., 2008], Mashup[Cho et al., 2016], and deepNF[Gligorijevi´c et al., 2018]. Multimodal-based methods: Graph2GO[Fan et al., 2020], NetQuilt[Barot et al., 2021], DeepGraphGO[You et al., 2021], and CFAGO[Wu et al., 2023].

Evaluation Metrics. In this study, we evaluate predictive performance using five metrics: micro-averaged AUPR (mAUPR) and macro-averaged AUPR (M-AUPR) [Peng et al., 2021], F1-score (F1) [Wu et al., 2023], accuracy (ACC), and F-max score (Fmax)[Lin et al., 2024], providing a comprehensive assessment of model accuracy and effectiveness.

3.2 Comparison with Unimodal-based and Multimodal-based Methods

Comparision with Unimodal-based Methods. Most of the previous methods are based on unimodal protein features, so to verify the performance of our multimodal-based method, we compare our method with unimodal-based methods. The experimental results are shown in Table 1. DSRPGO significantly outperforms unimodal-based methods across various metrics, except for M-AUPR in MFO. Compared to unimodal methods, DSRPGO improves Fmax by at least 6.4% in BPO, 7.7% in MFO, and 15.5% in CCO. This demonstrates the advantage of integrating multimodal data for protein function prediction.

Comparision with Multimodal-based Methods. To better evaluate our method, we also compare DSRPGO with other state-of-the-art multimodal-based methods, including CFAGO, DeepGraphGO, Graph2GO, and NetQuilt. The detailed results in Table 1 show that DSRPGO generally outperforms these methods. Compared to multimodal methods, DSRPGO improves the Fmax metric by at least 1.9% in BPO, 1.8% in MFO, and 8.6% in CCO. This indicates that DSRPGO's architecture is more effective in learning deep representations among multimodal features, thereby further enhancing overall performance. At the same time, we observe that DSRPGO does not perform optimally in M-AUPR. This is because M-AUPR evaluates each class equally, including those with fewer samples, which may not reflect the model's overall performance. In contrast, m-AUPR aggregates performance across all classes, offering a more comprehensive measure of predictive capability. In addition, we discuss the Structure-based and PLM-based comparison methods, as detailed in Appendix Section 4.

3.3 Feature Effectiveness Analysis

To further evaluate the distinguishing power of the multimodal features extracted by different components of DSRPGO, we use Davies-Bouldin (DB) [Wu et al., 2023] scores. In the calculation of DB scores, GO terms are set as the labels for protein clusters, meaning proteins sharing the same GO term set are grouped into the same cluster. A lower DB score indicates more compact clusters and clearer separation. As shown in Figure 3, DSRPGO components effectively capture multimodal features. Among them, DSM embedding performs best, indicating that DSM successfully integrates inputs from the MIL and MSL branches.

To further analyze the discriminative power of protein features, we visualize them using t-SNE [Chatzimparmpas et al., 2020], as shown in Figure 4. Raw input features (o PPI, o Attribute, o Sequence), which are not pre-trained, show distinct patterns but lack clear clustering boundaries. In contrast, the output of the feature by various modules of DSRPGO achieves better clustering results. Additionally, compared to the output of the feature by CFAGO (cf embedding), DSRPGO demonstrates significantly superior performance.

4 Ablation Studies

In this section, the contributions of each component in DSRPGO are evaluated, as shown in Table 2.

Analysis for Backbone Components. According to lines 1,2, and 3 of Table 2, the results of the backbone network only using MSL-Branch or MIL-Branch are not as good as those using combined branches.

Effectiveness of BInM. Considering the correlation of features among space and sequence, this method uses the BInM block to facilitate bidirectional multimodal feature interaction before DSM. As shown in rows 3 and 4 of Table 2, we verify the validity of BInM for the overall model by removing it.

Effectiveness of DSM. To enable effective feature selection and accurate prediction of protein functions, DSM is used to select channel features most relevant to specific functional labels adaptively. At the same time, it reduces the interference and conflict caused by redundant features. As shown in rows 3 and 5 of Table 2, DSM has a positive impact on protein function prediction.

Impact of Sequence and Spatial Structure Features. To verify the complementarity between sequence and spatial structure features, we perform an ablation study, retaining only spatial structure or sequence features. For the BInM module, it is removed as no interaction occurs with a single feature type. Rows 6 and 7 of Table 2 show that removing feature interaction significantly reduces model performance.

Impact of Pre-training. To evaluate the contribution of pre-training, we conduct an ablation study by removing it. As shown in the last row of Table 2, the model's performance drops significantly across all metrics without pre-training.

5 Conclusion

This paper proposes a dual-branched multimodal method for protein function prediction with reconstructive pre-training. The proposed method enhances the model's ability to integrate multimodal features through two key components: the BInM and the DSM, leading to significant performance gains. Experimental results show that the DSRPGO outperforms current state-of-the-art unimodal and multimodal methods across multiple metrics. These results underscore the importance of integrating multimodal data to enhance protein function prediction, and validate the superiority of the BInM and the DSM in multimodal protein data integration.

Acknowledgements

This work was supported in part by the National Natural Science Foundation of China under Grant No. 62302317, the Natural Science Foundation of Guangdong Province under Grant 2025A1515010184, the project of Shenzhen Science and Technology Innovation Committee under Grant JCYJ20240813141424032 and JCYJ20240813112420027, and the Foundation for Young innovative talents in ordinary universities of Guangdong under Grant 2024KQNCX042, the Stable Support Projects for Shenzhen Higher Education Institutions under grant 20231122005530001 and 20220715183602001, and Guangdong Basic and Applied Basic Research Foundation grant 2024A1515220079.

Contribution Statement

Xiaoling Luo and Peng Chen contributed equally to this work.