DEMEC: Drug Embedding & Multi-Effect Classification

By Diego Bustamante, Julian Allchin, and Georgios Mikos as part of the Stanford CS224W course project.

Assessing drug side effects represents a key challenge for pharmaceutical companies developing novel therapeutics because small changes in the chemical structure or the properties of a drug can result in drastically different side effects. Predicting drug side effects from the molecular structure would allow for the rapid screening of candidates to accelerate the drug development pipeline. Here, we explore graph neural networks (GNNs), given their inherent advantage with modeling molecular structures, to predict drug side effects and other molecular properties of varying difficulty.

All code developed can be found here: https://github.com/diegobus/DEMEC.

Motivation and Novelty

The drug-development pipeline narrows dramatically from discovery to approval, with unknown side effects accounting for many of the losses. Previous approaches to side effect prediction often include experimentally-determined molecular properties, such as binding affinities to known off-targets, plasma protein binding, and physicochemical descriptors. However, these properties can be expensive to determine, meaning that these models cannot aid in early-stage development when hundreds of thousands of compounds are considered in silico.

In this work, we explore whether these costly molecular properties can still guide representation learning without being provided as input features. Instead, we incorporate them as auxiliary supervision signals through a multitask learning framework [1]. A single GNN encoder is trained directly on molecular graphs and jointly supervised on three tasks: clinically reported side effects, Anatomical Therapeutic Chemical (ATC) classifications, and MACCS substructure fingerprints. At inference time, the model requires only the molecular structure, but its embedding has been shaped by richer biochemical signals encountered during training.

Our analysis includes extensive comparisons on model architecture, pooling mechanisms, and multitask configurations, allowing us to characterize how different sources of auxiliary supervision shape the learned molecular embeddings. This design aims not only to improve predictive accuracy, but also to illuminate when and why multitask learning is beneficial in molecular settings.

Data Description

We leveraged the SIDER 4.1 Side Effect Resource that provides information about 5868 side effects, 1430 drugs, and 139756 side effects — drug interactions [2]. SIDER drugs are identified with a Chemical ID (CID) and Anatomical Therapeutic Chemical Classification (ATC) code, which we leveraged alongside the tools RDKit and PubChemPy to generate SMILES strings, NetworkX molecular graph representations, and to scrape additional molecular properties such as molecular weight. Drugs with molecular weight over 2000 g/mol were excluded from the dataset as our focus was assessing drug effects for small-molecule therapeutic agents. This dataset and additional tools provide a comprehensive resource for our task.

Task

Given the molecular structure of a drug, our model predicts multiple molecular and clinical properties at the graph level. Concretely, we consider four tasks of increasing abstraction from structure:

1. Molecular weight (MW) — a simple regression baseline that is almost deterministic from the graph.
2. MACCS fingerprints — 166-bit molecular fingerprints that encode the presence of predefined chemical substructures. Predicting these forces the model to learn chemically interpretable structure [3].
3. ATC classification — multi-label classification using Level-3 ATC codes (167 classes). Unlike side effects, ATC codes are curated and much cleaner, making them a good test case for whether multitask learning can meaningfully shape the embedding. Examples of Level-3 ATC codes include biguanides, a class of molecules that reduce glucose output from the liver, and selective calcium channel blockers [4].
4. Side-effect prediction — multi-label prediction using SIDER (4251 classes). This is the noisiest and most challenging task because most side effects appear in only a handful of drugs, and each drug often has multiple side effects. Examples of side effects include heat stroke, scurvy, bladder infection, and gallbladder cancer.

Model Design

Our model consists of three components: (1) a molecular graph representation, (2) GNN encoder, and (3) multi-task prediction heads. Training runs can be configured with specific heads to use for training and inference.

Molecular Graph Representation:

To create a lossless representation of each molecule, each drug is represented as an undirected graph where atoms are nodes and bonds between atoms are edges. To create molecular features, we use one-hot vectors for each of the following chemical properties for each atom: element, total degree, formal charge, total valence, aromaticity, ring membership, hydrogens, hybridization, and chirality. These features are concatenated to create 154-dim node features. Edges are classified by their bond type, enabling the GNN to distinguish aromaticity, conjugation, and bond multiplicity.

