+1 (415) 366-1364

Explaining TensorFlow Model Predictions: Integrated Gradients, SHAP, and Model Cards

Every model we help a client put into production now arrives at the same question, usually from legal or risk rather than engineering: why did it decide that, and can you show us? Under the EU AI Act's phased obligations, and under the model-risk rules banks and insurers have lived with for years, "the network learned it" is not an answer you can file. This tutorial is the explainability workflow we set up on TensorFlow engagements: Integrated Gradients for deep models, SHAP-style attributions for tabular ones, built-in interpretability where the model type gives it to you for free, and a model card that turns all of it into a document a reviewer can sign.

None of our other tutorials cover this. Serving TensorFlow Models in Production gets predictions out the door; Catching Data Drift Early with TensorFlow Data Validation watches the inputs. Explainability is the third leg: it answers questions about a single prediction and about the model's global behaviour.

What "explainability" actually means in a review

Separate two requests before you write any code, because they need different tools:

  • Local explanation - "why was this loan application scored 0.12?" You need per-feature or per-pixel attributions for one input.
  • Global explanation - "what does the model rely on overall, and does it rely on anything we're not allowed to use?" You need aggregate attributions across a dataset, plus slice-level performance.

Reviewers almost always ask for both, and they will also ask the question that sinks unprepared teams: is the explanation stable? An attribution method that gives a different story on two near-identical inputs is worse than no attribution at all.

Setup

pip install -U tensorflow keras matplotlib
# optional, for the tabular section
pip install -U shap tensorflow_decision_forests

Everything below assumes a Keras 3 model you can call on a batch and differentiate through - the same model you exported in the serving tutorial.

1. Integrated Gradients from scratch

Integrated Gradients (IG) is the workhorse for differentiable models: it attributes a prediction to the input by integrating gradients along a straight path from a baseline input to the real one. It satisfies a completeness property - attributions sum to the difference in model output between input and baseline - which is exactly the kind of arithmetic a reviewer likes.

There is no need for a heavy library; it is about twenty lines of TensorFlow.

import tensorflow as tf

@tf.function
def _grads(model, inputs, target_idx):
    with tf.GradientTape() as tape:
        tape.watch(inputs)
        preds = model(inputs, training=False)
        if preds.shape[-1] > 1:
            preds = tf.gather(preds, target_idx, axis=-1, batch_dims=0)
        target = tf.reduce_sum(preds)
    return tape.gradient(target, inputs)


def integrated_gradients(model, x, baseline=None, target_idx=0, steps=128, batch=16):
    """x: single input tensor (no batch dim). Returns attributions with x's shape."""
    x = tf.cast(x, tf.float32)
    baseline = tf.zeros_like(x) if baseline is None else tf.cast(baseline, tf.float32)

    alphas = tf.linspace(0.0, 1.0, steps + 1)
    total = tf.zeros_like(x)

    for i in range(0, steps + 1, batch):
        a = alphas[i:i + batch]
        a = tf.reshape(a, [-1] + [1] * len(x.shape))
        interpolated = baseline + a * (x - baseline)
        g = _grads(model, interpolated, target_idx)
        total += tf.reduce_sum(g, axis=0)

    avg_grads = total / tf.cast(steps + 1, tf.float32)
    return (x - baseline) * avg_grads

Choosing the baseline (this is the whole game)

The baseline encodes the counterfactual "compared to what?" and it changes the answer:

  • Images. A black image is the default and is usually fine, but it makes genuinely black pixels unattributable. If that matters, average attributions over several baselines (black, white, blurred input, uniform noise) - this is expected gradients.
  • Tabular. Never use zeros: a zero income is not a neutral customer. Use the training-set median for numerics and the modal category for categoricals, or, better, average over a random sample of the training set.
  • Text. The padding or [MASK] token embedding, attributing at the embedding level and summing over the embedding dimension to get a per-token score.

Write the baseline choice down in the model card. It is the first thing a competent reviewer will challenge.

Verify completeness before you trust a single plot

attr = integrated_gradients(model, x, baseline=baseline, target_idx=cls, steps=256)

lhs = tf.reduce_sum(attr).numpy()
rhs = (model(x[None], training=False)[0, cls]
       - model(baseline[None], training=False)[0, cls]).numpy()
print(f"sum(attributions)={lhs:.4f}  f(x)-f(baseline)={rhs:.4f}  gap={abs(lhs-rhs):.4f}")

If the gap is more than a couple of percent of rhs, raise steps (256 or 512) before you conclude anything. A convergence check like this belongs in your test suite, not in a notebook you ran once.

Visualising image attributions

import numpy as np, matplotlib.pyplot as plt

heat = tf.reduce_sum(tf.abs(attr), axis=-1).numpy()
heat = heat / (np.percentile(heat, 99) + 1e-9)   # clip outliers, don't let one pixel own the scale

fig, ax = plt.subplots(1, 2, figsize=(9, 4))
ax[0].imshow(x.numpy()); ax[0].set_title("input"); ax[0].axis("off")
ax[1].imshow(x.numpy()); ax[1].imshow(np.clip(heat, 0, 1), cmap="inferno", alpha=0.6)
ax[1].set_title("integrated gradients"); ax[1].axis("off")
plt.tight_layout()

Two habits that prevent embarrassing meetings: percentile-clip the heatmap (one saturated pixel makes every map look identical), and always render the attribution for the runner-up class too. If the model highlights the same region for "melanoma" and "benign mole", the region is not the explanation.

2. Global explanations: aggregate, then slice

Run IG over a stratified sample - a few hundred inputs is plenty - and aggregate:

import collections

totals = collections.Counter()
for xi, yi in sample_ds:                      # unbatched
    a = integrated_gradients(model, xi, baseline=baseline, target_idx=int(yi), steps=128)
    for name, v in zip(feature_names, tf.reshape(a, [-1]).numpy()):
        totals[name] += abs(float(v))

