From ce0c6b1da102955160e2f51708345abac050fd50 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Mon, 23 Feb 2026 08:49:24 +0100 Subject: [PATCH 1/5] Add configurable MLP head depth for better confounding adjustment MLPHead now accepts num_hidden_layers parameter (default 1, matching previous architecture). Deeper heads give the outcome model more capacity to learn the exposure-outcome relationship. Config keys mlp_num_hidden_layers and mlp_hidden_size_ratio are read from head_config. Co-Authored-By: Claude Opus 4.6 --- corebehrt/modules/model/causal/heads.py | 23 ++++++---- corebehrt/modules/model/causal/model.py | 26 +++++++++-- .../test_model/test_causal/test_head.py | 45 +++++++++++++++++++ 3 files changed, 81 insertions(+), 13 deletions(-) diff --git a/corebehrt/modules/model/causal/heads.py b/corebehrt/modules/model/causal/heads.py index e38791e9..b58fe6b4 100644 --- a/corebehrt/modules/model/causal/heads.py +++ b/corebehrt/modules/model/causal/heads.py @@ -73,17 +73,22 @@ class MLPHead(nn.Module): """ def __init__( - self, input_size: int, hidden_size_ratio: int = 2, dropout_prob: float = 0.1 + self, + input_size: int, + hidden_size_ratio: int = 2, + dropout_prob: float = 0.1, + num_hidden_layers: int = 1, ): super().__init__() - intermediate_size = input_size // hidden_size_ratio - self.classifier = nn.Sequential( - nn.LayerNorm(input_size), - nn.Linear(input_size, intermediate_size), - nn.GELU(), - # nn.Dropout(dropout_prob), - nn.Linear(intermediate_size, 1, bias=True), - ) + layers = [nn.LayerNorm(input_size)] + current_size = input_size + for _ in range(num_hidden_layers): + next_size = current_size // hidden_size_ratio + layers.append(nn.Linear(current_size, next_size)) + layers.append(nn.GELU()) + current_size = next_size + layers.append(nn.Linear(current_size, 1, bias=True)) + self.classifier = nn.Sequential(*layers) def forward(self, x: torch.Tensor) -> torch.Tensor: """ diff --git a/corebehrt/modules/model/causal/model.py b/corebehrt/modules/model/causal/model.py index e2376554..d3f7080b 100644 --- a/corebehrt/modules/model/causal/model.py +++ b/corebehrt/modules/model/causal/model.py @@ -48,6 +48,7 @@ def __init__(self, config): self.shared_representation = self.head_config.get("shared_representation", True) self.bidirectional = self.head_config.get("bidirectional", True) self.bottleneck_dim = self.head_config.get("bottleneck_dim", 128) + self.exposure_loss_weight = self.head_config.get("exposure_loss_weight", 1.0) self.l1_lambda = self.head_config.get("l1_lambda", 0.0) if self.l1_lambda > 0: logger.info(f"Applying L1 regularization with lambda={self.l1_lambda}") @@ -103,12 +104,23 @@ def _setup_mlp_heads(self, config): else: head_input_size = config.hidden_size - self.exposure_head = MLPHead(input_size=head_input_size) + mlp_num_hidden_layers = self.head_config.get("mlp_num_hidden_layers", 1) + mlp_hidden_size_ratio = self.head_config.get("mlp_hidden_size_ratio", 2) + + self.exposure_head = MLPHead( + input_size=head_input_size, + num_hidden_layers=mlp_num_hidden_layers, + hidden_size_ratio=mlp_hidden_size_ratio, + ) # Create separate heads for each outcome self.outcome_heads = nn.ModuleDict() for outcome_name in self.outcome_names: - self.outcome_heads[outcome_name] = MLPHead(input_size=head_input_size + 1) + self.outcome_heads[outcome_name] = MLPHead( + input_size=head_input_size + 1, + num_hidden_layers=mlp_num_hidden_layers, + hidden_size_ratio=mlp_hidden_size_ratio, + ) def _get_loss_fn(self, loss_name: str, loss_params: dict): """Returns the loss function instance based on the name.""" @@ -256,9 +268,11 @@ def _compute_losses(self, outputs, batch): batch[EXPOSURE_TARGET].view(-1), ) outputs.exposure_loss = exposure_loss - total_loss += exposure_loss + total_loss += self.exposure_loss_weight * exposure_loss # Only compute outcome losses for available labels + outcome_loss_sum = 0 + n_outcomes = 0 for outcome_name in self.outcome_names: if outcome_name not in batch: continue @@ -269,7 +283,11 @@ def _compute_losses(self, outputs, batch): ) outputs.outcome_losses[outcome_name] = outcome_loss - total_loss += outcome_loss + outcome_loss_sum += outcome_loss + n_outcomes += 1 + + if n_outcomes > 0: + total_loss += outcome_loss_sum / n_outcomes # Add L1 regularization on the bottleneck representation if self.shared_representation and self.l1_lambda > 0: diff --git a/tests/test_modules/test_model/test_causal/test_head.py b/tests/test_modules/test_model/test_causal/test_head.py index 47e43966..9896533d 100644 --- a/tests/test_modules/test_model/test_causal/test_head.py +++ b/tests/test_modules/test_model/test_causal/test_head.py @@ -218,6 +218,51 @@ def test_output_range(self): # Check that outputs are not extremely large (arbitrary threshold) self.assertTrue(torch.abs(output).max() < 100) + def test_default_matches_original_architecture(self): + """Test that default (num_hidden_layers=1) produces 4 modules: LayerNorm, Linear, GELU, Linear.""" + head = MLPHead(self.input_size) + modules = list(head.classifier) + self.assertEqual(len(modules), 4) + self.assertIsInstance(modules[0], torch.nn.LayerNorm) + self.assertIsInstance(modules[1], torch.nn.Linear) + self.assertIsInstance(modules[2], torch.nn.GELU) + self.assertIsInstance(modules[3], torch.nn.Linear) + + def test_multi_layer_forward_pass(self): + """Test forward pass with different depths.""" + for depth in [1, 2, 3]: + with self.subTest(depth=depth): + head = MLPHead( + self.input_size, num_hidden_layers=depth + ).to(self.device) + x = torch.randn(self.batch_size, self.input_size, device=self.device) + output = head(x) + self.assertEqual(output.shape, (self.batch_size, 1)) + + def test_multi_layer_module_count(self): + """Test that deeper heads have more modules.""" + for depth in [1, 2, 3]: + with self.subTest(depth=depth): + head = MLPHead(self.input_size, num_hidden_layers=depth) + modules = list(head.classifier) + # LayerNorm + (Linear + GELU) * depth + Linear + expected = 1 + 2 * depth + 1 + self.assertEqual(len(modules), expected) + + def test_gradient_flow_deep_head(self): + """Test that gradients flow through a deep MLPHead.""" + head = MLPHead(self.input_size, num_hidden_layers=3).to(self.device) + x = torch.randn( + self.batch_size, self.input_size, device=self.device, requires_grad=True + ) + + output = head(x) + loss = output.sum() + loss.backward() + + self.assertIsNotNone(x.grad) + self.assertFalse(torch.allclose(x.grad, torch.zeros_like(x.grad))) + if __name__ == "__main__": unittest.main() From 5711724adb6b18f9612b470945eecbae5e3befe0 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Mon, 23 Feb 2026 08:51:58 +0100 Subject: [PATCH 2/5] Normalize outcome losses and add exposure loss weighting Outcomes are now averaged instead of summed, preventing exposure from being dominated when many outcomes are present. exposure_loss_weight config (default 1.0) controls the balance. PCGrad path applies the same weighting for consistency. Co-Authored-By: Claude Opus 4.6 --- corebehrt/modules/trainer/causal/trainer.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/corebehrt/modules/trainer/causal/trainer.py b/corebehrt/modules/trainer/causal/trainer.py index eb7d2d79..fc81304a 100644 --- a/corebehrt/modules/trainer/causal/trainer.py +++ b/corebehrt/modules/trainer/causal/trainer.py @@ -140,15 +140,18 @@ def _train_step(self, batch: dict): # Collect individual task losses for PCGrad task_losses = [] - # Add exposure loss if available + # Add exposure loss if available (weighted) if hasattr(outputs, "exposure_loss") and outputs.exposure_loss is not None: - task_losses.append(outputs.exposure_loss) + task_losses.append( + self.model.exposure_loss_weight * outputs.exposure_loss + ) - # Add outcome losses if available + # Add outcome losses if available (normalized by count) if hasattr(outputs, "outcome_losses") and outputs.outcome_losses: + n_outcomes = len(outputs.outcome_losses) for outcome_loss in outputs.outcome_losses.values(): if outcome_loss is not None: - task_losses.append(outcome_loss) + task_losses.append(outcome_loss / n_outcomes) # If we don't have individual losses, fall back to total loss if not task_losses and hasattr(outputs, "loss"): From 94d346c24775827825469a292e16b3494f7f3d96 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Mon, 23 Feb 2026 08:52:23 +0100 Subject: [PATCH 3/5] Add bootstrap sampling method for variance estimation sample_cohort now accepts method parameter: "resample" (default, without replacement) or "bootstrap" (with replacement). Bootstrap allows oversampling and enables proper variance estimation for causal effect estimates. Co-Authored-By: Claude Opus 4.6 --- corebehrt/functional/causal/cohort_sampler.py | 27 +++---- .../main_causal/simulate_with_sampling.py | 13 +++- .../test_causal/test_cohort_sampler.py | 75 +++++++++++++++++++ 3 files changed, 96 insertions(+), 19 deletions(-) create mode 100644 tests/test_functional/test_causal/test_cohort_sampler.py diff --git a/corebehrt/functional/causal/cohort_sampler.py b/corebehrt/functional/causal/cohort_sampler.py index 61038f1c..e2d26db8 100644 --- a/corebehrt/functional/causal/cohort_sampler.py +++ b/corebehrt/functional/causal/cohort_sampler.py @@ -14,9 +14,10 @@ def sample_cohort( sample_fraction: float = None, sample_size: int = None, seed: int = 42, + method: str = "resample", ) -> torch.Tensor: """ - Randomly sample patient IDs by fraction or absolute size without replacement. + Randomly sample patient IDs by fraction or absolute size. Args: full_pids: Tensor of all patient IDs in the full cohort @@ -24,24 +25,19 @@ def sample_cohort( Ignored if sample_size is provided. sample_size: Absolute number of patients to sample. Takes precedence over sample_fraction. seed: Random seed for reproducibility + method: Sampling method - "resample" (without replacement) or "bootstrap" (with replacement) Returns: Tensor of sampled patient IDs Raises: ValueError: If neither or both parameters are provided, or values are invalid - - Examples: - >>> full_pids = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) - >>> # Sample by fraction - >>> sampled = sample_cohort(full_pids, sample_fraction=0.5, seed=42) - >>> len(sampled) == 5 - True - >>> # Sample by absolute size - >>> sampled = sample_cohort(full_pids, sample_size=3, seed=42) - >>> len(sampled) == 3 - True """ + if method not in ("resample", "bootstrap"): + raise ValueError( + f"method must be 'resample' or 'bootstrap', got '{method}'" + ) + # Validate inputs if sample_size is None and sample_fraction is None: raise ValueError("Either 'sample_size' or 'sample_fraction' must be provided") @@ -54,12 +50,13 @@ def sample_cohort( # Calculate number of patients to sample n_total = len(full_pids) + replace = method == "bootstrap" if sample_size is not None: # Use absolute size if sample_size <= 0: raise ValueError(f"sample_size must be positive, got {sample_size}") - if sample_size > n_total: + if not replace and sample_size > n_total: raise ValueError( f"sample_size {sample_size} exceeds available patients {n_total}" ) @@ -76,8 +73,8 @@ def sample_cohort( f"sample_fraction {sample_fraction} results in 0 samples from {n_total} patients" ) - # Sample indices without replacement - sampled_indices = rng.choice(n_total, size=n_sample, replace=False) + # Sample indices + sampled_indices = rng.choice(n_total, size=n_sample, replace=replace) # Return sampled PIDs return full_pids[sampled_indices] diff --git a/corebehrt/main_causal/simulate_with_sampling.py b/corebehrt/main_causal/simulate_with_sampling.py index c471ff61..3ae2de1a 100644 --- a/corebehrt/main_causal/simulate_with_sampling.py +++ b/corebehrt/main_causal/simulate_with_sampling.py @@ -57,17 +57,22 @@ def main_simulate(config_path): # Sample subset of PIDs (either by size or fraction) sample_size = sampling_cfg.get("size", None) sample_fraction = sampling_cfg.get("fraction", None) + method = sampling_cfg.get("method", "resample") seed = cfg.get("seed", 42) if sample_size is not None: - logger.info(f"Sampling {sample_size} patients with seed {seed}") - sampled_pids = sample_cohort(full_pids, sample_size=sample_size, seed=seed) + logger.info( + f"Sampling {sample_size} patients with seed {seed} (method={method})" + ) + sampled_pids = sample_cohort( + full_pids, sample_size=sample_size, seed=seed, method=method + ) elif sample_fraction is not None: logger.info( - f"Sampling {sample_fraction * 100:.1f}% of patients with seed {seed}" + f"Sampling {sample_fraction * 100:.1f}% of patients with seed {seed} (method={method})" ) sampled_pids = sample_cohort( - full_pids, sample_fraction=sample_fraction, seed=seed + full_pids, sample_fraction=sample_fraction, seed=seed, method=method ) else: raise ValueError( diff --git a/tests/test_functional/test_causal/test_cohort_sampler.py b/tests/test_functional/test_causal/test_cohort_sampler.py new file mode 100644 index 00000000..edcfe3f0 --- /dev/null +++ b/tests/test_functional/test_causal/test_cohort_sampler.py @@ -0,0 +1,75 @@ +"""Unit tests for cohort sampling utilities.""" + +import unittest + +import torch + +from corebehrt.functional.causal.cohort_sampler import sample_cohort + + +class TestSampleCohort(unittest.TestCase): + + def setUp(self): + self.full_pids = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) + + def test_default_method_is_resample(self): + """Default method should be resample (without replacement).""" + sampled = sample_cohort(self.full_pids, sample_size=5, seed=42) + self.assertEqual(len(sampled), 5) + # No duplicates + self.assertEqual(len(set(sampled.tolist())), 5) + + def test_resample_no_duplicates(self): + """Resample method should produce no duplicate PIDs.""" + sampled = sample_cohort( + self.full_pids, sample_size=10, seed=42, method="resample" + ) + self.assertEqual(len(set(sampled.tolist())), 10) + + def test_resample_rejects_oversize(self): + """Resample method should reject sample_size > n_total.""" + with self.assertRaises(ValueError): + sample_cohort( + self.full_pids, sample_size=20, seed=42, method="resample" + ) + + def test_bootstrap_allows_duplicates(self): + """Bootstrap method should sample with replacement (duplicates possible).""" + # Use a large sample to make duplicates very likely + sampled = sample_cohort( + self.full_pids, sample_size=100, seed=42, method="bootstrap" + ) + self.assertEqual(len(sampled), 100) + # With 100 draws from 10 items, duplicates are virtually guaranteed + self.assertLess(len(set(sampled.tolist())), 100) + + def test_bootstrap_accepts_oversize(self): + """Bootstrap method should accept sample_size > n_total.""" + sampled = sample_cohort( + self.full_pids, sample_size=20, seed=42, method="bootstrap" + ) + self.assertEqual(len(sampled), 20) + + def test_invalid_method_raises(self): + """Invalid method should raise ValueError.""" + with self.assertRaises(ValueError): + sample_cohort( + self.full_pids, sample_size=5, seed=42, method="invalid" + ) + + def test_bootstrap_with_fraction(self): + """Bootstrap with fraction should work.""" + sampled = sample_cohort( + self.full_pids, sample_fraction=0.5, seed=42, method="bootstrap" + ) + self.assertEqual(len(sampled), 5) + + def test_resample_reproducibility(self): + """Same seed should produce same results.""" + s1 = sample_cohort(self.full_pids, sample_size=5, seed=123) + s2 = sample_cohort(self.full_pids, sample_size=5, seed=123) + self.assertTrue(torch.equal(s1, s2)) + + +if __name__ == "__main__": + unittest.main() From e4f05413fe8530b4c98bdb5bfbb2c66a636b3163 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Mon, 2 Mar 2026 10:27:30 +0100 Subject: [PATCH 4/5] ruff --- tests/test_functional/test_causal/test_cohort_sampler.py | 9 ++------- tests/test_modules/test_model/test_causal/test_head.py | 4 +--- 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/tests/test_functional/test_causal/test_cohort_sampler.py b/tests/test_functional/test_causal/test_cohort_sampler.py index edcfe3f0..760d07aa 100644 --- a/tests/test_functional/test_causal/test_cohort_sampler.py +++ b/tests/test_functional/test_causal/test_cohort_sampler.py @@ -8,7 +8,6 @@ class TestSampleCohort(unittest.TestCase): - def setUp(self): self.full_pids = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) @@ -29,9 +28,7 @@ def test_resample_no_duplicates(self): def test_resample_rejects_oversize(self): """Resample method should reject sample_size > n_total.""" with self.assertRaises(ValueError): - sample_cohort( - self.full_pids, sample_size=20, seed=42, method="resample" - ) + sample_cohort(self.full_pids, sample_size=20, seed=42, method="resample") def test_bootstrap_allows_duplicates(self): """Bootstrap method should sample with replacement (duplicates possible).""" @@ -53,9 +50,7 @@ def test_bootstrap_accepts_oversize(self): def test_invalid_method_raises(self): """Invalid method should raise ValueError.""" with self.assertRaises(ValueError): - sample_cohort( - self.full_pids, sample_size=5, seed=42, method="invalid" - ) + sample_cohort(self.full_pids, sample_size=5, seed=42, method="invalid") def test_bootstrap_with_fraction(self): """Bootstrap with fraction should work.""" diff --git a/tests/test_modules/test_model/test_causal/test_head.py b/tests/test_modules/test_model/test_causal/test_head.py index 9896533d..2e2b0cc8 100644 --- a/tests/test_modules/test_model/test_causal/test_head.py +++ b/tests/test_modules/test_model/test_causal/test_head.py @@ -232,9 +232,7 @@ def test_multi_layer_forward_pass(self): """Test forward pass with different depths.""" for depth in [1, 2, 3]: with self.subTest(depth=depth): - head = MLPHead( - self.input_size, num_hidden_layers=depth - ).to(self.device) + head = MLPHead(self.input_size, num_hidden_layers=depth).to(self.device) x = torch.randn(self.batch_size, self.input_size, device=self.device) output = head(x) self.assertEqual(output.shape, (self.batch_size, 1)) From d514b75282e916ef266f84fe1a1283173581c666 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Tue, 3 Mar 2026 08:57:47 +0100 Subject: [PATCH 5/5] format --- corebehrt/functional/causal/cohort_sampler.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/corebehrt/functional/causal/cohort_sampler.py b/corebehrt/functional/causal/cohort_sampler.py index e2d26db8..36afbb12 100644 --- a/corebehrt/functional/causal/cohort_sampler.py +++ b/corebehrt/functional/causal/cohort_sampler.py @@ -34,9 +34,7 @@ def sample_cohort( ValueError: If neither or both parameters are provided, or values are invalid """ if method not in ("resample", "bootstrap"): - raise ValueError( - f"method must be 'resample' or 'bootstrap', got '{method}'" - ) + raise ValueError(f"method must be 'resample' or 'bootstrap', got '{method}'") # Validate inputs if sample_size is None and sample_fraction is None: