Causal Framework for Model Robustness
By: Michael Larionov
Tags: Causality, Feature Engineering, Feature Selection, Causal Inference, Machine Learning

Model Robustness

We would all like our Machine Learning models generalize on the unseen data, but often find that our model performance drops when the new data do not look like the old data, that is have a different distribution. For example, a medical diagnostics system we trained on the data from one country does not generalize to another other country because the common diseases in that country are different than diseases in the country of the training set. Machine learning researches and practitioners developed various techniques to overcome this problem and name the model robust.

A Robust model is, by definition, not sensitive to distribution shift of the predictor or target variables. Model robustness is an active area of research and one of ways to analyze model robustness is using causal inference.

Causal Model of Robustness

In general, causal inference refers to answering question whether changing one variable affects the result. Examples could be studying how promotions and discounts improve sales or how a particular medical treatment improves patient’s outcome [1].

We all know a maxim: “Correlation does not imply causation”. Studying correlations is a very important part of the work of statisticians and data scientists, but are not sufficient when studying causal questions. In essence, causal model looks very similar to Bayesian network in that it expresses the joined distribution of all variables using conditional probabilities, for example:

This model becomes causal if we condition all probabilities on its direct causes:

An intervention on a variable Xi is when we change the distribution of Xi, for example, setting the variable to a constant. Causal modeling has extended syntax of conditional probabilities to include interventions. For example, the following equation

means the probability of the variable R when we set the variable G to T. This is also known as a causal mechanism.

How does causal modeling helps us build robust models? If we are predicting Y= f(X₁, X₂, … Xn), then the change in the distribution of a variable X₁ can be modeled as an intervention [2]. Thus, robust models are the models that are invariant to all possible interventions.

Causal models have a very useful property called independence between cause and mechanism, which roughly states that the mechanism mapping cause C to the effect E does not depend on the distribution of the cause variable C [3]. So if our predictor variables are causes and the target variable is the effect, then by this property the model P(Y|X) is robust, because it will work equally well when the marginal distribution P(X) changes.

However, if we are trying to predict causes based on effects, the situation becomes different. Consider the medical diagnostics case, when we have a set of symptoms and we are trying to predict the diagnosis. Since the diagnosis is the cause and the symptoms are the effects, we can use the independence assumption and claim that the we can robustly predict symptoms based on the disease E(X|Y). If we are trying to predict the diagnosis based on symptoms we can use Bayes’ rule:

We can see from this equation that the predicted diagnosis depends not only on the causal mechanism P(X|Y), but also on the prevalence of the disease in a given area P(Y). For example, the same symptoms can be used to diagnose malaria in one country and respiratory conditions in another country where malaria is not prevalent. The situation when we are trying to predict a cause from effects is called anticausal prediction [4].

Invariant sets

In the majority of cases the relationship between X and Y cannot be reduced to cause-effect. For example, in medical diagnostics the problem is inverted, because we need to predict cause based on effects. Rojas Carulla et al [3] asked a question, whether is is possible to derive a subset Xi of features X for which the distribution P(Y|Xi) is independent of P(Xi). This subset is called an invariant set, because the model prediction of Y given Xi does not depend on any intervention on any variables except the target variable. Let’s look at the following causal diagram:

Here we are trying to predict Y based on predictor variables X₁ through X₁₆. We can eliminate all indirect causes because they are redundant, and any variable that is not a cause of Y can change its distribution rendering predictions inaccurate. This leaves us with direct causes as the invariant set for this problem.

This means that the model is robust if Xi consist of all direct causes of Y .

You can see that we drastically eliminated the majority of the features, so while the model remains robust it has low predictive power. For completely anticausal models, like above mentioned medical diagnostics problem, the invariant set is empty.

Robustness of anticausal models

There are ways to improve robustness of the anticausal models without eliminating all predictor variables. One of he obvious one is to explicitly feed to the model the information about the distribution of Y as a new variable. In our example, we can add the prevalence of the disease in a given country as an additional predictor variable. If our training data contains several countries the model can learn to use this information to adjust the prediction for a new country. For example, if we are using a logistic regression, we can represent the updated model as:

This approach is equivalent to target encoding, a popular method of encoding categorical variables used very extensively in data science competitions. We can feed more information about the marginal distribution of the target variable to the model rather than its expectation. One possible technique is to use Bayesian Target Encoding I covered in an earlier article.

Conclusion

Causal inference became a necessary tool in the toolbox of every data scientist. While this is still an area of active research, there are techniques and tools to solve causal problems for the most common problems. However the role of the causal inference is not restricted to causal problems but is also extended to machine learning via causal theory of robustness. In this article, I attempted to share my thoughts on the causal considerations of the robust machine learning models and how it connects to common techniques of feature engineering.

References

[1] Judea Pearl. 2009. Causality: Models, Reasoning and Inference (2nd. ed.). Cambridge University Press, USA.

[2] R. Christiansen, N. Pfister, M. E. Jakobsen, N. Gnecco and J. Peters, “A causal framework for distribution generalization,” in IEEE Transactions on Pattern Analysis and Machine Intelligence, doi: 10.1109/TPAMI.2021.3094760.

[3] M. Rojas-Carulla, B. Schölkopf, R. Turner, and J. Peters. Causal transfer in machine learning. Journal of Machine Learning Research, 19(36):1–34, 2018

[4] B. Schölkopf, D. Janzing, J. Peters, E. Sgouritsa, K. Zhang, and J. M. Mooij. On causal and anticausal learning. In Proceedings of the 29th International Conference on Machine Learning (ICML), pages 1255–1262. Omnipress, 2012.