Data leakage between train and validation sets

[goepp]

Report

Hi,
I noticed that the train and validation split in MILClassifier allows for instances from one sample to be in both train and validation sets.

How to reproduce

For instance, we can add the following code below mil.train(lr=1e-3) in the Classification with multiple instance learning (MIL) tutorial:

train_patients = adata.obs.iloc[mil.train_indices_]['sample'].value_counts().index
val_patients = adata.obs.iloc[mil.validation_indices_]['sample'].value_counts().index
train_val_overlap = set(train_patients).intersection(val_patients)
print(len(train_val_overlap))

The output is 108, the number of samples in both train and validation split

Is this a bug?

I believe this creates data leakage, as on my data (with MILClassifier running a regression task) the validation error is consistently below the training error. This is also visible on the tutorial, altough less obvious because it's running a classification task.

Version information

---

anndata 0.10.8
gdown 5.2.0
multimil 0.2.0
numpy 1.26.4
pandas 2.2.3
scanpy 1.11.0
scvi 0.20.0
session_info 1.0.0

PIL 11.1.0
absl NA
asttokens NA
attr 25.1.0
bs4 4.13.3
certifi 2025.01.31
charset_normalizer 3.4.1
chex 0.1.88
comm 0.2.2
cycler 0.12.1
cython_runtime NA
dateutil 2.9.0.post0
debugpy 1.8.12
decorator 5.1.1
docrep 0.3.2
etils 1.12.0
exceptiongroup 1.2.2
executing 2.1.0
filelock 3.17.0
flax 0.10.3
fsspec 2025.2.0
h5py 3.12.1
idna 3.10
ipykernel 6.29.5
ipywidgets 8.1.5
jax 0.5.0
jaxlib 0.5.0
jedi 0.19.2
joblib 1.4.2
kiwisolver 1.4.8
legacy_api_wrap NA
lightning_fabric 1.9.5
lightning_utilities 0.12.0
llvmlite 0.44.0
matplotlib 3.10.0
matplotlib_inline 0.1.7
ml_collections 1.0.0
ml_dtypes 0.5.1
mpl_toolkits NA
mpmath 1.3.0
msgpack 1.1.0
mudata 0.3.1
multipledispatch 0.6.0
natsort 8.4.0
numba 0.61.0
numpyro 0.17.0
nvidia NA
opt_einsum 3.4.0
optax 0.2.4
packaging 24.2
parso 0.8.4
pickleshare 0.7.5
pkg_resources NA
platformdirs 4.3.6
prompt_toolkit 3.0.50
psutil 6.1.1
pure_eval 0.2.3
pydev_ipython NA
pydevconsole NA
pydevd 3.2.3
pydevd_file_utils NA
pydevd_plugins NA
pydevd_tracing NA
pygments 2.19.1
pyparsing 3.2.1
pyro 1.9.1
pytorch_lightning 1.9.5
pytz 2025.1
requests 2.32.3
rich NA
scipy 1.15.2
setuptools 59.5.0
simplejson 3.20.1
six 1.17.0
sklearn 1.5.2
socks 1.7.1
soupsieve 2.6
stack_data 0.6.3
sympy 1.13.1
threadpoolctl 3.5.0
toolz 1.0.0
torch 2.6.0+cu124
torchgen NA
torchmetrics 1.6.1
tornado 6.4.2
tqdm 4.67.1
traitlets 5.14.3
triton 3.2.0
typing_extensions NA
urllib3 2.3.0
vscode NA
wcwidth 0.2.13
yaml 6.0.2
zmq 26.2.1
zoneinfo NA

IPython 8.32.0
jupyter_client 8.6.3
jupyter_core 5.7.2

Python 3.10.16 | packaged by conda-forge | (main, Dec 5 2024, 14:16:10) [GCC 13.3.0]
Linux-4.18.0-553.36.1.el8_10.x86_64-x86_64-with-glibc2.28

Session information updated at 2025-02-21 19:01

[goepp]

To be clearer: here's an example of train and validation loss on my data:

[goepp]

Here's a proposal for modifications to have sample-level split between train and validation

Proposed solution

Quick-fix code

import logging
import warnings

import anndata as ad
import torch
from anndata import AnnData
from pytorch_lightning.callbacks import ModelCheckpoint
from scvi import REGISTRY_KEYS
from scvi.data import AnnDataManager, fields
from scvi.data._constants import _MODEL_NAME_KEY, _SETUP_ARGS_KEY
from scvi.model._utils import parse_use_gpu_arg
from scvi.model.base import ArchesMixin, BaseModelClass
from scvi.model.base._archesmixin import _get_loaded_data
from scvi.model.base._utils import _initialize_model
from scvi.train import AdversarialTrainingPlan, TrainRunner
from scvi.train._callbacks import SaveBestState