Graph Neural Network Encoder:

The molecular graph is fed into a shared GNN backbone that learns a continuous molecular embedding via message passing. We evaluate two architectures: GCN, which emphasizes local aggregation [5], and GAT, which applies learned attention weights to focus on chemically relevant neighbors [6]. Both models use five layers with 128 hidden units and ReLU activations, followed by one of three graph pooling strategies — mean, MLP, or attention pooling — to produce a fixed-length embedding G_{\text{emb}. This embedding serves as the unified representation of the molecule used by all downstream tasks. Each architecture handles the heterogeneous graphs based on different edge types by aggregating messages per edge type. We selected this approach because it inherently groups together aromatic rings and other chemical structures that share electrons and thus should be considered distinctly.

Heterogeneous Graph Construction

# From src/demec/models/gnn_backbone.py
class GNNBackbone(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, dropout=0.2, 
                 conv_type='gcn', heads=1, pooling='mean'):
        super().__init__()
        self.pooling = pooling

        # Heterogeneous edge type handling
        self.edge_types = ['single', 'double', 'triple', 'aromatic']

        # Separate convolution for each bond type
        self.convs = nn.ModuleList()
        for _ in range(num_layers):
            layer_convs = nn.ModuleDict()
            for edge_type in self.edge_types:
                if conv_type == 'gat':
                    layer_convs[edge_type] = GATConv(
                        input_dim if len(self.convs) == 0 else hidden_dim,
                        hidden_dim // heads,
                        heads=heads,
                        dropout=dropout
                    )
                else:  # gcn
                    layer_convs[edge_type] = GCNConv(
                        input_dim if len(self.convs) == 0 else hidden_dim,
                        hidden_dim
                    )
            self.convs.append(layer_convs)

Message Passing with Edge Types

# From src/demec/models/gnn_backbone.py
def forward(self, data: HeteroData):
    x = data['atom'].x
    batch = data['atom'].batch

    # Message passing over heterogeneous edges
    for layer_idx, layer_convs in enumerate(self.convs):
        x_new = []

        # Aggregate messages from each edge type
        for edge_type in self.edge_types:
            edge_attr = f'atom__to__atom_{edge_type}'
            if edge_attr in data.edge_index_dict:
                edge_index = data[edge_attr].edge_index
                x_conv = layer_convs[edge_type](x, edge_index)
                x_new.append(x_conv)

        # Sum contributions from all edge types
        if x_new:
            x = torch.stack(x_new, dim=0).sum(dim=0)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)

    return x

Attention-Based Graph Pooling

# From src/demec/models/gnn_backbone.py
if self.pooling == 'attention':
    # Learnable attention pooling
    attn_weights = self.attn_pool(x)  # [num_nodes, 1]
    attn_weights = scatter_softmax(attn_weights, batch, dim=0)
    graph_emb = scatter_sum(x * attn_weights, batch, dim=0)

Multi-Task Prediction Heads

The shared 128-dim embedding G_{\text{emb}} is passed into lightweight MLP classifiers or regressors, each responsible for a different prediction task:

• Molecular Weight: single-output regression.
• MACCS Fingerprint: multi-label prediction over 166 MACCS substructure keys.
• ATC: multi-label prediction over 167 Level-3 ATC classes.
• Side Effect: multi-label classification over all 4,251 SIDER side effects or a filtered Top-N subset.

During multitask training, all heads operate simultaneously, and their losses (described in Table 1) are combined into a single multitask objective. Molecular weight regression uses MSE loss and the multi-label classification tasks use BCEWithLogitsLoss, except for side-effect prediction, which requires Focal Loss due to extreme class imbalance [7]. We experiment with both equal weighting and manually tuned task weights.

# From src/demec/models/gnn_backbone.py
class MultiTaskGNN(nn.Module):
    def __init__(self, backbone, heads_dict):
        super().__init__()
        self.backbone = backbone
        self.heads = nn.ModuleDict(heads_dict)

    def forward(self, data: HeteroData):
        # Single forward pass through backbone
        graph_emb = self.backbone(data)

        # Separate prediction heads
        results = {}
        for task_name, head in self.heads.items():
            results[task_name] = head(graph_emb)

        return results

