GradientExplainer misattributes missing values (TabICLv2 over-attributes imputed cells, KumoTabular gives them zero gradient)

[skochrekar-hq]

Continuing from discussion here with @rusty1s : https://nvidia.slack.com/archives/C0BH1KUV1EC/p1790060033563569?thread_ts=1789170336.698109&cid=C0BH1KUV1EC

Quick intro - We’re identifying which telemetry fields distinguish problematic GeForce NOW gaming sessions from healthy ones, so engineers know where to investigate. We are using TFMs to distinguish these sessions across thousands of fields, then measure how replacing individual values with healthy-session references changes its predictions. This produces a prioritized list of diagnostic signals that engineers can investigate and validate as potential root causes.

Summary

sdm.explain.GradientExplainer differentiates the model output with respect to the preprocessed numerical block only. For columns with missing values this produces attributions that are wrong in opposite directions for the two tabular models, and for heavy-tailed columns it ranks uninformative fields above informative ones in both. Permutation importance on the same data ranks everything correctly. Reproduced on a fully synthetic table (script below), no real data involved.

Environment

• structured-data-models snapshot of main from 2026-09-14 (0.1.0.dev0), Python 3.12.3, torch 2.14.0+cpu, CPU only
• TabICLv2 classifier checkpoint jingang/TabICL tabicl-classifier-v2-20260212.ckpt; KumoTabular small, nvidia/Kumo-Tabular v1.0.3
• 3,000 synthetic rows, 1,500 context / 1,500 query; columns: x_signal (informative), cat (informative level), x_missing_when_pos (informative only through missingness: 70% missing on positives), tail_lognormal and tail_student_t (heavy-tailed, uninformative), noise_1..3

Observed

	TabICLv2 (AUROC 0.930)	KumoTabular small (AUROC 0.931)
mean |grad| at missing cells of x_missing_when_pos vs observed cells	350 vs 2.0	0.000 vs 0.027
gradient rank of x_missing_when_pos (8 columns)	1st, mean |grad| 130	8th (last), 0.017
permutation rank of the same column	1st, AUROC drop 0.30	1st, AUROC drop 0.29
gradient rank of uninformative tail_lognormal	3rd, above x_signal (4th)	1st, above cat and x_signal
permutation rank of tail_lognormal	4th, drop 0.0002	5th, drop 0.0000

Why it happens

1. TabICLv2: the default recipe applies ImputeMean (sdm/models/tabiclv2/recipe.py), so imputed cells reach the model as ordinary values. The gradient there is the sensitivity to the fill value, and it is ~175x larger than at observed cells, so the column's score is dominated by cells that carry no data.
2. KumoTabular: nothing is imputed in the recipe, but CellEmbedding.forward (sdm/models/kumo/tabular/cell_embedding.py) does x = torch.where(missing, impute_mean, x) and routes the NaN mask through nan_lin. torch.where sends the gradient to impute_mean, not to x, so d out / d x is exactly zero at missing cells. The missingness signal flows only through nan_lin, which GradientExplainer (sdm/explain/gradient.py, _GradientCallback._capture) never differentiates. Result: a column whose entire signal is its missingness becomes invisible to the explainer.
3. Heavy tails: after Standardize + Clip the slope with respect to a heavy-tailed column is large regardless of whether the column matters, so tail_lognormal outranks real signal in both models. This is a property of raw gradients in standardized space, not of either model.

Expected

Attribution that ranks x_missing_when_pos first (as permutation does) in both models, and ranks the uninformative heavy-tailed columns near zero.

Suggestions

• For TabICLv2-style recipes, return an imputed-cell mask with the attribution and report the "was missing" contribution separately from the value contribution (or zero it).
• For KumoTabular, also differentiate through the mask pathway (capture the mask tensor in _GradientCallback and report d out / d mask per column alongside d out / d x).
• Report categoricals per level (presence), not as a slope on an ordinal code.
• Scale per-column gradients by a robust spread of the preprocessed column, or offer Integrated Gradients with the imputation value as baseline.
• Ship permutation importance (column-family aware), leave-one-context-example-out and what-if re-prediction as first-class tools; they need no differentiability and are cheap for in-context models. In this repro they are correct in every case.

Also worth a docs line

KumoTabular.default_recipe() applies SelectColumns(500, method="first") immediately after ShuffleColumns(method="latin"), so each estimator sees a different 500-column subset of a wide table. Reasonable design, but silent for users with 1,000+ columns.

Reproduction script (synthetic data only, ~2 minutes on CPU)

"""Self-contained reproduction of gradient-attribution artefacts in sdm.

Synthetic table, no real data. Compares the gradient explainer with
permutation importance for TabICLv2 and KumoTabular, and measures the
gradient at missing cells versus observed cells.

Columns: x_signal (informative), cat (informative level "C"),
x_missing_when_pos (informative through missingness only), tail_lognormal
and tail_student_t (heavy-tailed, NOT informative), noise_1..3 (not
informative).

    $PY repro_gradient_artifacts.py
"""

