Lazy-fit model_mu in R-learner so it's only trained when return_components is used#936
Lazy-fit model_mu in R-learner so it's only trained when return_components is used#936aman-coder03 wants to merge 2 commits into
model_mu in R-learner so it's only trained when return_components is used#936Conversation
jeongyoonlee
left a comment
There was a problem hiding this comment.
Deferring the model_mu fit requires pinning the full training set on the instance (self._mu_fit_X, self._mu_fit_y), which inflates a fitted R-learner from 2.7 KiB to 4.8 MiB at n=20k/p=30 (pickle). With save()/load() added in #929/#932, every saved R-learner now carries its training data — a steep cost to avoid one model_mu.fit(X, y), which is cheap for the default learners. Either keep the eager fit, or exclude the cached X/y from serialization so on-disk size doesn't regress.
The new tests cover only BaseRLearner; BaseRClassifier and XGBRRegressor get the same lazy machinery but no test — please add at least the classifier path.
Note this overlaps #935, which also rewrites BaseRLearner.fit() (and keeps the eager model_mu fit this PR removes) — they'll need to be sequenced.
Proposed changes
follow up to #923
model_mu(the nuisance outcome model) was being fit on the full training data insidefit()every time, even when nobody ever callspredict(..., return_components=True),that's wasted work for the common casethis defers the fit:
fit()now just caches the training data, andmodel_muonly gets fit the first timepredict(return_components=True)is actually called, then cached so it's not re-fit on later callsalso fixes a related bug this uncovered: the bootstrap loops in
fit_predict(return_ci=True)andestimate_ate(bootstrap_ci=True)save/restoremodel_muandmodels_tauaround resampling, but weren't saving/restoring the new cached training data used for the lazy fit. without the fix, callingpredict(return_components=True)after a bootstrap run would silently fitmodel_muon stale bootstrap-resampled data instead of the real training setapplies to
BaseRLearner,BaseRClassifier, andXGBRRegressorTypes of changes
What types of changes does your code introduce to CausalML?
Put an
xin the boxes that applyChecklist
Put an
xin the boxes that apply. You can also fill these out after creating the PR. If you're unsure about any of them, don't hesitate to ask. We're here to help! This is simply a reminder of what we are going to look for before merging your code.Further comments
If this is a relatively large or complex change, kick off the discussion by explaining why you chose the solution you did and what alternatives you considered, etc. This PR template is adopted from appium.