bench-xECG · Data Splitting & Batch Flow Visualizer

Visual Guide: Pretraining Dataset Splits & DataLoader Shuffling

A step-by-step visual dissection explaining how patient-level dataset partitioning works, why string sorting caused the HEEDB validation anomaly, and how DataLoader batch composition controls self-supervised representation geometry.

Inside a Single Dataset: Patient Indexing & 90/10 Splitting

Level 1: Dataset Partitioning

In bench-xecg, datasets like HEEDB, CODE, and Chapman are indexed by unique patients (not raw ECG recording files). When __getitem__(idx) is called, it retrieves recordings specifically belonging to that patient.

PATIENT LIST PARTITIONING: 10 PATIENT TOY EXAMPLE _split_train_val(dataset, config)

✗ Old Contiguous Slice (Range 0 → 0.9N) No Permutation

Original Patient Array (Lexicographically Sorted):

P01

P02

P03

P04

P05

P06

P07

P08

P09

P10*

↓ Contiguous 90% Slice ↓

Train Set (First 9 / 90%): 100% Hospital 1

P01

P02

P03

P04

P05

P06

P07

P08

P09

—

Val Set (Last 1 / 10%): 100% Hospital 2 (P10*)

P10 (Hospital 2)

Because string sorting grouped all "I0006_..." patients at the very end, taking the tail 10% sent 100% of Hospital 2 to validation and starved training.

✓ Seeded Random Split (_split_train_val) Permutation

Seeded Random Permutation (Seed 0):

P04

P10*

P01

P07

P09

P02

P08

P05

P03

P06

↓ Seeded Permutation 90% / 10% Split ↓

Train Set (90% Sample): Hosp 1 + Hosp 2 Mixed

P04

P10*

P01

P07

P09

P02

P08

P05

P03

—

Val Set (10% Sample): True Representative Sample

P06 (Hosp 1)

Every hospital and demographic cohort is uniformly distributed across both splits with a reproducible fixed seed.

Multi-Dataset Concatenation & DataLoader Streaming

Level 2: Cross-Dataset Batch Construction

In pretraining, multiple distinct datasets are concatenated: ConcatDataset([CODE, Chapman, INCART, HEEDB]). Here is what happens when the DataLoader constructs batches of size \(B=512\) under shuffle=False vs shuffle=True:

CONCATENATED DATASET MEMORY LAYOUT val_dataset = ConcatDataset([val_code, val_chapman, val_incart, val_heedb])

1. CODE Subset (15%)

2. Chapman Subset (20%)

3. INCART (8%)

4. HEEDB Subset (57%)

HOW BATCHES ARE FED INTO THE MODEL

Case A: DataLoader(shuffle=False) — Monolithic Siloed Batches Causes Validation Loss Explosion

Batch 01 (Val) B = 512

100% CODE

Batch Variety: 1 Dataset Only

Batch 25 (Val) B = 512

100% Chapman

Batch Variety: 1 Dataset Only

Batch 42 (Val) B = 512

100% INCART

Batch Variety: 1 Dataset Only

Batch 70 (Val) B = 512

100% HEEDB

Batch Variety: 1 Dataset Only

Case B: Training (shuffle=True) OR Interleaved Validation — Balanced Mixtures True i.i.d. Representation

Batch 01 (Train/Val) B = 512

HEEDB

Batch Variety: Ideal 4-Way Mix

Batch 25 (Train/Val) B = 512

HEEDB

Batch Variety: Ideal 4-Way Mix

Batch 42 (Train/Val) B = 512

HEEDB

Batch Variety: Ideal 4-Way Mix

Batch 70 (Train/Val) B = 512

HEEDB

Batch Variety: Ideal 4-Way Mix

Latent Space Geometry: Why Batch-Level Losses Blow Up

Level 3: Loss Function Mathematics

Why does batch mixing matter so much for LeJEPA (SIGReg) and SimDINOv2 (Expansion), while standard losses don't care?

Homogeneous / Unmixed Batch SIGReg Loss: 216+

Zero Variance in other dimensions • Extreme Skewness

SIGReg evaluates: \(\mathcal{L}_{\text{SIGReg}} = \|\text{Cov}(\mathbf{Z}_B) - \mathbf{I}\|_F^2 + \dots\)
When 512 samples belong to one single hospital, all points collapse into one sector of the sphere. The covariance matrix has near-zero eigenvalues, and SIGReg flags this as collapsed variance.

Balanced / Mixed Batch SIGReg Loss: ~11–15

Full Rank Isotropic Covariance • \(\mathbf{Z}_B \sim \mathcal{N}(0, \mathbf{I})\)

Expected Geometry: With a uniform mixture across datasets and patients, the batch fills all orthogonal dimensions of the embedding space, matching standard Gaussian properties.

The Implemented Solution in bench-xecg

Resolved

Pipeline Component	Implementation Code	What It Solves
1. Seeded Patient Split	`_split_train_val()` with `np.random.default_rng(seed).permutation(n)`	Prevents alphabetical ID sorting artifacts in HEEDB, CODE, and Chapman. Ensures both train and val receive equal proportional slices of all hospital cohorts.
2. Val Dataset Interleaving	`val_dataset = Subset(val_dataset, val_order)` in `load_datasets()`	Ensures every validation batch contains a proportional mixture across all concatenated datasets, preventing batch-level covariance collapse.
3. Training DataLoader	`DataLoader(shuffle=True)` in `pretrain.py:170`	Continuously randomizes batches across all patients and datasets every training epoch.