import numpy as np
import pandas as pd
import torch
from sklearn.metrics import roc_auc_score

import sdm
from sdm import CategoricalTensor, Stype, TableTensor
from sdm.explain import GradientExplainer

torch.manual_seed(0)
rng = np.random.default_rng(0)
n, n_ctx = 3000, 1500
x_signal = rng.normal(size=n)
cat = rng.choice(list("ABCDEF"), n)
logit = 1.2 * x_signal + 1.5 * (cat == "C")
y = (rng.random(n) < 1 / (1 + np.exp(-logit))).astype(int)
x_missing = rng.normal(size=n)
x_missing[(y == 1) & (rng.random(n) < 0.7)] = np.nan  # missing mostly on positives
df = pd.DataFrame({
    "x_signal": x_signal,
    "x_missing_when_pos": x_missing,
    "tail_lognormal": rng.lognormal(0, 2.0, n),      # heavy tail, uninformative
    "tail_student_t": rng.standard_t(1.5, n),        # heavy tail, uninformative
    "noise_1": rng.normal(size=n), "noise_2": rng.normal(size=n), "noise_3": rng.normal(size=n),
    "cat": cat, "y": y,
})
num_cols = [c for c in df.columns if c not in ("cat", "y")]
table = TableTensor.from_pandas(df=df, stypes={**{c: "numerical" for c in num_cols}, "cat": "categorical", "y": "categorical"})
ctx, qry = table[:n_ctx], table[n_ctx:]
ctx_x, ctx_y, qry_x, y_qry = ctx.drop_columns("y"), ctx["y"], qry.drop_columns("y"), y[n_ctx:]
missing_mask = np.isnan(df["x_missing_when_pos"].to_numpy()[n_ctx:])

for name, model in (("TabICLv2", sdm.models.TabICLv2(task="classification", device="cpu")),
                    ("KumoTabular-small", sdm.models.KumoTabular(task="classification", size="small", device="cpu"))):
    torch.manual_seed(0)
    model.fit(x=ctx_x, y=ctx_y)

    def p1(x):
        with torch.inference_mode():
            out = model.predict(x)
        return out.numerical[:, list(out.columns[Stype.numerical]).index("1")].numpy()

    base = roc_auc_score(y_qry, p1(qry_x))
    # gradient explainer (float32) on all 1,500 query rows
    out0 = model.predict(qry_x)
    col1 = list(out0.columns[Stype.numerical]).index("1")
    res = GradientExplainer(output=lambda o: o.numerical[..., col1]).explain(model, qry_x)
    g = res.x.numerical.detach()
    gnames = list(res.x.columns[Stype.numerical])
    grad = pd.Series(g.abs().mean(0).numpy(), index=gnames)
    # permutation importance (AUROC drop, 3 seeds) on raw columns
    perm = {}
    for c in num_cols + ["cat"]:
        drops = []
        for s in range(3):
            gen = torch.Generator().manual_seed(100 + s)
            idx = torch.randperm(len(qry_x), generator=gen)
            st = qry_x.stype(c)
            j = list(qry_x.columns[st]).index(c)
            if st == Stype.numerical:
                blk = qry_x.numerical.clone(); blk[:, j] = blk[idx, j]; xs = qry_x.replace_blocks(numerical=blk)
            else:
                codes = qry_x.categorical.code.clone(); codes[:, j] = codes[idx, j]
                xs = qry_x.replace_blocks(categorical=CategoricalTensor(codes, categories=qry_x.categorical.categories))
            drops.append(base - roc_auc_score(y_qry, p1(xs)))
        perm[c] = float(np.mean(drops))
    perm = pd.Series(perm)
    tab = pd.DataFrame({"mean_abs_grad": grad, "grad_rank": grad.rank(ascending=False).astype(int)})
    tab["perm_auroc_drop"] = perm.reindex(tab.index)
    tab["perm_rank"] = tab.perm_auroc_drop.rank(ascending=False).astype(int)
    tab["truth"] = ["informative" if c in ("x_signal", "cat") else "informative via missingness" if c.startswith("x_missing") else "uninformative" for c in tab.index]
    print(f"\n=== {name}: AUROC {base:.3f} | gradient explainer vs permutation importance ===")
    print(tab.sort_values("grad_rank").to_string(float_format=lambda v: f"{v:.4f}"))
    # gradient at missing vs observed cells of the missing column
    j = gnames.index("x_missing_when_pos")
    gm = g[:, j].abs().numpy()
    print(f"gradient on x_missing_when_pos: mean |grad| at MISSING cells {gm[missing_mask].mean():.4f} vs OBSERVED cells {gm[~missing_mask].mean():.4f} "
          f"({missing_mask.mean():.0%} of query cells missing)")
    model.clear()