MTGIB-UNet: A Multi-Task Graph Information Bottleneck and Uncertainty Weighted Network for ADMET Prediction

1 Introduction

Computer-aided drug design (CADD) has emerged as a key area of focus at the intersection of artificial intelligence and scientific domains [Rentzsch et al., 2019; Zhao et al., 2020; Yang and Du, 2022]. In the costly drug discovery process, approximately half of clinical trial failures are attributed to an insufficient understanding of the absorption, distribution, metabolism, excretion, and toxicity (ADMET) properties of candidate drugs, which are crucial for their efficacy and safety as therapeutic agents [Norinder and Bergström, 2006; Feinberg et al., 2020]. Poor pharmacokinetic properties and unacceptable potential characteristics, such as toxicity, pose significant risks to human health and the environment and are primary reasons for the exclusion of candidate drugs. In this context, ADMET prediction lays the foundation for drug-candidate selection, thereby supporting advancements in CADD [Du et al., 2023a].

Today, machine learning (ML) methods are widely employed for this purpose. Typically, such methods involve capturing compound structures by molecular descriptors such as SMILES strings, molecular topological graphs, molecular fingerprints, etc., and then analyzing them using elaborately designed algorithms in classification and regression tasks to explore potential quantitative structure-activity relationships (QSAR) [Fang et al., 2024b; Du et al., 2023b]. For example, CORAL utilized Monte Carlo models to encode SMILES for predicting pIC50 values [Azimi et al., 2023; Du et al., 2024]. HRGCN+ utilized molecular graphs as the input to capture the intricate QSAR [Wu et al., 2021]. However, when applied to multi-tasks scenarios, their efficacy diminishes significantly. One primary reason for this decline is the pervasive issue of optimization in multi-task prediction datasets [Zhang and Yang, 2021], manifesting in two major aspects.

Intrinsic connections in multi-tasks. In the realm of molecular field, many predictive models focus solely on individual molecular properties (Figure 1(a)), neglecting the interdependencies among these properties [Wang et al., 2024; Xu et al., 2025]. In reality, there are complex interactions between ADMET properties, while single-task models are limited in capturing [Kim et al., 2024]. For instance, studies have shown that a drug's solubility is closely related to its lipophilicity, which in turn affects the drug's ability to permeate cell membranes [Gleeson et al., 2011; Du et al., 2023c]. This relationship directly influences the drug's absorption characteristics and subsequently impacts its distribution, metabolism, and excretion [Swanson et al., 2023] (Figure 1 (b)). Although multitask learning has been effectively implemented in this field, previous studies ignore the potential conflicts between tasks during training, which may induce negative task interference and result in suboptimal model performance [Xu et al., 2017; Guo et al., 2018]. Dynamic weighting strategies appear to be a potential approach to mitigate optimization discrepancies between tasks [Sener and Koltun, 2018; Vandenhende et al., 2021], but optimizing the weights for different tasks requires consideration of the true physical meaning of chemical tasks to be applicable for ADMET prediction.

Intrinsic connections in molecular substructure. Molecular representation plays a crucial role in determining the predictive performance of machine learning models [Gao et al., 2024a]. A single molecule possesses multiple physicochemical properties, each of which is often influenced by distinct substructural fragments. For instance, the -NH2 group is frequently associated with toxicity, while the -COOH group exhibits hydrophilicity. These substructures hold significant chemical relevance, suggesting that identifying substructures unique to specific properties could enhance model interpretability and performance [Fang et al., 2024a; Du et al., 2025] (Figure 1 (c)). Furthermore, the similarity between different substructures may also provide insights into the underlying commonalities among related tasks, such as -OH is both alkaline and has a certain hydrophilicity.

In response to the limitations of existing models, we propose the Multi-Task Graph Information Bottleneck and Uncertainty Weighted NETwork (MTGIB-UNet) for ADMET prediction, a novel framework that carefully considers inter-task and inter-subgraph correlations to enhance the performance of the primary task. This framework significantly improves the accuracy of ADMET predictions. As illustrated in Figure 2, MTGIB-UNet begins by grouping tasks based on their initial similarities. Following this, an improved Graph Information Bottleneck (GIB) theory is employed to capture the core substructures relevant to different tasks, which are then used to further refine the task groupings. Based on these criteria, a select number of auxiliary tasks are chosen to enhance the learning of the primary task. Finally, an uncertainty-weighting method is employed to balance the contributions of the primary and auxiliary tasks, leading to precise predictive outcomes. Our contributions in this work are summarized as follows:

