diff --git a/docs/api/metrics/pyhealth.metrics.interpretability.rst b/docs/api/metrics/pyhealth.metrics.interpretability.rst index e4e0a275e..25fdd10bc 100644 --- a/docs/api/metrics/pyhealth.metrics.interpretability.rst +++ b/docs/api/metrics/pyhealth.metrics.interpretability.rst @@ -29,6 +29,14 @@ For binary classifiers, a sample filter can mark class-0 predictions as from the class-0 perspective. Each percentage is evaluated independently, so a percentage's score does not depend on the other requested percentages. +Integer inputs, such as the code ids of ``sequence`` and ``nested_sequence`` +fields, are always ablated by setting the removed codes to the padding id 0, +whatever the ``ablation_strategy``; averaging or adding noise to code ids has +no meaning. A sample whose codes would all be removed keeps its first code, so +every sequence has at least one non-padding position. Float inputs use the +chosen strategy (``"zero"``, ``"mean"`` or ``"noise"``). See +``examples/interpretability/comprehensiveness_code_sequences.py``. + Base Class ^^^^^^^^^^ diff --git a/examples/interpretability/comprehensiveness_code_sequences.py b/examples/interpretability/comprehensiveness_code_sequences.py new file mode 100644 index 000000000..45b9970ac --- /dev/null +++ b/examples/interpretability/comprehensiveness_code_sequences.py @@ -0,0 +1,63 @@ +"""Comprehensiveness and sufficiency for a model on code sequences. + +Removal-based metrics ablate the code ids of ``sequence`` inputs by setting +them to padding. This example trains an MLP on a small synthetic dataset and +scores Integrated Gradients attributions; it runs in seconds on a CPU and needs +no downloads. + +Usage: + python examples/interpretability/comprehensiveness_code_sequences.py +""" + +import random + +import torch + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.interpret.methods import IntegratedGradients +from pyhealth.metrics.interpretability import ( + evaluate_attribution, + threshold_sample_filter, +) +from pyhealth.models import MLP +from pyhealth.trainer import Trainer + + +def main(): + random.seed(0) + torch.manual_seed(0) + # The label is 1 when code "risk" is present, so a faithful explanation + # should point at it. + samples = [] + for i in range(400): + codes = random.sample([f"c{k}" for k in range(30)], 6) + label = i % 2 + if label: + codes[random.randrange(6)] = "risk" + samples.append({"patient_id": f"p{i}", "codes": codes, "label": label}) + dataset = create_sample_dataset( + samples=samples, + input_schema={"codes": "sequence"}, + output_schema={"label": "binary"}, + dataset_name="synthetic_codes", + ) + loader = get_dataloader(dataset, batch_size=64, shuffle=True) + + model = MLP(dataset=dataset) + Trainer(model=model, enable_logging=False).train( + train_dataloader=loader, epochs=5 + ) + model.eval() + + scores = evaluate_attribution( + model, + get_dataloader(dataset, batch_size=64), + IntegratedGradients(model, steps=16), + percentages=[10, 20, 50], + sample_filter=threshold_sample_filter(0.5), + ) + print(scores) + + +if __name__ == "__main__": + main() diff --git a/pyhealth/metrics/interpretability/base.py b/pyhealth/metrics/interpretability/base.py index ea0ea66c0..23634a408 100644 --- a/pyhealth/metrics/interpretability/base.py +++ b/pyhealth/metrics/interpretability/base.py @@ -266,37 +266,9 @@ def _apply_ablation( mask = masks[key] - # Check if values are integers (discrete features) - is_discrete = x_values.dtype in [ - torch.long, - torch.int, - torch.int32, - torch.int64, - ] - # Apply ablation to values part - if is_discrete: - # For discrete features (codes), multiply by (1-mask) - # Where mask=1 (ablate): set to 0 (padding index) - # Where mask=0 (keep): preserve original value - ablated_values = x_values * (1 - mask).long() - - # Safety: prevent complete ablation of sequences - # Complete ablation (all zeros) causes issues in - # StageNet: - # - All embeddings become zero (padding_idx=0) - # - Mask becomes all zeros - # - get_last_visit() tries to index with -1 - # Solution: keep at least one non-zero element - for b in range(ablated_values.shape[0]): - if ablated_values[b].sum() == 0: - non_zero_mask = x_values[b] != 0 - if non_zero_mask.any(): - # Keep first non-zero element - first_idx = non_zero_mask.nonzero()[0] - ablated_values[b][tuple(first_idx)] = x_values[b][ - tuple(first_idx) - ] + if _is_discrete(x_values): + ablated_values = _ablate_codes(x_values, mask) else: # For continuous features, apply standard ablation if self.ablation_strategy == "zero": @@ -332,8 +304,13 @@ def _apply_ablation( mask = masks[key] + # Code ids can only be removed (set to padding), whatever the + # strategy; averaging or adding noise to ids is meaningless. + if _is_discrete(x): + ablated_inputs[key] = _ablate_codes(x, mask) + # Apply ablation strategy - if self.ablation_strategy == "zero": + elif self.ablation_strategy == "zero": # Set ablated features to 0 ablated_inputs[key] = x * (1 - mask) @@ -527,3 +504,39 @@ def compute( # Average across percentages metric_scores = metric_scores / len(self.percentages) # type: ignore return metric_scores, val_mask + + +def _is_discrete(x: torch.Tensor) -> bool: + """Whether ``x`` holds discrete values such as code ids.""" + return not (x.is_floating_point() or x.is_complex() or x.dtype == torch.bool) + + +def _ablate_codes(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Sets the masked code ids to padding (0), keeping the dtype. + + A sample whose codes would all be removed keeps its first non-padding + code, because models such as StageNet and RNN need at least one + non-padding position per sequence. + + Args: + x: Integer code ids, shape ``(batch, ...)``. + mask: Same shape as ``x``; 1 marks the codes to remove. + + Returns: + A tensor like ``x`` with the masked codes set to 0. + + Examples: + >>> import torch + >>> _ablate_codes(torch.tensor([[3, 4, 5]]), torch.tensor([[1.0, 0.0, 1.0]])) + tensor([[0, 4, 0]]) + >>> _ablate_codes(torch.tensor([[3, 4, 0]]), torch.ones(1, 3)) + tensor([[3, 0, 0]]) + """ + ablated = torch.where(mask.bool(), torch.zeros_like(x), x) + for b in range(ablated.shape[0]): + if not ablated[b].any(): + nonzero = (x[b] != 0).nonzero() + if len(nonzero): + first = tuple(nonzero[0]) + ablated[b][first] = x[b][first] + return ablated diff --git a/tests/core/test_interp_ablation_codes.py b/tests/core/test_interp_ablation_codes.py new file mode 100644 index 000000000..d7e4cb902 --- /dev/null +++ b/tests/core/test_interp_ablation_codes.py @@ -0,0 +1,89 @@ +"""Removal-based metrics on code-sequence inputs (integer code ids).""" + +import math +import unittest + +import torch + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.interpret.methods import IntegratedGradients +from pyhealth.metrics.interpretability import ( + ComprehensivenessMetric, + evaluate_attribution, + threshold_sample_filter, +) +from pyhealth.metrics.interpretability.base import _ablate_codes +from pyhealth.models import MLP, Transformer + + +def _dataset(): + samples = [ + { + "patient_id": f"p{i}", + "codes": [f"c{(i + j) % 7}" for j in range(3 + i % 3)], + "label": i % 2, + } + for i in range(20) + ] + return create_sample_dataset( + samples=samples, + input_schema={"codes": "sequence"}, + output_schema={"label": "binary"}, + dataset_name="ablation_codes", + ) + + +class TestAblateCodes(unittest.TestCase): + def test_masked_codes_become_padding_and_keep_dtype(self): + x = torch.tensor([[3, 4, 5, 0]]) + out = _ablate_codes(x, torch.tensor([[1.0, 0.0, 1.0, 0.0]])) + self.assertEqual(out.dtype, torch.long) + self.assertEqual(out.tolist(), [[0, 4, 0, 0]]) + + def test_fully_ablated_sample_keeps_first_code(self): + out = _ablate_codes(torch.tensor([[0, 6, 7]]), torch.ones(1, 3)) + self.assertEqual(out.tolist(), [[0, 6, 0]]) + + def test_float_inputs_still_use_the_strategy(self): + model = MLP(dataset=_dataset()) + metric = ComprehensivenessMetric( + model, ablation_strategy="zero", sample_filter=threshold_sample_filter() + ) + x = torch.tensor([[1.5, 2.5]]) + out = metric._apply_ablation({"x": x}, {"x": torch.tensor([[1.0, 0.0]])}) + self.assertTrue(torch.equal(out["x"], torch.tensor([[0.0, 2.5]]))) + + +class TestRemovalMetricsOnCodeSequences(unittest.TestCase): + """Comprehensiveness/sufficiency used to crash: code ids became floats.""" + + @classmethod + def setUpClass(cls): + torch.manual_seed(0) + cls.dataset = _dataset() + + def _check(self, model_class): + model = model_class(dataset=self.dataset) + model.eval() + loader = get_dataloader(self.dataset, batch_size=8) + for strategy in ["zero", "mean", "noise"]: + with self.subTest(model=model_class.__name__, strategy=strategy): + scores = evaluate_attribution( + model, + loader, + IntegratedGradients(model, steps=5), + ablation_strategy=strategy, + sample_filter=threshold_sample_filter(0.0), + ) + for name, value in scores.items(): + self.assertTrue(math.isfinite(value), (name, value)) + + def test_mlp(self): + self._check(MLP) + + def test_transformer(self): + self._check(Transformer) + + +if __name__ == "__main__": + unittest.main()