Log In to continue learning
Tabular Foundation Models for Data Science Agents
Intermediate
1h32m
You can now make TabPFN predictions efficiently, but you may also want to know what is driving them. In this lesson, you'll use SHAP values and partial dependence plots to examine feature influence and explain individual predictions. Let's get to it. Interpretability helps us inspect how the model arrived at the output and what patterns its predictions reflect. We will use SHAP values to explain a single prediction and understand global feature importance. We will also build partial dependence plots to analyze the relationships between modelled features and predicted probabilities. SHAP values are estimated contributions of the feature value to a model prediction. We start by explaining the predicted probability of disease of a single patient with a SHAP waterfall plot. The waterfall starts at a reference prediction of about 0.45 and ends well above 0.9. Red contributions, positive SHAP values, move the model output upward, Blue contributions, negative SHAP values, move it downward. For this patient, three Blocked vessels contribute about 30 percentage points upward, chest pain and Max heart rate also push the predicted probability of disease higher, while the normal Thallium scan contributes downward. All the contributions together bridge the gap from the reference to the final prediction. This view can help investigate unexpected outputs and support decision making. But which features does the model actually use? To understand which features the model uses in general, we can visualize all SHAP values for all features together. This is a SHAP beeswarm plot. Each dot represents one patient's attribution for one feature. A dot to the right indicates a contribution towards a higher predicted disease probability. A dot on the left indicates a contribution in the other direction. Features are ordered by their average absolute contribution across the records shown. Blocked vessels is at the top. Its higher values, shown in red, push the predicted probabilities up, while lower values in blue mostly appear on the negative side. Color is the feature value, not the prediction. And this, overall, is a summary over all the displayed samples. An individual patient can have a different leading contributor. Now that we know which features matter most for the model, let's understand their impact on predictions. To better understand how the model responds to changes in input features, We can use Partial Dependence Plots. To build a Partial Dependence Plot, we keep all feature values constant and vary the value of a single selected feature. On the top-left chart, we can see that the predicted probability of disease increases as the number of blocked vessels increases, too. On the other hand, the predicted probability of disease decreases as the max heart rate increases, as a sign of good cardiovascular health. Let's now apply all these interpretability methods to our data. Let's start by importing the libraries we'll use in this lesson, along with the helper functions. Let's now load the heart disease dataset. and split it between train and test. Here, we show the top rows of the training set, along with some features. We can also inspect the top three rows of the target variable. We then instantiate the TabPFNClassifier object and call fit with the train features and the train target variables. We then use the predict_proba method to generate probabilities over the test set. We score these probabilities with the ROC AUC, which you see on the screen. This ROC AUC, above .9, is an indication that TabPFN is a good predictor for this use case. We now come to the most important part of this notebook. We instantiate an explainer object from the TabPFN classifier and the training data. We then use this explainer object to generate SHAP values for 10 rows of the test set. This may take a while to complete. Now that we have generated these SHAP values, we can use them to explain individual predictions. We will explain two predictions: First, a high predicted probability of disease, and then a predicted probability of disease close to 0.5. We can then plot the SHAP values for a given patient, here, the high-risk patient, that will show us which features contribute to the final prediction. On this chart, we can see that the high predicted probability is driven by four main features: the Thallium scan, the Exercise angina, Chest pain type, and ST depression. On the other hand, this patient has zero blocked vessels, which brings the predicted probability down. This is the only feature with a negative contribution here. We can then move to a patient with the predicted probability closer to 0.5. As you can see on the chart, the prediction is mostly driven by the Thallium scan and the Chest pain type. A lot of the other feature contributions cancel out, leading to a predicted probability of around 0.7. Beyond individual predictions, we can plot the SHAP values of all rows and all features on a single image. This is the beeswarm plot. There, the features are ranked by importance and by how much they drive the prediction. Looking at the Max heart rate, the color of the dot represents the feature value. Here, a red dot means that the Max heart rate is high. And as you can see, a high Max heart rate is a sign of good health and pushes the predicted probability of disease down. On the other hand, a low Max heart rate pushes the predicted probability of disease higher. To better understand how the model responds to changes in the value of individual features, we can build partial dependence plots. I can now run this cell. And it would show the partial dependence plots for four different features: the Blocked vessels, Max heart rate, Cholesterol, and ST depression. As you can see, the predicted probability of disease increases as the number of blocked vessels increases, too. On the other hand, the predicted probability of disease decreases along with the max heart rate. A high max heart rate is a sign of good cardiovascular health. In this notebook, we explored three methods to interpret TabPFN predictions. We first explain individual predictions with SHAP values and a waterfall plot. We then plotted all of these SHAP values on a beeswarm plot to get a better idea of feature importance. We also explored partial dependence plots that show us how the model responds to changes in individual feature values. In the next lesson, you will learn how to accelerate TabPFN inference with the KV cache. All right, see you there.