• Introduction of MTGIB-UNet: We present the MultiTask Graph Information Bottleneck and Uncertainty Weighted Network (MTGIB-UNet), a pioneering framework designed to capture inter-task and inter-subgraph correlations, enabling more accurate predictions across multiple ADMET properties.
• Auxiliary Task Grouping: Our approach integrates GIB theory into the model to effectively filter out irrelevant information. By combining this with task similarity, we optimize the selection of auxiliary tasks, and UW strategy is introduced to keep balance between tasks, thereby enhancing the overall performance of the model.
• Validation through Extensive Experiments: We conducted comprehensive experiments to validate the effectiveness of our proposed model on real-world ADMET datasets, demonstrating significant improvements over existing state-of-the-art methods.

2 RelatedWork

2.1 ADMET Property Prediction

In recent years, the rapid development of deep learning has significantly advanced ADMET property prediction. Initially, researchers employed various machine learning (ML) algorithms to predict ADMET properties [Kim et al., 2024]. With the success of deep learning across various fields, deep learning models have shown strong potential in ADMET research [Tao and Abe, 2025]. For example, [Yang et al., 2019] applied a GCNN model to a dataset from Amgen Inc., significantly improving the prediction accuracy for multiple endpoints. [De Carlo et al., 2024] predict ADMET properties from molecule SMILES using a bottom-up approach with attention-based graph neural networks. These methods often fail to capture the complex interdependencies between ADMET properties, which is crucial for accurate prediction, we propose an improved method that integrates task weighting and grouping strategies in multitask learning to better capture these connections and enhance prediction accuracy.

2.2 Multitask Learning

Multitask learning (MTL) leverages knowledge from related tasks to improve performance, optimizing weights through a combined loss function of multiple tasks. While this can enhance predictions, training different tasks on a shared model can increase optimization challenges and risk negative transfer [Caruana, 1997]. Task weighting addresses this by assigning different weights to task losses, suppressing noise and promoting beneficial training signals. Homoscedastic uncertainty [Kendall et al., 2018] and Dynamic Weight Averaging (DWA) [Liu et al., 2019] are common methods to balance task losses and learning speeds. Task grouping can further mitigate these risks by balancing task losses and clustering tasks into subsets. Techniques such as differentiable pruning [Gao et al., 2024b] and status theory with maximum flow [Du et al., 2023a] have been developed to optimize task grouping and model weights simultaneously. However, these methods still face limitations in addressing task conflicts. We build upon the GIB theory with an uncertainty weighting strategy to better manage inter-task dynamics and minimize negative impact.

2.3 Preliminaries

Problem Definition

Given a set of molecular compounds, the goal is to predict multiple ADMET (Absorption, Distribution, Metabolism, Excretion, and Toxicity) properties. These properties correspond to K endpoint tasks T = {t1, t2, . . . , tK}, where each task represents a specific ADMET property. These include both classification (binary) and regression (continuous) tasks.

Definition 1. (IB-Graph) Molecules are represented as a graph G = (V, E), where V is the set of nodes (atoms) and E is the set of edges (bonds).

The GIB method [Yu et al., 2021] identifies and extracts the most relevant substructures for a given task Y. By compressing the graph G to obtaion the core subgraph GIB, which is determined by:

where I(·, ·) denotes the mutual information, and β is a Lagrangian multiplier that controls the trade-off between the prediction and compression.

Definition 2. (UW) The Uncertainty Weighting (UW) strategy [Kendall et al., 2018] dynamically adjusts task weights in multi-task learning based on each task's uncertainty. This strategy utilizes a learnable parameter σ2 k for each task, modifying the loss calculation as follows:

where σ2 k adjusts the weight of each task's loss Lk. The regularization term δk = log (1 + σ2 k ) ensures stability by preventing the assignment of negative weights.

3 Methodology

The architecture of our model is illustrated in Figure 2. The model consists of five components: the Molecular Graph Encoder, Task-oriented Subgraph Extraction, Auxiliary Task Grouping, Task-Centered Gating, and Multi-task Prediction. The Molecular Graph Encoder processes the molecular graphs to obtain atomic embeddings using a three-layer GCN. The Task-oriented Subgraph Extraction module applies the GIB principle to extract task-specific subgraphs. Next, these subgraphs are then utilized in the Auxiliary Task Grouping module to identify and group tasks with similar subgraph features. The Task-Centered Gating module optimizes the contribution of the grouped auxiliary tasks to the primary task. Finally, the Multi-task Prediction module generates predictions for each task, integrating an uncertainty weighting strategy to dynamically adjust the task weights. Detailed explanation of each component is following:

3.1 Molecular Graph Encoder Module

Each node vi is encoded as a 40-dimensional feature vector xi ∈ Rd , where d is the dimensionality. The set of all node feature vectors constitutes the feature matrix X ∈ R|V|×d , where each row corresponds to a node's feature vector.The specific details of the node and edge features are provided in Appendix A.

The feature matrix X is then passed through three layers of GCN [Kipf and Welling, 2017] to update information:

Here, H represents the updated graph.

3.2 Task-oriented Subgraph ExtractionModule

We introduce the GIB theory to extract the task k's core substructure Gk, as defined in Equation (1). Specifically, We optimize the following two items to obtain Gk:

Minimizing: -I (Y; Gk) We consider Pθ (Y | Gk) as the variational estimation of P (Y | Gk). Thus, we derive:

where H(Y) is a constant, and thus omitted in the model optimization process. The proof is given in Appendix B. Minimizing: I(Gk; G) Following the information-bottleneck principle, we encourage the encoder to discard superfluous structures by perturbing node embeddings. For each node i with hidden state Hi we first obtain a retention logit:

where σ(·) is the sigmoid. A Bernoulli mask λi ∼ Bernoulli(pi) is then relaxed with the Gumbel-Sigmoid trick [Maddison et al., 2017] to keep gradients:

The resulting perturbed embedding is

so uninformative nodes (small pi) are replaced by Gaussian noise. This operation yields a pruned subgraph Gk that retains the salient information of G.

Finally, the mutual information I(Gk; G) admits the following upper bound, a detailed proof is in Appendix B.2:

Therefore, the following objective function must be optimized:

where Lpre is the cross-entropy (classification) or MSE (regression) loss, and LMI is the Kullback-Leibler (KL) divergence between the extracted core subgraphs and random noise graphs.The parameter β acts as a factor to balance the contributions of these two loss components. Finally, the subgraph Gk is pooled by the Set2Set network [Vinyals et al., 2016], yielding the representation vector zk:

This vector serves as a compact representation that encodes the essential information of the subgraph.

3.3 Auxiliary Task GroupingModule (ATG)

Task Similarity Calculation

First, we calculate the similarity between each auxiliary task and the primary task by comparing their core subgraph representations zpri and zk. The similarity score Sk could be calculated as:

Here, sim(·, ·) denotes the cosine similarity. A higher Sk indicates a greater similarity.

Auxiliary Task Selection

Tasks are ranked in descending order by their similarity scores Sk. A threshold τ = 0.7 is applied to obtain a filtered set Tfiltered, which includes tasks with scores above the threshold:

To ensure an adequate number of auxiliary tasks, a minimumm = 2 is enforced. The final number of selected tasks Nselected is the maximum of m and the size of the filtered set:

The primary task representation zpri and the selected auxiliary task representations zk (where k ∈ Tfiltered) are then passed to the next module.

3.4 Task-Centered GatingModule

The task-centric module aims to optimize the contribution of each auxiliary task to the primary task. Each gating network in the module consists of a multi-layer feed-forward neural network and a sigmoid activation function. The input of each network is zk of the auxiliary task, and it outputs two scalar weights: (1) θk, for the primary task and (2) 1 - θk, for the auxiliary task.

The final embedding for the primary task is obtained by summing the weighted embeddings from all auxiliary tasks:

This combined embedding ˆ zpri is then passed through fully connected layers to predict the primary task. The embeddings of the other auxiliary tasks zk remain unchanged.

3.5 Multi-task PredictionModule

Each task's embedding zk is processed through a task-specific three-layer fully-connected network to generate the corresponding predictions:

This formula captures the process of generating the final predicted value for each task. These predicted values are then used to calculate the task-specific losses.

Uncertainty Weighting

Considering potential optimizing conflicts in multi-task training, an uncertainty weighting strategy is applied, as described in Definition 2. This strategy adjusts the weight of each task dynamically, reflecting the uncertainty associated with each task.

Loss Function

The total loss Ltotal for multi-task learning is calculated by aggregating the losses from all tasks. The loss L for each task is defined as shown in Equation (9). After introducing the uncertainty weighting strategy, the total loss function is expressed as:

where α scales the primary-task loss, σ2 k are learnable uncertainty parameters that adjust the weight of each task's loss based on its uncertainty, and δk = log (1+σ2 k ) acts as a regularization term to maintain the stability of the learning process. Lpri and Lk represent the losses for the primary task and auxiliary tasks, respectively.

4 Experiment

4.1 Experimental Setup