from multimil.dataloaders import GroupAnnDataLoader, GroupDataSplitter
from multimil.module import MILClassifierTorch
from multimil.utils import (
    get_bag_info,
    get_predictions,
    plt_plot_losses,
    prep_minibatch,
    save_predictions_in_adata,
    select_covariates,
    setup_ordinal_regression,
)

import numpy as np
from scvi.data import AnnDataManager
from multimil.dataloaders import GroupDataSplitter
from typing import Optional
import torch
from multimil.model import MILClassifier
# from multimil.training import AdversarialTrainingPlan, TrainRunner
# from multimil.callbacks import SaveBestState, ModelCheckpoint
import warnings


class CustomGroupDataSplitter(GroupDataSplitter):
    """
    Data splitter that ensures samples (patients) are not split between train/val sets.
    Maintains the 80-20 split at the sample_ID level.
    """

    def setup(self, stage: Optional[str] = None):
        """
        Split indices ensuring all data from the same sample_ID stays together.
        Default split is 80% train, 20% validation.
        """
        # Get unique sample IDs
        adata = self.adata_manager.adata
        unique_samples = adata.obs[self.group_column].unique()
        n_samples = len(unique_samples)

        # Calculate number of samples for validation
        n_val_samples = int(np.ceil(n_samples * 0.2))  # 20% of samples

        # Randomly shuffle sample IDs
        rng = np.random.RandomState(seed=42)  # For reproducibility
        shuffled_samples = rng.permutation(unique_samples)

        # Split sample IDs
        val_samples = shuffled_samples[:n_val_samples]
        train_samples = shuffled_samples[n_val_samples:]

        # Get indices for each split
        self.val_idx = adata.obs[adata.obs[self.group_column].isin(
            val_samples)].index
        self.train_idx = adata.obs[adata.obs[self.group_column].isin(
            train_samples)].index

        # Convert to numpy array indices
        self.val_idx = np.where(adata.obs.index.isin(self.val_idx))[0]
        self.train_idx = np.where(adata.obs.index.isin(self.train_idx))[0]

        # Empty test set as we're only doing train/val split
        self.test_idx = np.array([], dtype=int)

        # Verify no sample overlap between splits
        val_samples_check = set(
            adata.obs.iloc[self.val_idx][self.group_column].unique())
        train_samples_check = set(
            adata.obs.iloc[self.train_idx][self.group_column].unique())
        assert len(val_samples_check.intersection(train_samples_check)) == 0, \
            "Found samples that appear in both train and validation sets!"

        # Set up device information
        if self.use_gpu and torch.cuda.is_available():
            self.device = torch.device("cuda")
            self.pin_memory = True
        else:
            self.device = torch.device("cpu")
            self.pin_memory = False


def verify_patient_split(mil, adata):
    # 1. Check no patient overlap between train and validation
    train_patients = set(
        adata.obs.iloc[mil.train_indices_]['sample_ID'].unique())
    val_patients = set(
        adata.obs.iloc[mil.validation_indices_]['sample_ID'].unique())
    patient_overlap = train_patients.intersection(val_patients)
    print("1. Patient overlap between train/val:", len(patient_overlap))
    assert len(
        patient_overlap) == 0, f"Found overlapping patients: {patient_overlap}"

    # 2. Check all tiles from same patient stay together
    for patient in adata.obs['sample_ID'].unique():
        patient_indices = adata.obs['sample_ID'] == patient
        patient_tiles = set(np.where(patient_indices)[0])

        tiles_in_train = patient_tiles.intersection(set(mil.train_indices_))
        tiles_in_val = patient_tiles.intersection(set(mil.validation_indices_))

        assert (len(tiles_in_train) == 0 or len(tiles_in_val) == 0), \
            f"Patient {patient} has tiles split between train ({len(tiles_in_train)}) and val ({len(tiles_in_val)})"

    # 3. Check splitting proportions
    total_patients = len(adata.obs['sample_ID'].unique())
    expected_train_size = int(total_patients * 0.9)
    print(
        f"3. Train patients: {len(train_patients)} (expected ~{expected_train_size})")
    print(
        f"   Val patients: {len(val_patients)} (expected ~{total_patients - expected_train_size})")

    return {
        'train_patients': train_patients,
        'val_patients': val_patients,
        'train_indices': mil.train_indices_,
        'val_indices': mil.validation_indices_
    }


