GBD-DRP: Generalizable Bias-Disentangled Learning for Drug Response Prediction with Missing Modalities

1 Introduction

Cancer remains one of the most serious threats to human health, and drug-based therapy is currently the main approach for cancer treatment. Due to the genetic diversity of cancer cell lines, the same drug can show very different responses across different cell lines. As a result, accurately predicting drug response is an important problem and also a key component of future personalized medication. With the development of sequencing technologies and the construction of large-scale databases, machine learning methods have shown clear progress on this task. In particular, resources such as the Genomics of Drug Sensitivity in Cancer (GDSC) [Yang et al., 2012] and the Cancer Cell Line Encyclopedia (CCLE) [Barretina et al., 2012] have enabled data-driven modeling of drug response, and recent deep learning methods have achieved encouraging results [Baptista et al., 2021; Partin et al., 2023].

∗Corresponding author

Due to the irregular structure of drug molecules and the heterogeneity of cell-line features, many approaches adopt graph-based models. For example, GraphDRP [Nguyen et al., 2021] and DeepCDR [Liu et al. ] represent drugs as molecular graphs and use Graph Convolutional Networks (GCN) to encode drug structures. Another line of work applies transformer-based models to drug SMILES representations. Methods such as DeepTTA [Jiang et al., 2022] decompose canonical SMILES strings into functional groups, encode them as vocabularies, and use transformers to obtain drug representations. These models focus on capturing structural information from drugs in a flexible manner.

To further improve prediction accuracy, some studies pay more attention to modeling interactions between drugs and cell lines. Instead of encoding drugs and cell lines separately and simply concatenating their representations for prediction, these methods introduce interaction modeling during the encoding stage. FA2GL [Xu et al., 2025] constructs graphs for both drugs and cell lines and connects them through pipe nodes, allowing the heterogeneous graph to capture drug-cell interactions more explicitly. PASO [Wu et al., 2025] use attention network to get drug representations and applies cross-attention with cell-line features to model the effects of drugs on different biological pathways. In addition, these approaches often incorporate multiple cell-line modalities, including gene expression (GEP), gene mutation profiles (MUT), and copy number variations (CNV), to better describe cell-line-specific responses.

