Different distributions for the training and test sets in Section 4.5

[airalcorn2]

Hi. Thanks a lot for this book. I'm really enjoying it.

I was, however, pretty surprised when I came to Figure 4 in Section 4.5 ("Outcome model-based methods"). While I understand the reasoning for why a supervised learning model will produce biased estimates of the treatment effect under the conditions described in the book, I was struggling to understand why the estimates were so far off. I decided to run my own experiments using a neural network (code here), and the first thing I noticed was that the test loss was extremely divergent from the training loss—even when training on a large number of samples. When I looked into the code for how the datasets were generated, I noticed that generate_dataset is called separately to generate the training and testing datasets, which means the confounding variables are generated from different distributions, and so the training and testing data are not sampled from the same distribution. If I instead use the same distribution to generate the confounding variables for the training and testing data, the training loss and test loss are similar.

Next, I decided to use the test set to estimate the treatment effect because the test loss is supposed to be our ultimate measure of how well the model is approximating the target function. Using this strategy, I did find that the neural network's estimates of the treatment effect were biased, but they were not nearly as biased as the estimates shown in Figure 4. For the 1,000 sample dataset, the average estimate of the treatment effect over 20 runs was 7.5 for the neural network and 11.3 for a linear model trained on the same datasets. However, when I increased the number of samples to 10,000, the average estimate of the treatment effect was 9.6 for the neural network and 11.0 for a linear model. Further, the neural network's estimate was closer to the true effect than the linear model in 90% of the 20 runs when using the 10,000 sample dataset, while it was only closer than the linear model in 30% of the runs for the 1,000 sample dataset. Lastly, interestingly, the accuracy of the neural network's estimate is strongly correlated with its test loss. The R^2 for a logarithm trendline predicting the absolute error for the effect estimate from the test loss for the neural network is 0.83 (all experiment results can be found in the spreadsheet here).

These results fit my intuition better. If there aren't any omitted variables and the causal direction is correct, a model with sufficient capacity will eventually be an accurate approximation of the true function given enough data. If the model's estimate of the treatment effect is far off, then we would expect its test loss to reflect that because the spurious relationships the model learned to compensate for the incorrect estimate of the treatment effect would lead to considerably wrong predictions. All in all, I think these results are actually pretty encouraging. Another interesting experiment would be to impose a linearity assumption on the treatment variable by only using it as an input in the final layer of the neural network.

[airalcorn2]

Almost forgot about this. Any thoughts, @amit-sharma?