MoleculeNet Classification Data : Modularization, dynamic splits, Clean Code

[aditya0by0]

Changes related to Molecule Net classification introduced in #130

• Remove Redundant Code logic for molecule net classification
• Modularize the molecule net classification code
• Enable dynamic splitting functionality
• Use deepchem to access data (https://deepchem.readthedocs.io/en/latest/api_reference/moleculenet.html)
• Use the recommended split type ("scaffold" or "random") for underlying data specified by https://arxiv.org/abs/1703.00564
• Disable split data size specification (we shouldn't use custom split size for these data rather we should follow the standards used in community which is 80/10/10, for fair comparison)
• Add Two more missing data classes PCBA and ToxCast for classification
• Delete Tox21 data class and test, instead new modularized class is created

General Changes

• Check before dataloader to assert if the data is empty (as pytorch or lightning doesn't check this)

• Generalize _DynamicDataset class to all datasets by replacing _graph_to_raw_dataset to _preprocess_data_into_dataframe

• Modularize splitting logic into separate classes (so we dont need to specifiy the split logic for every new data class, rather just inherent the required splitting logic class )

Suggested Metric as per data type by https://arxiv.org/abs/1703.00564 for Molecule Net dataset

Dataset	Task Type	Classes	Recommended Metric	Why
BACE	Binary classification	2	ROC-AUC	Standard benchmark metric; moderate class imbalance
BBBP	Binary classification	2	ROC-AUC	MoleculeNet standard; threshold-independent
HIV	Binary classification	2	ROC-AUC	Highly imbalanced dataset
MUV	Multi-task binary (17 tasks)	2	PR-AUC (Average Precision)	Extremely imbalanced; ROC-AUC can be misleading
PCBA	Multi-task binary (128 tasks)	2	PR-AUC (Average Precision)	Very sparse positives; PR-AUC is preferred
SIDER	Multi-task binary (27 tasks)	2	ROC-AUC	Standard MoleculeNet evaluation
ClinTox	Multi-task binary (2 tasks)	2	ROC-AUC	Standard benchmark metric
Tox21	Multi-task binary (12 tasks)	2	ROC-AUC	Widely used benchmark metric
ToxCast	Multi-task binary (hundreds of tasks)	2	ROC-AUC (sometimes PR-AUC as secondary)	Standard MoleculeNet metric

[aditya0by0]

@sfluegel05, once I complete the gnn experiments with molecule net dataset, then we could merge the PR. Until then we could keep this PR in the draft state.

[aditya0by0]

@sfluegel05, as I want to train the augmented GNN on the MolNet dataset, I looked into reusing the missing-label handling implemented in #130.

Comparing it with HiMol's implementation, I noticed two critical discrepancies in AUROC computation and Loss normalization that may cause behavioral divergence. Please let me know if I am wrong about them or if I am missing something.

1. AUROC Computation

• HiMol's approach: For each task column, it builds a boolean mask (is_valid) and removes missing entries entirely before computing the metric via roc_auc_score(y_true[is_valid, i], y_scores[is_valid, i]). The sample size $N$ dynamically shrinks to include only annotated compounds for that task.
  • HiMol reference: optimization.py#L95-L104
• PR New regression and classification datasets for ontology pre-training #130's approach: Keeps all samples, but forces missing entries to prediction = 0 and label = 0 in _get_prediction_and_labels (chebai/models/electra.py).

Why this is problematic: (Detailed by Claude)
Zeroing out missing values does not exclude them; instead, it injects them as artificial, trivially correct negative samples ranked at the bottom of the prediction distribution. This introduces two issues:

1. Inflated Denominator: It artificially alters the negative sample count in the ROC denominator ($\text{positives} \times \text{negatives}$) rather than evaluating the true sample population.
2. Artificial Score Inflation: These synthetic $(0, 0)$ pairs act as "free wins" in ranking, systematically pushing AUROC scores higher whenever real negatives have prediction scores $> 0$ (which is practically always true for sigmoid outputs). The inflation gap scales directly with label sparsity.

2. Loss Computation & Denominator Normalization

• HiMol's approach: Computes elementwise loss, zeros out invalid positions with torch.where(is_valid, loss_mat, zeros), and normalizes strictly by the valid count:

  • HiMol reference: optimization.py#L32-L51

• PR New regression and classification datasets for ontology pre-training #130's approach: Masks before the loss computation in _process_for_loss by corrupting inputs (setting logits to -10000 and labels to 0), then passing them to BCEWithLogitsLoss.

Why this is problematic: (Detailed by Claude)
Unless normalized downstream by the valid entry count (which is not) averages across all elements . This leads to:

1. Denominator Dilution: Highly sparse datasets (e.g., MUV 84% missing, Tox21 15% missing) have their loss and gradients diluted proportionally to the number of missing labels in each batch.

2. Inconsistent Effective Learning Rates: Batches or tasks with high missingness produce artificially diminished gradients compared to dense batches, distorting training dynamics based purely on sparsity patterns rather than model performance.

[aditya0by0]

@sfluegel05, as I want to train the augmented GNN on the MolNet dataset, I looked into reusing the missing-label handling implemented in #130.

Comparing it with HiMol's implementation, I noticed two critical discrepancies in AUROC computation and Loss normalization that may cause behavioral divergence. Please let me know if I am wrong about them or if I am missing something.

1. AUROC Computation

• HiMol's approach: For each task column, it builds a boolean mask (is_valid) and removes missing entries entirely before computing the metric via roc_auc_score(y_true[is_valid, i], y_scores[is_valid, i]). The sample size
  N
  dynamically shrinks to include only annotated compounds for that task.

  • HiMol reference: optimization.py#L95-L104

• PR New regression and classification datasets for ontology pre-training #130's approach: Keeps all samples, but forces missing entries to prediction = 0 and label = 0 in _get_prediction_and_labels (chebai/models/electra.py).

Why this is problematic: (Detailed by Claude) Zeroing out missing values does not exclude them; instead, it injects them as artificial, trivially correct negative samples ranked at the bottom of the prediction distribution. This introduces two issues:

1. Inflated Denominator: It artificially alters the negative sample count in the ROC denominator (
   positives
   ×
   negatives
   ) rather than evaluating the true sample population.
2. Artificial Score Inflation: These synthetic
   (
   0
   ,
   0
   )
   pairs act as "free wins" in ranking, systematically pushing AUROC scores higher whenever real negatives have prediction scores
   0
   (which is practically always true for sigmoid outputs). The inflation gap scales directly with label sparsity.

2. Loss Computation & Denominator Normalization

• HiMol's approach: Computes elementwise loss, zeros out invalid positions with torch.where(is_valid, loss_mat, zeros), and normalizes strictly by the valid count:

  • HiMol reference: optimization.py#L32-L51

• PR New regression and classification datasets for ontology pre-training #130's approach: Masks before the loss computation in _process_for_loss by corrupting inputs (setting logits to -10000 and labels to 0), then passing them to BCEWithLogitsLoss.

Why this is problematic: (Detailed by Claude) Unless normalized downstream by the valid entry count (which is not) averages across all elements . This leads to:

1. Denominator Dilution: Highly sparse datasets (e.g., MUV 84% missing, Tox21 15% missing) have their loss and gradients diluted proportionally to the number of missing labels in each batch.
2. Inconsistent Effective Learning Rates: Batches or tasks with high missingness produce artificially diminished gradients compared to dense batches, distorting training dynamics based purely on sparsity patterns rather than model performance.

I have implemented this change, to test that this change doesn't affect the models trained on datasets without any missing labels, I trained same graph based model before and after this change for BBBP and the results were identical down to the last decimal.

Before Change: https://wandb.ai/chebai/chebai/runs/zuxii7ad/overview?nw=nwuseraditya0by0
After Change: https://wandb.ai/chebai/chebai/runs/1doijnz3/overview?nw=nwuseraditya0by0

Right now I have implemented the change for graph models, hence this PR will break the missing labels implementation for ELECTRA. To fix it I will create a separate issue for it and it should be fixed in separate PR.