for name, v in sorted(totals.items(), key=lambda kv: -kv[1])[:15]:
    print(f"{v/len(feature_names):10.4f}  {name}")

Read the ranking with suspicion, and look for these three failure modes:

  1. A proxy for a protected attribute near the top - postcode, first name, device locale. This is where explainability turns into a compliance finding.
  2. A leaked feature - something only known after the outcome. High attribution on a field like days_since_collections_call in a default model usually means leakage, not insight.
  3. An ID-like feature with real attribution. The model has memorised.

Then split performance by slice, because a model can be accurate overall and useless for a subgroup:

for slice_name, ds in slices.items():
    m = model.evaluate(ds, verbose=0, return_dict=True)
    print(f"{slice_name:24s} n={cardinality[slice_name]:6d} auc={m['auc']:.3f}")

Report slice metrics with sample counts. A slice AUC computed on 40 rows is noise, and presenting it as a finding damages your credibility for the slices that do matter.

3. Tabular models: use the structure you have

For tabular problems, don't reach for gradients first.

TensorFlow Decision Forests gives you variable importances directly from the trained model, at no compute cost:

import tensorflow_decision_forests as tfdf

model = tfdf.keras.GradientBoostedTreesModel()
model.fit(train_ds)

inspector = model.make_inspector()
print(inspector.variable_importances().keys())
for feature, score in inspector.variable_importances()["SUM_SCORE"][:15]:
    print(f"{score:12.2f}  {feature.name}")

NUM_AS_ROOT and SUM_SCORE answer different questions - which features drive the top splits, versus total contribution across the ensemble - so quote the one you mean and name it. As our Tabular ML on TensorFlow tutorial argues, on a lot of business datasets a decision forest matches a neural net and hands you interpretability for free. That is a legitimate architectural reason to choose it.

For a Keras tabular network, shap.GradientExplainer gives you attributions with the additive property reviewers understand, using a background sample rather than a single baseline:

import shap

background = next(iter(train_ds.batch(200)))[0]
explainer = shap.GradientExplainer(model, background)
shap_values = explainer.shap_values(x_batch)     # list per output class
shap.summary_plot(shap_values[1], x_batch, feature_names=feature_names)

Keep the background sample fixed and versioned. Change it and every previously published explanation shifts, which is impossible to defend in an audit.

4. Sanity-check the explanations themselves

Attribution methods can produce confident nonsense. Three cheap tests, all worth automating:

  • Randomised-model test. Re-initialise the weights and re-run IG. If the heatmap looks broadly the same, your method is describing the input, not the model.
  • Stability test. Perturb the input slightly (small noise, or one irrelevant field changed). Attribution rankings should move a little, not reshuffle.
  • Deletion / insertion curve. Mask the top-k attributed features and measure how fast the score falls; compare against masking random features. A method whose curve tracks random is not explaining anything.
def deletion_curve(model, x, attr, cls, ks=(0, 5, 10, 20, 50)):
    order = tf.argsort(tf.reshape(tf.abs(attr), [-1]), direction="DESCENDING")
    flat = tf.reshape(tf.identity(x), [-1]).numpy()
    out = []
    for k in ks:
        masked = flat.copy()
        masked[order[:k].numpy()] = 0.0
        xm = tf.reshape(masked, x.shape)
        out.append(float(model(xm[None], training=False)[0, cls]))
    return list(zip(ks, out))

5. Write the model card, and put it in CI

The deliverable is not a heatmap in a slide deck; it is a document that travels with the model version. We keep a model_card.md next to the exported SavedModel with these sections:

  • Intended use and out-of-scope use. Explicitly name uses the model must not be put to.
  • Training data. Source, date range, size, known gaps, and the preprocessing that was baked into the export.
  • Performance. Headline metrics plus a slice table with sample counts.
  • Explainability method. IG or SHAP, the baseline or background sample, step count, and the measured completeness gap.
  • Known limitations and failure modes, including anything the randomised-model and stability tests surfaced.
  • Version and lineage. Model version, dataset hash, training commit - the same fields your serving layer logs with every prediction.

Then make the pipeline enforce it. In a TFX continuous-training pipeline, the natural home is a custom component after Evaluator, run before the model is blessed:

def validate_explainability(model, sample, baseline, thresholds):
    attr = integrated_gradients(model, sample, baseline=baseline, steps=256)
    gap = abs(float(tf.reduce_sum(attr)) - float(delta_output(model, sample, baseline)))
    assert gap < thresholds["completeness"], f"IG did not converge: gap={gap:.4f}"

    top = top_features(attr, feature_names, k=5)
    banned = set(thresholds["prohibited_features"]) & set(top)
    assert not banned, f"prohibited features in top-5 attributions: {sorted(banned)}"

Two gates, both cheap: the explanation converged, and no prohibited or proxy feature is in the top ranks. A model that fails either does not get blessed. Regenerating the card automatically on every retrain is what turns explainability from a one-off audit exercise into a property of the system.

What this does not do

Attributions are correlational, not causal. IG tells you what this model, with these weights, responded to on this input; it does not tell you what would happen in the world if the feature changed. For causal questions you need a causal design, not a saliency map. And explainability is not fairness: a clean-looking attribution ranking is entirely compatible with a model that performs badly for a subgroup, which is why the slice table sits next to it in the card.

Getting help

We do this as a fixed-scope piece of work: pick the model, agree the explanation questions your reviewers will actually ask, build the attribution and slice pipeline, and hand over a model card plus the CI gates that keep it honest. It is usually two to three weeks, and it is far cheaper before a launch than after a regulator's letter. See our Services page, or get in touch with a short description of the model and who is asking the questions.