Explaining ML Models with Shapley Values

Recently it was in the news that MIT'S NANDA (Networked AI Agents in Decentralized Architecture) initiative published a report finding that only five percent (5%) of AI pilot programs generate significant revenue, with the majority failing to have any meaningful impact. Buried within that report are two factors that lead to success with current technology: working with an expert vendor and automating back-office processes.

Explainabiltiy is the field of understanding and interpreting why a machine learning model makes a particular decision. For many AI/ML applications, a crucial factor to success is gaining control of the internals of a model from a mechanistic standpoint. This could be to enforce a sense of fairness or enforce some other concept on the model. Part of this process is understanding how a model arrived at a particular set of outcomes. Cooperative game theory provides a solution concept called Shapley values that is useful for gaining that understanding.

SHAP is a Python library for explaining the output of machine learning models. It provides sample datasets and can integrate with matplotlib to provide visualized explanations.

First, ensure any needed dependencies are present in the jupyter environment.

%pip install torch tensorflow tf_keras transformers matplotlib shap numpy scipy

In this example, a BERT (Bidirectional Encoder Representations for Transformers) model, which is ideal for text classification is used. A random review from the IMDB dataset provided with the SHAP library is used as a sample input.

import shap
import numpy as np
import scipy as sp
from torch import tensor
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification

model_id = "lxyuan/distilbert-base-multilingual-cased-sentiments-student"
tokenizer = DistilBertTokenizerFast.from_pretrained(model_id)
model = DistilBertForSequenceClassification.from_pretrained(model_id).cuda()

# BERT uses word-piece tokenization which can create additional tokens so
# additional parameters are specified to ensure fixed length tokenization
def f(x):
    tv = tensor([tokenizer.encode(v, padding="max_length", max_length=500, truncation=True) for v in x]).cuda()
    outputs = model(tv)[0].detach().cpu().numpy()
    scores = (np.exp(outputs).T / np.exp(outputs).sum(-1)).T
    val = sp.special.logit(scores[:, 1])
    return val

explainer = shap.Explainer(f, tokenizer)

imdb_train = shap.datasets.imdb()[0]
shap_values = explainer(imdb_train[:10], fixed_context=1)

shap.plots.text(shap_values[2], display=True)
shap.plots.waterfall(shap_values[2])

Using SHAP's text plot we're able to see how tokens overlay on top of the text along with the importance of those tokens. Red regions increase the output of the model while blue regions decrease the output giving the overall sentiment.

A Waterfall plot also helps to visualize what tokens had the most influence on the model. In this sample we can see how the BERT model broke apart the word "dissapointed" and the tokens comprising this word had an outsized impact on the output.

Another simple way to glance at the internals of a model in a superficial way is to have the model complete a sentence, and then look at the other possible outcomes the model considered. In this example we're asking GPT-2 to complete a common english idiom.

import shap
from transformers import AutoModelForCausalLM, AutoTokenizer

text = "costs an arm and a"

model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
shap_model = shap.models.TopKLM(model, tokenizer, k=5)
masker = shap.maskers.Text(tokenizer)
explainer = shap.Explainer(shap_model, masker)
shap_values = explainer([text])

shap.plots.text(shap_values, display=True)