Example: Reusing base models

Motivation

In our Why MetaLearners section we praise the modularity of MetaLearners. Part of the reason why modularity is useful is because we can actively decouple different parts of the CATE estimation process.

Concretely, this decoupling allows for saving lots of compute resources: if we know that we merely want to change some parts of a MetaLearner, we may as well reuse the parts that we don’t want to change. Enabling this kind of base model reuse was one of the requirements on metalearners, see Why not causalml or econml.

For instance, imagine trying to tune an R-Learner’s – consisting of two nuisance models, a propensity model and an outcome model – propensity model with respect to its R-Loss. In such a scenario we would like to reuse the same outcome model because it isn’t affected by the propensity model and thereby save a lot of redundant compute.

Example

Loading the data

Just like in our example on estimating CATEs with a MetaLearner, we will first load some experiment data:

import pandas as pd
from pathlib import Path
from git_root import git_root

df = pd.read_csv(git_root("data/learning_mindset.zip"))
outcome_column = "achievement_score"
treatment_column = "intervention"
feature_columns = [
    column
    for column in df.columns
    if column not in [outcome_column, treatment_column]
]
categorical_feature_columns = [
    "ethnicity",
    "gender",
    "frst_in_family",
    "school_urbanicity",
    "schoolid",
]
# Note that explicitly setting the dtype of these features to category
# allows both lightgbm as well as shap plots to
# 1. Operate on features which are not of type int, bool or float
# 2. Correctly interpret categoricals with int values to be
#    interpreted as categoricals, as compared to ordinals/numericals.
for categorical_feature_column in categorical_feature_columns:
    df[categorical_feature_column] = df[categorical_feature_column].astype(
        "category"
    )

Now that we’ve loaded the experiment data, we can train a MetaLearner.

Training a first MetaLearner

Again, mirroring our example on estimating CATEs with a MetaLearner, we can train an RLearner as follows:

from metalearners import RLearner
from lightgbm import LGBMRegressor, LGBMClassifier
rlearner = RLearner(
    nuisance_model_factory=LGBMRegressor,
    propensity_model_factory=LGBMClassifier,
    treatment_model_factory=LGBMRegressor,
    is_classification=False,
    n_variants=2,
    nuisance_model_params={"verbose": -1},
    propensity_model_params={"verbose": -1},
    treatment_model_params={"verbose": -1},
)

rlearner.fit(
    X=df[feature_columns],
    y=df[outcome_column],
    w=df[treatment_column],
)

<metalearners.rlearner.RLearner at 0x7fe784db12b0>

By virtue of having fitted the ‘overall’ MetaLearner, we fitted the base model, too. Thereby we can now reuse some of them if we wish to.

Extracting a basel model from a trained MetaLearner

In order to reuse a base model from one MetaLearner for another MetaLearner, we first have to from the former. If, for instance, we are interested in reusing the outcome nuisance model of the RLearner we just trained, we can access it via its _nuisance_models attribute:

rlearner._nuisance_models