Training and Evaluation

All models were trained under a standardized protocol to ensure comparability across architectures and task configurations. The dataset was partitioned into training, validation, and test sets using a fixed 80/10/10 split over unique drug CIDs, with a deterministic random seed (42) and no stratification. Model parameters were optimized using the Adam optimizer with a learning rate of 1x10^-3 and a batch size of 32. Training proceeded for 300–500 epochs based on validation performance.

To evaluate model performance, individual metrics appropriate to their label structure were used (Table. 1)

These task–metric pairings reflect the statistical nature of each prediction objective: R² is standard for continuous regression, Tanimoto similarity is the canonical metric for evaluating binary chemical fingerprints [3], while mAP and AUROC appropriately measure performance for multi-label pharmacological classifications with highly skewed label distributions.

Experimental Design

Step 1 — Pooling and Architecture Design (18 experiments). We first assessed how graph-level aggregation interacts with backbone choice and task type. For molecular weight, MACCS fingerprints, ATC, and side effects, we trained single-task models using GCN or GAT with mean, MLP, or attention pooling, yielding 18 backbone–pooling–task combinations. This phase established strong single-task baselines and identified preferred pooling strategies for each prediction task.

Step 2 — Multitask Learning (4 experiments). Next, we investigated whether auxiliary tasks improve performance on clinically relevant targets. Using the best-performing backbone/pooling settings from Phase 1, we compared single-task models against multitask variants that jointly predict side effects, ATC, and MACCS under different loss weightings and task groupings (e.g., SE+ATC+MACCS with weighted vs. equal losses, ATC as the primary task, and SE+ATC without MACCS).

Step 3 — Label Quality and Side-Effect Subsets (5 experiments). Finally, we examined the impact of side-effect label quality by varying the number of predicted side effects. Starting from a baseline model trained on all 4,251 SIDER side effects, we trained additional single-task and multitask models on prevalence-filtered subsets containing the Top-50, Top-100, or Top-200 most common side effects. This phase isolates how much of the difficulty in side-effect prediction arises from noisy, sparse labels versus limitations of the model architecture.

Results

Step 1 — Pooling and Architecture Design

In Step 1, we examined how backbone architecture (GCN vs. GAT) and pooling strategy (mean, MLP, or attention) influence performance across the four prediction tasks. The pooling comparison (Figure 1) shows that no single pooling method dominates; instead, each task favors a different form of graph-level aggregation. MLP pooling performs best for molecular weight, which is almost entirely determined by global structural properties. Attention pooling excels on MACCS fingerprints and side-effect prediction, both of which depend on identifying local functional groups. Mean pooling performs best for ATC classification, suggesting that therapeutic class signals are distributed across the molecular scaffold rather than localized to specific substructures.

The architecture comparison in Figure 2 reinforces the task-dependent nature of model performance. GAT outperforms GCN on ATC and side-effect prediction — tasks requiring fine-grained sensitivity to functional motifs — while GCN performs best on molecular weight, where global averaging is beneficial. MACCS fingerprints show near parity between architectures. These findings establish that architectural choices meaningfully affect representation learning, but their impact depends on the nature of the task.

Step 2 — Multitask Learning

Step 2 evaluates whether multitask learning improves predictive performance by enriching the supervision available to the shared encoder. For side-effect prediction, multitask learning provides only marginal benefit. As shown in Figure 3 (left), adding ATC and MACCS supervision improves mAP by just +1.7% (0.420 → 0.427), and removing MACCS supervision actually decreases performance. Even when restricting the prediction space to the 100 most frequent side effects (Figure 3, right), the multitask gain remains small (0.669 → 0.679 mAP). These results suggest that the extreme label sparsity and noise in SIDER fundamentally limit the usefulness of additional supervision for this task.

In sharp contrast, ATC classification benefits substantially from multitask learning. Figure 4 shows that performance increases from 0.176 mAP (GAT single-task) and 0.252 mAP (GCN single-task) to 0.300 mAP under multitask training — a 70% improvement relative to the weaker baseline. Because ATC labels are clean, well-curated, and densely populated, they provide a reliable supervisory signal that synergizes with the MACCS and SE tasks. Together, these results demonstrate that multitask learning is effective for molecular property prediction, provided that the task leverages well-curated labels.