class PatientLevelMILClassifier(MILClassifier):
    def train(
        self,
        max_epochs: int = 200,
        lr: float = 5e-4,
        use_gpu: str | int | bool | None = None,
        train_size: float = 0.9,
        validation_size: float | None = None,
        batch_size: int = 256,
        weight_decay: float = 1e-3,
        eps: float = 1e-08,
        early_stopping: bool = True,
        save_best: bool = True,
        check_val_every_n_epoch: int | None = None,
        n_epochs_kl_warmup: int | None = None,
        n_steps_kl_warmup: int | None = None,
        adversarial_mixing: bool = False,
        plan_kwargs: dict | None = None,
        early_stopping_monitor: str | None = "accuracy_validation",
        early_stopping_mode: str | None = "max",
        save_checkpoint_every_n_epochs: int | None = None,
        path_to_checkpoints: str | None = None,
        **kwargs,
    ):
        if len(self.regression) > 0:
            if early_stopping_monitor == "accuracy_validation":
                warnings.warn(
                    "Setting early_stopping_monitor to 'regression_loss_validation' and early_stopping_mode to 'min' as regression is used.",
                    stacklevel=2,
                )
                early_stopping_monitor = "regression_loss_validation"
                early_stopping_mode = "min"
        if n_epochs_kl_warmup is None:
            n_epochs_kl_warmup = max(max_epochs // 3, 1)
        update_dict = {
            "lr": lr,
            "adversarial_classifier": adversarial_mixing,
            "weight_decay": weight_decay,
            "eps": eps,
            "n_epochs_kl_warmup": n_epochs_kl_warmup,
            "n_steps_kl_warmup": n_steps_kl_warmup,
            "optimizer": "AdamW",
            "scale_adversarial_loss": 1,
        }
        if plan_kwargs is not None:
            plan_kwargs.update(update_dict)
        else:
            plan_kwargs = update_dict

        if save_best:
            if "callbacks" not in kwargs.keys():
                kwargs["callbacks"] = []
            kwargs["callbacks"].append(SaveBestState(monitor=early_stopping_monitor, mode=early_stopping_mode))

        if save_checkpoint_every_n_epochs is not None:
            if path_to_checkpoints is not None:
                kwargs["callbacks"].append(
                    ModelCheckpoint(
                        dirpath=path_to_checkpoints,
                        save_top_k=-1,
                        monitor="epoch",
                        every_n_epochs=save_checkpoint_every_n_epochs,
                        verbose=True,
                    )
                )
            else:
                raise ValueError(
                    f"`save_checkpoint_every_n_epochs` = {save_checkpoint_every_n_epochs} so `path_to_checkpoints` has to be not None but is {path_to_checkpoints}."
                )

        # The only difference is here - using CustomGroupDataSplitter instead of GroupDataSplitter
        data_splitter = CustomGroupDataSplitter(
            self.adata_manager,
            group_column=self.sample_key,
            train_size=train_size,
            validation_size=validation_size,
            batch_size=batch_size,
            use_gpu=use_gpu,
        )

        training_plan = AdversarialTrainingPlan(self.module, **plan_kwargs)
        runner = TrainRunner(
            self,
            training_plan=training_plan,
            data_splitter=data_splitter,
            max_epochs=max_epochs,
            use_gpu=use_gpu,
            early_stopping=early_stopping,
            check_val_every_n_epoch=check_val_every_n_epoch,
            early_stopping_monitor=early_stopping_monitor,
            early_stopping_mode=early_stopping_mode,
            early_stopping_patience=50,
            enable_checkpointing=True,
            **kwargs,
        )
        return runner()

Usage

For illustration, just replace the old code in the "Classification" tutorial with:

PatientLevelMILClassifier.setup_anndata(
    adata,
    categorical_covariate_keys=categorical_covariate_keys,
)
mil = PatientLevelMILClassifier(
    adata,
    classification=classification_keys,
    z_dim=z_dim,
    sample_key=sample_key,
    class_loss_coef=0.1,
)
mil.train(lr=1e-3, max_epochs=5) ## FEW EPOCHS TO SAVE TIME
adata.obs["sample_ID"] = adata.obs["sample"] ## For compatibility with the function verify_patient_split
verify_patient_split(mil, adata) ## CHECK THAT TRAIN/VAL SPLIT IS AT THE PATIENT LEVEL

[alitinet]

Hi, apologies for the late reply. Indeed during training the train/validation split is based on cells and does not take samples into consideration. We are aware of this, and it will be changed in the next release early next year. Thanks for the suggestion!

From our experience, though, setting max_epoch to 50 usually prevents overfitting, so we suggest doing that for now. Also, of course, checking the prediction accuracy on the query data.