{'outcome_model': [CrossFitEstimator(n_folds=10, estimator_factory=<class 'lightgbm.sklearn.LGBMRegressor'>, estimator_params={'verbose': -1}, enable_overall=True, random_state=None, _estimators=[LGBMRegressor(verbose=-1), LGBMRegressor(verbose=-1), LGBMRegressor(verbose=-1), LGBMRegressor(verbose=-1), LGBMRegressor(verbose=-1), LGBMRegressor(verbose=-1), LGBMRegressor(verbose=-1), LGBMRegressor(verbose=-1), LGBMRegressor(verbose=-1), LGBMRegressor(verbose=-1)], _estimator_type='regressor', _overall_estimator=LGBMRegressor(verbose=-1), _test_indices=(array([   39,    43,    51, ..., 10367, 10371, 10373]), array([    6,    28,    34, ..., 10353, 10357, 10383]), array([    9,    14,    19, ..., 10384, 10385, 10386]), array([    0,     4,     5, ..., 10382, 10387, 10388]), array([    7,    10,    25, ..., 10376, 10377, 10381]), array([    3,    21,    23, ..., 10350, 10354, 10380]), array([    2,     8,    31, ..., 10351, 10356, 10366]), array([    1,    12,    15, ..., 10352, 10369, 10372]), array([   24,    27,    36, ..., 10348, 10364, 10389]), array([   11,    30,    35, ..., 10363, 10378, 10390])), _n_classes=None, classes_=None)],
 'propensity_model': [CrossFitEstimator(n_folds=10, estimator_factory=<class 'lightgbm.sklearn.LGBMClassifier'>, estimator_params={'verbose': -1}, enable_overall=True, random_state=None, _estimators=[LGBMClassifier(verbose=-1), LGBMClassifier(verbose=-1), LGBMClassifier(verbose=-1), LGBMClassifier(verbose=-1), LGBMClassifier(verbose=-1), LGBMClassifier(verbose=-1), LGBMClassifier(verbose=-1), LGBMClassifier(verbose=-1), LGBMClassifier(verbose=-1), LGBMClassifier(verbose=-1)], _estimator_type='classifier', _overall_estimator=LGBMClassifier(verbose=-1), _test_indices=(array([   39,    43,    51, ..., 10367, 10371, 10373]), array([    6,    28,    34, ..., 10353, 10357, 10383]), array([    9,    14,    19, ..., 10384, 10385, 10386]), array([    0,     4,     5, ..., 10382, 10387, 10388]), array([    7,    10,    25, ..., 10376, 10377, 10381]), array([    3,    21,    23, ..., 10350, 10354, 10380]), array([    2,     8,    31, ..., 10351, 10356, 10366]), array([    1,    12,    15, ..., 10352, 10369, 10372]), array([   24,    27,    36, ..., 10348, 10364, 10389]), array([   11,    30,    35, ..., 10363, 10378, 10390])), _n_classes=2, classes_=array([0, 1]))]}

We notice that the RLearner has two kinds of nuisance models: "propensity_model" and "outcome_model". Note that we could’ve figured this out by calling its nuisance_model_specifications() method, too.

Therefore, we now know how to fetch our outcome model:

outcome_model = rlearner._nuisance_models["outcome_model"]

Training a second MetaLearner by reusing a base model

Given that we know have an already trained outcome model, we can reuse for another ‘kind’ of RLearner on the same data. Concretely, we will now want to use a different propensity_model_factory and nuisance_model_factory. Note that this time, we do not specify a nuisance_model_factory in the initialization of the RLearner since the RLearner only relies on a single non-propensity nuisance model. This might vary for other MetaLearners, such as the DRLearner.

from sklearn.linear_model import LinearRegression, LogisticRegression

rlearner_new = RLearner(
    propensity_model_factory=LogisticRegression,
    treatment_model_factory=LinearRegression,
    is_classification=False,
    fitted_nuisance_models={"outcome_model": outcome_model},
    propensity_model_params={"max_iter": 500},
    n_variants=2,
)

rlearner_new.fit(
    X=df[feature_columns],
    y=df[outcome_column],
    w=df[treatment_column],
)

<metalearners.rlearner.RLearner at 0x7fe797e1dd00>

What’s more is that we can also reuse models between different kinds of MetaLearner architectures. A propensity model, for instance, is used in many scenarios. Let’s reuse it for a DRLearner:

from metalearners import DRLearner

trained_propensity_model = rlearner._nuisance_models["propensity_model"][0]

drlearner = DRLearner(
    nuisance_model_factory=LGBMRegressor,
    treatment_model_factory=LGBMRegressor,
    fitted_propensity_model=trained_propensity_model,
    is_classification=False,
    n_variants=2,
)

