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:

In [1]:

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"
    )

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:

In [2]:

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],
)

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], )

Out[2]:

<metalearners.rlearner.RLearner at 0x11470d110>

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:

In [3]:

rlearner._nuisance_models

rlearner._nuisance_models

Out[3]:

{'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([   13,    27,    37, ..., 10364, 10379, 10386]), array([    6,    10,    16, ..., 10373, 10374, 10389]), array([    5,     9,    42, ..., 10363, 10369, 10370]), array([   43,    49,    50, ..., 10366, 10387, 10388]), array([    0,    18,    31, ..., 10361, 10372, 10375]), array([   11,    14,    20, ..., 10345, 10368, 10378]), array([    4,    12,    35, ..., 10377, 10384, 10390]), array([   17,    24,    52, ..., 10350, 10362, 10380]), array([    2,     7,    19, ..., 10376, 10382, 10385]), array([    1,     3,     8, ..., 10365, 10381, 10383])), _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([   13,    27,    37, ..., 10364, 10379, 10386]), array([    6,    10,    16, ..., 10373, 10374, 10389]), array([    5,     9,    42, ..., 10363, 10369, 10370]), array([   43,    49,    50, ..., 10366, 10387, 10388]), array([    0,    18,    31, ..., 10361, 10372, 10375]), array([   11,    14,    20, ..., 10345, 10368, 10378]), array([    4,    12,    35, ..., 10377, 10384, 10390]), array([   17,    24,    52, ..., 10350, 10362, 10380]), array([    2,     7,    19, ..., 10376, 10382, 10385]), array([    1,     3,     8, ..., 10365, 10381, 10383])), _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:

In [4]:

outcome_models = rlearner._nuisance_models["outcome_model"]
outcome_models

outcome_models = rlearner._nuisance_models["outcome_model"] outcome_models

Out[4]:

[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([   13,    27,    37, ..., 10364, 10379, 10386]), array([    6,    10,    16, ..., 10373, 10374, 10389]), array([    5,     9,    42, ..., 10363, 10369, 10370]), array([   43,    49,    50, ..., 10366, 10387, 10388]), array([    0,    18,    31, ..., 10361, 10372, 10375]), array([   11,    14,    20, ..., 10345, 10368, 10378]), array([    4,    12,    35, ..., 10377, 10384, 10390]), array([   17,    24,    52, ..., 10350, 10362, 10380]), array([    2,     7,    19, ..., 10376, 10382, 10385]), array([    1,     3,     8, ..., 10365, 10381, 10383])), _n_classes=None, classes_=None)]

Note that outcome_models is a sequence of models - in this case of length 1.

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 treatment_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.

In [5]:

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_models},
    propensity_model_params={"max_iter": 500},
    n_variants=2,
)

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

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_models}, propensity_model_params={"max_iter": 500}, n_variants=2, ) rlearner_new.fit( X=df[feature_columns], y=df[outcome_column], w=df[treatment_column], )

Out[5]:

<metalearners.rlearner.RLearner at 0x1146d41d0>

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:

In [6]:

from metalearners import DRLearner

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

drlearner = DRLearner(
    nuisance_model_factory=LGBMRegressor,
    treatment_model_factory=LGBMRegressor,
    nuisance_model_params={"verbose": -1},
    treatment_model_params={"verbose": -1},
    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],
)

from metalearners import DRLearner trained_propensity_model = rlearner._nuisance_models["propensity_model"][0] drlearner = DRLearner( nuisance_model_factory=LGBMRegressor, treatment_model_factory=LGBMRegressor, nuisance_model_params={"verbose": -1}, treatment_model_params={"verbose": -1}, 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], )

Out[6]:

<metalearners.drlearner.DRLearner at 0x11471fc10>

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.