Feat/mlp depth loss weight bootstrap - #167
Open
kirilklein wants to merge 5 commits into
Open
Conversation
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review infoConfiguration used: defaults Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a sampling method option ("resample" vs "bootstrap") to cohort sampling and propagates it through simulation; refactors MLPHead to support configurable hidden layers; updates model and trainer to weight exposure loss and average outcome losses; and adds tests for cohort sampling and deeper MLPHead behavior. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Tip Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs). Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
corebehrt/modules/trainer/causal/trainer.py (1)
149-175:⚠️ Potential issue | 🟠 MajorPCGrad path currently drops non-task loss terms (e.g., L1 regularization).
When
task_lossesis populated, backprop only covers exposure/outcome losses. Ifoutputs.lossincludes extra terms (likeoutputs.l1_lossfrom the model), enabling PCGrad silently changes the optimized objective.Proposed fix
# 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 / n_outcomes) + + # Preserve regularization terms included in total loss + if hasattr(outputs, "l1_loss") and outputs.l1_loss is not None: + task_losses.append(outputs.l1_loss)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@corebehrt/modules/trainer/causal/trainer.py` around lines 149 - 175, The PCGrad branch currently only backprops per-task losses (`task_losses`) and thus drops any extra terms present in `outputs.loss` (e.g., L1 regularization); compute the residual extra loss as `residual = outputs.loss - sum(task_losses)` when `hasattr(outputs, "loss")` and the residual is non-negligible, append that residual to `task_losses` before scaling (`scaled_losses = [self.scaler.scale(loss) for loss in task_losses]`) so `self.optimizer.pc_backward(scaled_losses)` includes those terms, and keep the existing return logic (`return outputs.loss if hasattr(outputs, "loss") else sum(task_losses)`) to log the full unscaled objective; reference symbols: task_losses, outputs.loss, outputs.outcome_losses, self.scaler.scale, self.optimizer.pc_backward.
🧹 Nitpick comments (1)
corebehrt/modules/model/causal/heads.py (1)
79-91:dropout_probis currently dead config.The parameter is exposed but never applied in
self.classifier. Either wirenn.Dropout(dropout_prob)into the stack or remove the arg to avoid misleading config behavior.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@corebehrt/modules/model/causal/heads.py` around lines 79 - 91, The parameter dropout_prob is unused in the constructor building self.classifier; either wire it into the Sequential or remove it. Fix by inserting nn.Dropout(dropout_prob) into the layers list (e.g., append nn.Dropout(dropout_prob) immediately after each nn.GELU() inside the for-loop that builds hidden layers in the class constructor, and if num_hidden_layers==0 ensure a dropout is added before the final nn.Linear(current_size, 1) so the config is honored). Alternatively, if you choose to remove the parameter, delete dropout_prob from the signature and any references; ensure the unique symbols referenced are dropout_prob, self.classifier, and the constructor that builds layers.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@corebehrt/functional/causal/cohort_sampler.py`:
- Around line 52-54: Add an explicit empty-cohort validation by checking
full_pids (or n_total) before proceeding: if full_pids is empty (n_total == 0)
raise a clear ValueError so callers fail fast; place this check before the
existing lines that compute n_total = len(full_pids) and replace = method ==
"bootstrap" (or immediately after computing n_total) and include a descriptive
message referencing the cohort being empty.
In `@corebehrt/modules/model/causal/heads.py`:
- Around line 85-90: The loop computing next_size = current_size //
hidden_size_ratio can produce zero or invalid widths; add guards in the MLP
builder that validate inputs (ensure current_size and num_hidden_layers are
positive integers and hidden_size_ratio >= 1) and after computing next_size
check if next_size < 1 (or next_size == 0) and raise a clear ValueError
referencing next_size/current_size/hidden_size_ratio, or clamp next_size to at
least 1 before creating nn.Linear; update the code that appends layers (layers
and nn.Linear) to rely on these validated/adjusted sizes so no zero-width Linear
is created.
In `@tests/test_modules/test_model/test_causal/test_head.py`:
- Around line 221-264: The test module (functions like
test_default_matches_original_architecture, test_multi_layer_forward_pass,
test_multi_layer_module_count, test_gradient_flow_deep_head) needs to be
auto-formatted with Ruff to satisfy CI; run `ruff format` on this test file,
review the resulting whitespace/line-break changes, stage the formatted file,
and commit the updates so the CI no longer reports a reformatting error.
---
Outside diff comments:
In `@corebehrt/modules/trainer/causal/trainer.py`:
- Around line 149-175: The PCGrad branch currently only backprops per-task
losses (`task_losses`) and thus drops any extra terms present in `outputs.loss`
(e.g., L1 regularization); compute the residual extra loss as `residual =
outputs.loss - sum(task_losses)` when `hasattr(outputs, "loss")` and the
residual is non-negligible, append that residual to `task_losses` before scaling
(`scaled_losses = [self.scaler.scale(loss) for loss in task_losses]`) so
`self.optimizer.pc_backward(scaled_losses)` includes those terms, and keep the
existing return logic (`return outputs.loss if hasattr(outputs, "loss") else
sum(task_losses)`) to log the full unscaled objective; reference symbols:
task_losses, outputs.loss, outputs.outcome_losses, self.scaler.scale,
self.optimizer.pc_backward.
---
Nitpick comments:
In `@corebehrt/modules/model/causal/heads.py`:
- Around line 79-91: The parameter dropout_prob is unused in the constructor
building self.classifier; either wire it into the Sequential or remove it. Fix
by inserting nn.Dropout(dropout_prob) into the layers list (e.g., append
nn.Dropout(dropout_prob) immediately after each nn.GELU() inside the for-loop
that builds hidden layers in the class constructor, and if num_hidden_layers==0
ensure a dropout is added before the final nn.Linear(current_size, 1) so the
config is honored). Alternatively, if you choose to remove the parameter, delete
dropout_prob from the signature and any references; ensure the unique symbols
referenced are dropout_prob, self.classifier, and the constructor that builds
layers.
ℹ️ Review info
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
corebehrt/functional/causal/cohort_sampler.pycorebehrt/main_causal/simulate_with_sampling.pycorebehrt/modules/model/causal/heads.pycorebehrt/modules/model/causal/model.pycorebehrt/modules/trainer/causal/trainer.pytests/test_functional/test_causal/test_cohort_sampler.pytests/test_modules/test_model/test_causal/test_head.py
Comment on lines
+85
to
+90
| 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)) |
Contributor
There was a problem hiding this comment.
Add guardrails for invalid/degenerated MLP dimensions.
next_size = current_size // hidden_size_ratio is unchecked. Bad configs can create zero-width layers or runtime errors.
Proposed fix
def __init__(
self,
input_size: int,
hidden_size_ratio: int = 2,
dropout_prob: float = 0.1,
num_hidden_layers: int = 1,
):
super().__init__()
+ if hidden_size_ratio <= 0:
+ raise ValueError(f"hidden_size_ratio must be > 0, got {hidden_size_ratio}")
+ if num_hidden_layers < 0:
+ raise ValueError(f"num_hidden_layers must be >= 0, got {num_hidden_layers}")
+
layers = [nn.LayerNorm(input_size)]
current_size = input_size
- for _ in range(num_hidden_layers):
+ for layer_idx in range(num_hidden_layers):
next_size = current_size // hidden_size_ratio
+ if next_size < 1:
+ raise ValueError(
+ f"Invalid MLP config at hidden layer {layer_idx}: "
+ f"current_size={current_size}, hidden_size_ratio={hidden_size_ratio}"
+ )
layers.append(nn.Linear(current_size, next_size))
layers.append(nn.GELU())
current_size = next_size🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@corebehrt/modules/model/causal/heads.py` around lines 85 - 90, The loop
computing next_size = current_size // hidden_size_ratio can produce zero or
invalid widths; add guards in the MLP builder that validate inputs (ensure
current_size and num_hidden_layers are positive integers and hidden_size_ratio
>= 1) and after computing next_size check if next_size < 1 (or next_size == 0)
and raise a clear ValueError referencing
next_size/current_size/hidden_size_ratio, or clamp next_size to at least 1
before creating nn.Linear; update the code that appends layers (layers and
nn.Linear) to rely on these validated/adjusted sizes so no zero-width Linear is
created.
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/test_modules/test_model/test_causal/test_head.py (1)
240-249: Consider asserting module order/types, not only count.Count-only checks can pass even if layer ordering regresses. A light type-pattern assertion would strengthen this test.
Optional test hardening
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) + self.assertIsInstance(modules[0], torch.nn.LayerNorm) + for i in range(depth): + self.assertIsInstance(modules[1 + 2 * i], torch.nn.Linear) + self.assertIsInstance(modules[2 + 2 * i], torch.nn.GELU) + self.assertIsInstance(modules[-1], torch.nn.Linear)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_modules/test_model/test_causal/test_head.py` around lines 240 - 249, The test test_multi_layer_module_count only checks the number of modules; update it to also assert the expected module types and order for MLPHead: verify head.classifier[0] is LayerNorm, then for i in range(depth) assert the pair head.classifier[1 + 2*i] is Linear and head.classifier[1 + 2*i + 1] is GELU, and finally assert the last module is Linear; use isinstance checks against torch.nn.LayerNorm, torch.nn.Linear, and torch.nn.GELU so the test will fail if ordering or types regress while still supporting varying depth.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tests/test_modules/test_model/test_causal/test_head.py`:
- Around line 240-249: The test test_multi_layer_module_count only checks the
number of modules; update it to also assert the expected module types and order
for MLPHead: verify head.classifier[0] is LayerNorm, then for i in range(depth)
assert the pair head.classifier[1 + 2*i] is Linear and head.classifier[1 + 2*i +
1] is GELU, and finally assert the last module is Linear; use isinstance checks
against torch.nn.LayerNorm, torch.nn.Linear, and torch.nn.GELU so the test will
fail if ordering or types regress while still supporting varying depth.
ℹ️ Review info
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
tests/test_functional/test_causal/test_cohort_sampler.pytests/test_modules/test_model/test_causal/test_head.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/test_functional/test_causal/test_cohort_sampler.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Tests