drlearner.fit(
    X=df[feature_columns],
    y=df[outcome_column],
    w=df[treatment_column],
)

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000376 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 6306, number of used features: 11
[LightGBM] [Info] Start training from score -0.157471
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000255 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 6306, number of used features: 11
[LightGBM] [Info] Start training from score -0.159419

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000255 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 6306, number of used features: 11
[LightGBM] [Info] Start training from score -0.164724
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000255 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 6306, number of used features: 11
[LightGBM] [Info] Start training from score -0.153577

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000326 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 6306, number of used features: 11
[LightGBM] [Info] Start training from score -0.151392
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000254 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 6306, number of used features: 11
[LightGBM] [Info] Start training from score -0.147012

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000255 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 6306, number of used features: 11
[LightGBM] [Info] Start training from score -0.147527
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000270 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 6307, number of used features: 11
[LightGBM] [Info] Start training from score -0.156169

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000281 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 6307, number of used features: 11
[LightGBM] [Info] Start training from score -0.148469
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000262 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 6307, number of used features: 11
[LightGBM] [Info] Start training from score -0.152270

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000324 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 7007, number of used features: 11
[LightGBM] [Info] Start training from score -0.153803
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000124 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 3045, number of used features: 11
[LightGBM] [Info] Start training from score 0.326819
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000127 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 3045, number of used features: 11
[LightGBM] [Info] Start training from score 0.312533

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000125 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 3045, number of used features: 11
[LightGBM] [Info] Start training from score 0.319549
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000141 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 3045, number of used features: 11
[LightGBM] [Info] Start training from score 0.312079
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000124 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 3046, number of used features: 11
[LightGBM] [Info] Start training from score 0.311451

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000124 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 3046, number of used features: 11
[LightGBM] [Info] Start training from score 0.320800
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000125 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 3046, number of used features: 11
[LightGBM] [Info] Start training from score 0.323098
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000124 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 3046, number of used features: 11
[LightGBM] [Info] Start training from score 0.322901

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000128 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 3046, number of used features: 11
[LightGBM] [Info] Start training from score 0.313838
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000128 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 491
[LightGBM] [Info] Number of data points in the train set: 3046, number of used features: 11
[LightGBM] [Info] Start training from score 0.321617
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000144 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 3384, number of used features: 11
[LightGBM] [Info] Start training from score 0.318469

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000368 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 9351, number of used features: 11
[LightGBM] [Info] Start training from score 0.382315
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000366 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 9352, number of used features: 11
[LightGBM] [Info] Start training from score 0.380391

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000367 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 9352, number of used features: 11
[LightGBM] [Info] Start training from score 0.381870
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000379 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 9352, number of used features: 11
[LightGBM] [Info] Start training from score 0.387085

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000397 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 9352, number of used features: 11
[LightGBM] [Info] Start training from score 0.391571
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000368 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 9352, number of used features: 11
[LightGBM] [Info] Start training from score 0.369922

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000367 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 9352, number of used features: 11
[LightGBM] [Info] Start training from score 0.382253
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000392 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 9352, number of used features: 11
[LightGBM] [Info] Start training from score 0.380612

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000399 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 9352, number of used features: 11
[LightGBM] [Info] Start training from score 0.376130
[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000397 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 9352, number of used features: 11
[LightGBM] [Info] Start training from score 0.384191

[LightGBM] [Info] Auto-choosing row-wise multi-threading, the overhead of testing was 0.000449 seconds.
You can set `force_row_wise=true` to remove the overhead.
And if memory is not enough, you can set `force_col_wise=true`.
[LightGBM] [Info] Total Bins 496
[LightGBM] [Info] Number of data points in the train set: 10391, number of used features: 11
[LightGBM] [Info] Start training from score 0.381634

<metalearners.drlearner.DRLearner at 0x7fe797e1e000>

Further comments

Note that the nuisance models are always expected to be of type CrossFitEstimator. More precisely, the when extracting or passing a particular model kind, we pass a list of CrossFitEstimator unless it is the propensity model.
In the examples above we reused nuisance models trained as part of a call to a MetaLearners overall fit() method. If one wants to train a nuisance model in isolation (i.e. not through a MetaLearner) to be used in a MetaLearner afterwards, one should do it by instantiating CrossFitEstimator.
Additionally, individual nuisance models can be trained via a MetaLearner’s fit_nuisance() method.
We strongly recommend only reusing base models if they have been trained on exactly the same data. If this is not the case, some functionalities will probably not work as hoped for.
Note that only nuisance models can be reused, not treatment models.