Although these methods perform well under standard evaluation settings, they also introduce new challenges. As discussed by Codice et al. [Codic`e et al., 2025], many preclinical studies show limited generalization when tested on truly unseen cell lines, which are more relevant to real clinical scenarios. This highlights the need to explicitly evaluate and improve generalization to unseen cell lines. Moreover, the use of additional modalities increases the risk of missing data. Compared with gene expression data, far fewer cell lines have complete MUT and CNV profiles. As a result, methods that rely on these features often exclude cell lines with missing modalities during testing. This can introduce bias and further worsen generalization to new cell lines, while also requiring additional data collection before predictions can be made.

To address these issues, we propose GBD-DRP, a generalizable bias-disentangled learning framework for drug response prediction. We decompose drug response into three channels. The first channel captures drug-invariant properties and represents the baseline efficacy of a drug. We model this component using an EdgeGAT-based graph encoder on drug molecular graphs. The second channel represents explicit cell-line-specific effects that can be modeled using available features such as GEP, MUT, and CNV. This branch captures drug responses under specific cellular environments and remains generalizable when features are available. The third channel models residual bias, which includes unobserved mechanisms and noise that cannot be explained by existing features. This component corresponds to domain bias in traditional domain generalization settings. By separating the residual bias from explicit cell-line-specific representations, our method preserves cell-line specificity while improving generalization to unseen cell lines.

For potentially missing modalities, we use the explicit cellspecific branch to process the modality when it is available, treating it as a supplement to the standard cell-line features. When the modality is not available, the residual branch is used to estimate its effect. By applying a counterfactual consistency constraint on the residual branch, the model can leverage existing samples of optional modalities to infer the influence of the missing samples. This allows the framework to utilize incomplete additional modality information effectively.

The main contributions of this work are as follows:

• We introduce a Triple-Branch Disentanglement framework. It separates explicit biological specific features from spurious biases, which helps the model generalize better to new cell lines.
• We provide a solution for missing modalities using a Counterfactual Consistency mechanism. The model can use existing features to estimate the effects of missing modalities through the residual branch, without relying on imputation.
• Experiments shows GBD-DRP outperforms current state-of-the-art methods. And also maintains strong accuracy and robustness, especially for unseen cell lines and cases with missing modalities.

2 RelatedWorks

2.1 Domain Generalization

Domain Generalization (DG) is a method for handling outof-distribution (OOD) data [Zhou et al., 2022]. It aims to train a model on source domains and allow it to generalize to target domains. Methods such as domain alignment, metalearning, and data augmentation are commonly used. DG either models domain shifts directly or reduces the influence of domain-specific features through data or training techniques. One important branch is domain-agnostic learning [Peng et al., 2019]. This approach disentangles input features into domain-invariant true signals and domain-specific bias. The domain-specific bias is then removed, allowing the model to make predictions that do not depend on the domain.

2.2 The Generalization Challenge in DRP

In principle, domain-agnostic learning seems suitable for predicting drug response in new cell lines, because it separates the interference caused by cell-line-specific features. However, drug response is highly specific to cell lines. Removing all domain-specific information often reduces prediction accuracy and can even cause model collapse. In this biochemical setting, it is more important to explicitly model cell-linespecific features to generate accurate predictions for each cell line. But at the same time, focusing too much on cell-linespecific features can increase bias in the model.

To address this dilemma, GBD-DRP further disentangles cell-line-specific features from residual biases that cannot be generalized. By separating these two parts, the model can accurately use cell-line features for new cell lines while removing non-generalizable bias. This design allows us to use domain generalization principles to improve generalization without losing important cell-line-specific information.

3 Methodology

3.1 ProblemFormulation and Latent Decomposition

Let D and C denote the set of drugs and cell lines. For a pair (d, c), the input consists of drug molecular features xd and cellular genomic features xbio c (composed of base and optional modalities). The ground truth response is y ∈ R. The goal of Drug Response Prediction (DRP) is to learn a mapping f : (xd, xbio c ) → y that generalizes to unseen cell identities c /∈ Ctrain.

Standard regression often overfits to cell-specific spurious correlations (e.g., batch effects or unmeasured confounds). To address this, we do not treat the response as a monolithic target. Instead, we formulate the expected response ˆ y as a sum of three disentangled components, motivated by the separation of global pharmacological effects from local biological interactions and domain noise:

This decomposition guides our Triple-Branch Disentanglement architecture (Figure 1), where each term corresponds to a distinct functional module:

• Intrinsic Potency (µ): The baseline effect of drug d (e.g., general toxicity) roughly invariant across cell lines. Modeled by finv(zinv).

• Explicit Specificity (τ): The variation attributable to specific interactions between drug structure and observed genomic features xbio c . This represents the mechanistic, generalizable signal. Modeled by fexp(zexp).
• Residual Bias (ϵ): The discrepancy arising from unobserved factors coupled with cell identity c (e.g., experimental noise, unmeasured omics). Modeled by fres(zres).

GBD-DRP Overview

Guided by Eq. (1), GBD-DRP implements a Triple-Branch Disentanglement architecture to model these components explicitly. As illustrated in Figure 1, the model maps the input space into three orthogonal latent representations:

where zinv, zexp, and zres encode the invariant drug features, mechanism-aware biological interactions, and residual cell-identity biases, respectively. During the training phase, we employ a Conditional Orthogonality Constraint to force these representations to be statistically independent. Crucially, during the inference phase, the spurious branch fres(zres) is discarded. This ensures that the final prediction relies solely on the invariant potency and explicit biological interactions, thereby achieving robust generalization.

3.2 Channel I: Invariant Drug Encoder (Einv)

To model the intrinsic potency µ(d) defined in Sec. 3.1, the Invariant Channel extracts structural features that remain consistent across different cellular environments. This contextindependent baseline is captured via a dual-pathway architecture.

Graph-level Encoding. We utilize an EdgeGAT encoder to process the molecular graph G = (V, E). To capture nuanced chemical motifs, we explicitly integrate bond features ej,i (e.g., bond type, conjugation) into the attention mechanism. The atom-level update for node i at layer ℓ is formulated as:

where α(ℓ) ij is the attention coefficient computed over concatenated node-edge features, and σ is an activation function. These local features are then aggregated via a self-attention pooling layer to yield a permutation-invariant graph embedding gemb.

Global Descriptor. In parallel, 2048-bit Morgan fingerprints are projected through an MLP to generate a descriptor fpemb, capturing long-range functional groups. The final invariant representation zinv is obtained by:

By anchoring zinv to the drug's mean response across cell line (see Sec. 3.6), we ensure it captures the general pharmacological baseline rather than cell-specific variations.

3.3 Channel II: Cell Line-Aware Explicit Encoder (Eexp)

To model the mechanistic interaction τ (xd, xbio c ), the Explicit Channel aligns drug structural motifs with cell-line-specific genomic features. Unlike the global bias isolated in Channel III, this channel captures generalizable patterns.

Cross-Modality Interaction. We utilize a multi-head cross-attention mechanism to identify dependencies between drug atoms and multi-omics features. For a specific omics modality k ∈ {base, opt1, . . . , optK}, let V be the drug atom embeddings and xk be the genomic feature vector. The interaction is captured as:

where the queries are derived from chemical structures and keys/values from biological features. This identifies specific genomic sites sensitive to particular functional groups.

Modality-Specific Encoding. To facilitate the reconstruction requirement (Lrecon in Sec. 3.6) and the counterfactual consistency, we generate a latent code for each observed modality:

The total explicit contribution is then the masked summation: fexp(zexp) = ∑k mk · f (k) exp(z(k) exp). By maintaining modalityspecific latent paths, the model ensures that the biological semantics of each omics layer are preserved and can be reconstructed, addressing the ambiguity regarding feature fusion.

3.4 Channel III: Residual Bias Encoder (Eres)

The final component of our decomposition addresses the unobserved cell-specific factors ϵ(xd, c), such as experimental batch effects or unmeasured omics. To prevent the model from "hallucinating" these non-generalizable variations into the explicit branch, Channel III acts as a residual buffer. It isolates domain-specific noise during training to shield the mechanistic integrity of the zexp representations.

Drug-Conditioned Bias Encoding. We model the residual bias as a joint interaction between drug properties and cell identity. Let ec ∈ Rdemb be a learnable embedding representing the unique signature of cell line c.

Modality-Aware Conditioning. To address the ambiguity of mask incorporation, we explicitly condition the residual encoder on the modality availability mask m ∈ {0, 1}K . This allows the branch to dynamically adjust its bias estimation based on which omics data are missing. The residual latent code is synthesized as:

where ∥ denotes concatenation. By incorporating m, zres can learn to compensate for missing biological signals (detailed in Sec. 3.5).

To ensure zres captures identity-specific noise rather than generalizable signals, we regularize it with an identity classifier Cres (Sec. 3.6). Crucially, this entire channel is deactivated during inference. This force-filters all residual domain biases, ensuring that predictions for unseen cell lines are driven exclusively by invariant and explicit biological mechanisms.

3.5 Missing-Aware Learning via Counterfactual Compensation

To handle incomplete multi-omics data without relying on noise-prone imputation, we propose a Counterfactual Compensation mechanism. We reframe modality absence as a dynamic signal migration: when a modality is missing, its biological effect is not "lost" but instead migrates from explicit interaction features (zexp) to implicit residual representations (zres).

Dynamic Routing Mechanism

We define cellular features as a set of K + 1 modalities, where k = 0 is the stable base modality (e.g., GEP) and k ∈ {1, . . . ,K} are optional modalities (e.g., CNV, Mutation). Modality availability is controlled by a binary mask m ∈ {0, 1}K . The drug response prediction is dynamically routed:

By conditioning fres on m (via concatenation as described in Sec. 3.4), the model adaptively adjusts the residual contribution. This architecture ensures that when mk = 0, the system does not zero out the underlying biological contribution but instead shifts its estimation to the residual latent space.

Counterfactual Consistency Learning

To formally train zres to recover missing signals, we introduce the Counterfactual Consistency loss (Lcons). For a sample where modality k is observed (mk = 1), we artificially simulate a missing state by applying a mask m\k.

We define the target bias contribution Tk under the fullinformation state as the sum of the implicit residual and the explicit interaction of modality k:

We then enforce a consistency constraint using Mean Squared Error:

where sg(·) denotes the stop-gradient operator. This treats Tk as a stable reference, compelling the residual branch to "learn to compensate" for the missing modality by aligning its partial-information output with the full-information target.

Through this mechanism, knowledge is transferred from the explicit mechanistic branch (fopt) to the implicit residual branch (fres) under simulated missing contexts. Consequently, during inference on incomplete data, GBD-DRP can recover missing biological signals through the residual latent space while shielding the explicit branch from zero-filling artifacts.

3.6 Training Objective and Statistical Constraints

To ensure that the decomposed latent representations (zinv, zexp, zres) capture distinct functional aspects of the drug response, we formulate a joint optimization objective. This composite loss integrates the predictive task with statistical constraints designed to enforce representational fidelity and geometric disentanglement.

1. Primary Prediction (Ltask) & Mean Anchoring (Lmean): The total prediction follows our tri-channel summation: ˆ y = finv(zinv) + fexp(zexp) + fres(zres). The task loss is Ltask = 1N ∑(yi - ˆ yi)2 . To anchor zinv as the context-independent potency µ(d), we introduce Lmean = 1N ∑(finv(z(i) inv ) - ¯ yd(i))2 , where ¯ yd(i) is the drug-specific average response in the training set. This prevents zinv from absorbing cell-specific variations.
2. Semantic Reconstruction (Lrecon): To prevent zexp from losing biological meaning during fusion, we require it to reconstruct the original omics features. For each modality k, we employ a decoder Dec(k) (a 2-layer symmetric MLP):

Beyond semantic grounding, this constraint is critical for preventing posterior collapse. Specifically, the conditional orthogonality constraint (Lorth) may inadvertently drive zexp to become uninformative to minimize overlap with the residual space. Lrecon acts as an information-preservation bottleneck, forcing the mechanistic branch to retain sufficient biological specificity even under strict disentanglement pressure.

1. Residual Identity Prior (Lclass): To force zres to isolate domain-specific noise (e.g., identity-linked batch effects), we employ an auxiliary classifier Cres (a single linear layer with Softmax) to predict the cell identity c from the residual code:

1. Conditional Orthogonality (Lorth): To eliminate redundant information, we minimize the squared cosine similarity between branches. Crucially, the orthogonality between zexp and zres is conditional (gated by mk):

This gating allows zres to capture information from missing modalities (mk = 0) while prohibiting information leakage when full data is observed.

1. Joint Optimization: The final objective is an end-toend weighted sum:

where Lcons is the counterfactual consistency loss from Sec. 3.5. The penalty term λn∥fres(zres)∥2 forcing the model to prioritize explicit channels. Hyperparameters λ were determined via grid search; a detailed sensitivity analysis is provided in Section 4.5.

4 Experiment

4.1 Datasets and Experimental Setup

Following the experimental framework of PASO [Wu et al., 2025], we utilized a large-scale pharmacogenomic dataset consolidated from the GDSC and CCLE databases. The final curated dataset comprises 141,222 drug-response pairs, covering 233 unique drugs and 688 cancer cell lines.

Drug Representations

To capture both global and local molecular semantics, we extracted two complementary features for each compound: (1) Molecular Graphs for structural message passing, and (2) Morgan Fingerprints for fixed-length substructure encoding.

Cell Line Features

We integrated three distinct omics modalities to characterize cellular heterogeneity: Gene Expression Profiles (GEP), Copy Number Variations (CNV), and Somatic Mutations (MUT).

Evaluation Protocol

We adopted a 10-fold cross-validation strategy to benchmark performance. To comprehensively assess model robustness, we designed three distinct experimental scenarios:

• Mixed Setting: Drug-cell pairs are randomly split into training and testing sets.
• Cell-Blind Setting: The dataset is split based on cell line identities, ensuring no cell line in the test set appears during training. This simulates a "precision medicine" scenario for unseen patients.
• Missing Modality Setting: We randomly masked a certain ratio of the optional modalities (CNV and MUT), testing the model's ability to maintain performance relying primarily on the stable modality (GEP) and utilize unstable modalities.

4.2 Baseline Methods

To evaluate the performance of our proposed framework, we compare it against a comprehensive set of baselines ranging from classical machine learning algorithms to state-of-the-art deep learning models:

Random Forest (RF): Classical machine learning baselines that utilize concatenated drug features (e.g., fingerprints or encodings) and cell line genomic profiles to model non-linear response patterns via decision tree ensembles and kernel-based regression, respectively.

XGBoost [Chen, 2016]: An optimized gradient boosting decision tree algorithm that iteratively trains weak learners to capture complex, non-linear interactions between highdimensional omics features and drug descriptors.

PathDSP [Tang and Gottlieb, 2021]: A deep neural network approach that integrates drug molecular fingerprints with pathway-level features derived from gene expression enrichment analysis to predict drug sensitivity.

• FA2GL [Xu et al., 2025]: A feature aggregation and graph learning network that captures fine-grained drug-cell interactions via cross-attention and models network-level associations using a Graph Neural Network (GNN) with auxiliary pipe nodes.

PASO [Wu et al., 2025]: A hybrid deep learning model that combines pathway-based multi-omics difference features with drug SMILES sequences, utilizing multi-scale CNNs, Transformers, and attention mechanisms to capture molecular interactions.

4.3 Results

Mixed Cross Validation Results

Table 1 presents the quantitative comparison of GBD-DRP against baseline methods. GBD-DRP consistently outperforms all competing methods across all three evaluation metrics. Specifically, our model achieves the lowest RMSE of 0.9288 and the highest R2 of 88.99%, surpassing the second-best method, PASO (RMSE: 0.9400, R2 : 88.38%). While PASO already establishes a strong baseline by utilizing pathway-based difference features and transformer encoders, GBD-DRP further reduces the prediction error by 1.2% in terms of RMSE. This performance gain is attributed to our explicit bias disentanglement strategy, which effectively isolates spurious domain noise from true biological signals.

Cell Blind Cross Validation Results

As summarized in Table 2, although the absolute improvement in the Cell-Blind setting is more modest than in the Mixed setting, GBD-DRP's advantages lie in its predictive stability and noise resilience.

Compared to PASO, GBD-DRP reduces the RMSE from 1.3166 to 1.2872 while significantly lowering the prediction variance (±0.0631 vs. ±0.1778). Since PASO lacks an explicit mechanism to isolate domain-specific biases, its attention mechanism inadvertently captures non-generalizable experimental noise (e.g., batch effects or lab-specific artifacts) and treats them as predictive biological signals.

In contrast, GBD-DRP addresses this through its Residual Disentanglement strategy. By delegating cell-specific residual variations to the dedicated buffer (zres), our model "filters" the information entering the explicit branch (zexp). This ensures that the learned "lock-and-key" interactions are grounded in invariant biological mechanisms rather than spurious correlations. Consequently, when transitioning to unseen cell lines, GBD-DRP maintains a more robust performance by relying on these filtered, generalizable patterns, whereas PASO's performance fluctuates due to its entanglement with the unobserved biases of the training domains.

Missing Modality Cross Validation Results

In clinical practice, multi-omics profiles are frequently incomplete due to cost or sample degradation. To evaluate GBD-DRP's robustness under such data-constrained scenarios, we simulate a challenging environment where 50% of the optional modalities (CNV and MUT) are randomly masked during inference.

To ensure a fair and rigorous comparison, we implemented three representative strategies for the baseline models that lack inherent missing-aware mechanisms: (1) Mean Imputation: Missing features are replaced by the global mean of the respective modality in the training set. (2) Zero Filling: Missing entries are set to zero, a common practice in deep learning to neutralize missing inputs. (3) Feature Exclusion: The models are restricted to using only the stable modality (GEP), effectively bypassing the unobserved features.

As shown in Table 5, GBD-DRP demonstrates a decisive advantage, achieving an RMSE of 0.9322 and a Pearson correlation of 94.20%, significantly outperforming all baseline methods. Notably, the performance of GBD-DRP under 50% modality missingness remains remarkably close to its full-modality performance, whereas baselines experience a sharper decline.

Intriguingly, baselines perform better with Feature Exclusion than naive Mean/Zero Imputation, suggesting that inputlevel infilling introduces detrimental "phantom signals" and overfitting to imputation artifacts. In contrast, GBD-DRP leverages Counterfactual Consistency to adaptively compensate for missing explicit signals in the latent space. By dynamically reallocating unobserved variance to the residual channel (zres), GBD-DRP "recovers" the biological essence of missing modalities without noise injection, ensuring superior resilience in incomplete data environments.

4.4 Ablation Study and Sensitivity Analysis

To rigorously justify the necessity of GBD-DRP's components, we conduct a fine-grained ablation study across both Mixed and Cell-Blind settings. We evaluate four variants: (1) w.o. Attn: Replacing cross-attention in Channel II with concatenation; (2) w.o. Lorth: Removing orthogonality constraints; (3) w.o. Lcons: Disabling the counterfactual consistency mechanism; (4) w.o. Res: Removing the entire Residual Channel and its associated losses (Lclass, Lorth).

Mechanistic Interaction via Attention: The w.o. Attn variant yields inferior results, particularly in R2 . This suggests that simple concatenation fails to capture the "lock-andkey" dependencies between chemical motifs and genomic sites, whereas attention provides the necessary granularity for modeling drug sensitivity.

Generalization to Unseen Domains (Cell-Blind): As shown in Table 4, the performance of the w.o. Res variant collapses significantly in the Cell-Blind setting, with RMSE increasing from. This confirms that without a dedicated residual channel to sequester domain-specific noise, non-generalizable biases leak into the predictive branches, severely hindering zero-shot transfer.

Fine-grained Analysis of the Residual Branch: The degradation in the w.o. Lorth variant indicates that representational disentanglement is not achieved by structural separation alone; explicit geometric constraints are required to prevent information redundancy.

Counterfactual Consistency (Lcons): Removing Lcons under 50% missingness increases RMSE from 0.9337 to 0.9475 (Table 5). This confirms Lcons is vital for signal compensation: it compels the residual branch (zres) to actively recover functional information lost from omitted modalities.

These results validate that our mechanism successfully reallocates predictive logic from the explicit branch (zexp) to the residual buffer during data depletion, ensuring a robust "information safety net."

4.5 Implementation Details

Our model is implemented using PyTorch and Deep Graph Library (DGL). The choice of hyper-parameter are based on grid search and shown in Table 6

5 Conclusion

In this study, we presented GBD-DRP, a novel tri-channel disentanglement framework designed to enhance the generalizability and robustness of drug response prediction. By explicitly decomposing the latent space into invariant pharmacological effects, explicit biological interactions, and residual domain-specific biases, our model effectively mitigates the "shortcut learning" common in conventional DRP methods. Through the integration of a counterfactual consistency mechanism, GBD-DRP not only isolates spurious experimental noise but also provides a principled solution for the missing modality problem, enabling robust "signal compensation" without relying on data imputation. Extensive experiments demonstrate that GBD-DRP consistently achieves state-of-the-art performance in both Cell-Blind generalization and missing modality data scenarios. Future research can focus on extending the current domain generalization (DG) paradigm. While this work primarily addresses cellular-level biases, we aim to investigate a composite DG framework that treats both cell lines and drugs as distinct yet interacting domains. By modeling "Drug-Domain" shifts-such as those arising from chemical scaffold variations or novel drug classes-we anticipate developing a more holistic system capable of achieving simultaneous zero-shot generalization for both unseen patients and unseen therapeutic compounds. Additionally, integrating single-cell multi-omics data into this disentangled architecture could further refine the precision of mechanistic explanations at the sub-clonal level.

Ethical Statement

There are no ethical issues.

Acknowledgements

This work was supported by the Science and Technology Commission of Shanghai Municipality (grants No. 24510714300), and the Shanghai Municipal Science and Technology Major Project, China (Grant No. 2021SHZDZX0102).