Datasets. In this study, we utilized the dataset constructed by [Du et al., 2023a], consisting of 24 endpoints, including 18 classification tasks and 6 regression tasks, collected from 8 published studies related to ADMET properties. It comprises 43,291 drug-like compounds, with 28,153 compounds involved in classification tasks and 16,545 in regression tasks. Depending on the specific tasks, each molecular may be associated with one or more endpoint labels [Yang et al., 2018; Wang et al., 2020; Delaney, 2004].

Evaluation Metrics. For the regression tasks, including Caco-2 permeability, PPB, LD50, IGC50, ESOL, and logD7.4, the coefficient of determination (R2 ) is used as the evaluation metric, covering a total of 6 tasks. For the remaining 18 classification tasks, including HIA, OB, P-gp inhibitor, P-gp substrates, BBB, CYP1A2 inhibitor, and others, the area under the receiver operating characteristic curve (ROC-AUC) is used as the evaluation metric.

Baselines. We compare our model with several state-of-theart models including: Uni-Mol [Zhou et al., 2023], ChemBFN [Tao and Abe, 2025], and MTGL [Du et al., 2023a], MT-GCN [Montanari et al., 2019] and MGA [Peng et al., 2020], along with single-task variants like ST-GCN and STMGA [Montanari et al., 2019; Peng et al., 2020]. Detailed descriptions are in Appendix C.

Implementation Details

In the model implementation, the input feature dimension is set to 40, and the core computations are performed via three GCN layers of dimensions 64, 128, and 128.

During training, the Adam optimizer was employed with a learning rate of 0.001 and a weight decay of 10-5 . The batch size was set to 128, and the model was trained for 200 epochs.

Additionally, the UW coefficient for the primary task, α, was set to 1.2. All experiments were conducted on a Tesla V100PCIE-16GB GPU. Each experiment was repeated eight times, with the mean and variance reported.

4.2 Model Performance

Table 1 presents the quantitative results for ADMET property prediction using various baseline models. MTGIB-UNet consistently outperformed other models, achieving superior results in both classification and regression tasks. Specifically, MTGIB-UNet demonstrated significant improvements across 14 classification tasks and 5 regression tasks, surpassing the second-best model by 0.79 % and 0.85 %, respectively. Overall, MTGIB-UNet secured the first position in 19 out of 24 tasks and the second position in the remaining 5 tasks.

In particular, for the BBB endpoint, MTGIB-UNet improved the ROC-AUC score from 0.973 to 0.979, representing a 0.62 % increase compared to MTGL. For the IGC50 endpoint, the ROC-AUC score was elevated from 0.819 to 0.824, marking a 0.61 % improvement. When considering all endpoints, MTGIB-UNet achieved an overall performance enhancement of 1.03% compared to MTGL. Notably, when compared to single-task models like ST-GCN and ST-MGA, MTGIB-UNet demonstrated an average improvement of 7.32 % to 4.28 %, underscoring its superiority within a multi-task learning framework.

4.3 Ablation Studies

We conduct a series of ablation studies to evaluate the contribution of various components in our proposed model, specifically the Task-oriented Subgraph Extraction Module, Uncertainty Weighting Analysis, and Auxiliary Task Grouping Module. The primary results of these studies are summarized in Table 2, more detailed results provided in Appendix D.

Impact of the Task-oriented Subgraph Extraction Module

This module constitutes a critical component of our model. To assess the impact of accurate subgraph extraction on the target prediction properties, we conducted experiments by excluding the GIB module from the model. The results, presented in Table 2, reveal a notable degradation in model performance across both classification and regression tasks when this module is removed.

Impact of Uncertainty Weighting (UW)

The UW strategy is designed to dynamically adjust task weights during training, prioritizing tasks based on their uncertainty. We conducted an ablation study by setting uniform task weights to 1, as detailed in Table 2. The results clearly demonstrate that the absence of UW leads to decreased model performance, especially on tasks with higher uncertainty. This underscores the importance of adaptive task weighting in improving overall model efficacy.

Impact of the Auxiliary Task Grouping (ATG) Module

The ATG module is a crucial component of our multi-task learning model. To evaluate its impact, we performed an ablation study by removing the ATG strategy and treating auxiliary tasks as independent, without grouping based on their relatedness. This configuration effectively represents a singletask strategy. The results, summarized in Table 2, indicate a noticeable decline in performance across all tasks when ATG is removed, highlighting the significance of strategic auxiliary task grouping in enhancing model accuracy.