Step 3 — Label Quality and Side-Effect Subsets

Step 3 investigates the extent to which side-effect prediction is limited by label quality rather than architectural or multitask factors. Figure 5 shows that filtering side-effect labels by prevalence dramatically improves performance: mAP increases from 0.427 on all 4,251 labels to 0.601 on the Top-200 labels, 0.669 on the Top-100, and 0.757 on the Top-50. This nearly 80% improvement reveals that the greatest barrier to accurate side-effect prediction is not model capacity but the sparsity, heterogeneity, and incomplete nature of clinical reporting. Multitask learning on these filtered subsets provides small additional gains, but the magnitude of improvement is modest compared to the gains achieved simply by removing rare, noisy labels.

To contextualize these results, we compared the improvements observed here with our previous approach of multi-task learning. For side effect prediction, multitask learning improved the performance by 1.7% only, whereas removing rare and noisy side effects improved the performance by 77%, reflecting the reality that rare clinical outcomes are far removed for the specific molecular structure of a drug and possibly tied to environmental factors or individual difference that cannot be captured without additional biological or patient-specific information. For a task that is more closely tied to the chemical structure, namely the prediction of ATC codes, we observed that multitask learning provides a 70% improvement relative to the weaker baseline, demonstrating that our models can learn meaningful representations from the molecular graphs.

Conclusion

Here, we leveraged the strength of GNNs to model molecular graphs in order to predict drug side effects. Considering multiple different GNN architectures and pooling strategies, we demonstrate that the best architecture is task-dependent: predicting ATC codes, which requires a holistic molecular picture, benefits from mean pooling, whereas MACCS predictions, which requires a consideration of distinct functional groups, benefits from attention pooling and a GCN. Considering the effect of multitask learning, we illustrate minor improvements to side effect prediction but more pronounced improvements for ATC predictions, owing to their well-curated nature. Finally, we demonstrate that removing rare side effects drastically improves the performance, indicating that label sparsity and a relative lack of data (only 1430 drugs for 4251 side effects) limited the effectiveness of our approach.

Considering future directions, beyond attempting to increase the dataset size, we envision two potential approaches. First, graph-representation of our drugs could be included in a larger, hierarchical, heterogeneous graph that explicitly models drug-side effect relationships as edges. Additionally, this graph could be augmented with additional biological data, such as drug-protein interactions, protein-protein interactions, and protein-side effect association, thus creating a rich knowledge graph that allows for information to be passed across different scales and provides more biological context. Second, future work could integrate molecular dynamics simulations to develop richer node and edge embeddings. By encoding conformational flexibility and interactions with common binding pockets, the model could better capture functional drug behavior, as opposed to static chemical representation. These approaches both aim to exploit additional biological and physical properties of drugs to aid in the overarching goal of predicting drug side effects to enable rapid drug development.

References

[1] Dey, V. & Ning, X. Enhancing molecular property prediction with auxiliary learning and task-specific adaptation. J Cheminform 16, 85 (2024).

[2] Kuhn, M., Letunic, I., Jensen, L. J. & Bork, P. The SIDER database of drugs and side effects. Nucleic Acids Res 44, D1075–D1079 (2016).

[3] Cereto-Massagué, A. et al. Molecular fingerprint similarity search in virtual screening. Methods 71, 58–63 (2015).

[4] Anatomical Therapeutic Chemical (ATC) Classification. World Health Organization https://www.who.int/tools/atc-ddd-toolkit/atc-classification.

[5] Kipf, T. N. & Welling, M. Semi-Supervised Classification with Graph Convolutional Networks. Preprint at https://doi.org/10.48550/ARXIV.1609.02907 (2016).

[6] Veličković, P. et al. Graph Attention Networks. Preprint at https://doi.org/10.48550/ARXIV.1710.10903 (2017).

[7] Lin, T.-Y., Goyal, P., Girshick, R., He, K. & Dollár, P. Focal Loss for Dense Object Detection. Preprint at https://doi.org/10.48550/ARXIV.1708.02002 (2017).