To further investigate the role of the ATG module, we shuffled the learned groupings from the current model in Appendix D. This additional analysis helps to validate the importance of the ATG module.

4.4 Sensitivity Analysis

In this section, we explore the impact of various hyperparameters on the performance of our model.

First, we examine the effect of the parameter β in the GIB layer, which governs the trade-off between preserving task-specific information and reducing redundant information. When β is set to lower values, the model tends to retain more noise, leading to suboptimal performance. On the other hand, higher values of β result in excessive compression of subgraphs, leading to the loss of critical information and, consequently, diminished performance. Optimal performance is achieved at β = 10-3 , where the balance between noise reduction and information preservation is well-maintained, as illustrated in Table 3.

We also conducted a sensitivity analysis on α, which adjusts the weighting of the primary task relative to auxiliary tasks. A higher α increases the model's focus on the primary task, thereby enhancing its influence on overall learning and performance. As α increases, the model's performance initially improves due to this enhanced focus. However, setting α too high leads to diminishing returns, where the model may overly prioritize the primary task at the expense of valuable insights from auxiliary tasks, potentially resulting in a decline in overall performance, as shown in Table 3.

Overall, our sensitivity analysis highlights that carefully chosen values of α and β are critical for optimizing the performance of our model. These parameters significantly influence the model's ability to effectively capture core subgraphs and maintain a balance between different tasks.

4.5 Qualitative Analysis

In this section, we evaluate the interpretability and effectiveness of our model on the primary task Caco-2 permeability and its auxiliary tasks OB and ESOL. We examine the dynamic changes in task weights and total loss over epoch, and analyze the extracted subgraphs for compounds.

As shown in Figures 3 (a) and (b), we examine the core substructures of the compounds clonidine and propylthiouracil across three different tasks. It is evident that the core substructures extracted for the primary task, Caco-2 permeability, bear a closer resemblance to those for the auxiliary task OB. Both tasks are related to important absorption properties, highlighting the role of auxiliary tasks in enhancing the primary task. Interestingly, for the propylthiouracil molecule, the substructures extracted for the primary task, Caco-2 permeability, are more similar to those for the auxiliary task ESOL. The model accurately identifies significant groups like sulfonamide and amide, suggesting that similar substructures can also contribute to enhancing the primary task. Therefore, the enhancement of the primary task is jointly determined by the type of auxiliary tasks and the similarity of the core substructures extracted by GIB.

The model's dynamic weights during training are evident in Figure 3 (c), where the weight of the primary task, Caco-2 permeability, increases rapidly in the early stages, reflecting the model's prioritization of this task. Simultaneously, the auxiliary tasks OB and ESOL receive refined weight adjustments, ensuring balanced multi-task learning. As shown in Figure 3 (d), the total loss decreases consistently over time, with a sharp decline in the first 50 epochs, indicating the model's effective learning and adaptability across tasks.

4.6 Resource Usage Analysis

In this section, we evaluate the computational resource usage of our model on the PPB dataset. As shown in Table 4, although the proposed method requires longer training time compared to other baselines, the total number of parameters remains relatively small. Meanwhile, the performance shows a noticeable improvement over other approaches, indicating that our design achieves a favorable balance between model complexity and predictive accuracy.

5 Conclusion and Broader Impact

In this paper, we address the intricate challenge of ADMET property prediction, a pivotal aspect of drug discovery and development. We introduce a novel multi-task GNN model, MTGIB-UNet, that integrates UW and auxiliary learning strategies to enhance predictive accuracy across multiple tasks. Our model functions in two stages: the first stage employs a grouping method based on structural and task similarity, to identify auxiliary task groups for each primary task. The second stage employs a multi-task learning framework incorporating a graph information bottleneck and UW strategies. The graph information bottleneck effectively filters out irrelevant information by constraining the information flow, thereby reducing noise interference, while the UW strategy dynamically adjusts task weights to mitigate potential conflicts between tasks. Our experimental results demonstrate that MTGIB-UNet outperforms existing models in the ADMET prediction domain, underscoring the effectiveness of our integrated approach. Meanwhile, there remains potential for further refinement as in Appendix E.

Acknowledgements

This paper is partially supported by the Project of Stable Support for Youth Team in Basic Research Field, CAS (YSBR005), Natural Science Foundation of China National Major Research Instrument Development Project (No.12227901), the Innovation Program for Quantum Science and Technology (2021ZD0303303), the National Natural Science Foundation of China (22025304, 22033007, 62072427).

Contribution Statement

Xuqiang Li and Wenjie Du are co-frst authors; Yang Wang and Yang Yang are corresponding authors.