Weight distributed attribution aggregation by samples (#1925)

captum/_utils/common.py

@@ -36,6 +36,24 @@
 from torch.nn import Module
 
 
+def _generate_derangement(n: int, device: torch.device | None = None) -> Tensor:
+    """Generate a random cyclic donor ordering without fixed points.
+
+    A random cycle gives every row a uniformly distributed donor among the
+    other rows while requiring only one device-side ``randperm`` and no
+    device-to-host synchronization. The result is a valid derangement for
+    every ``n > 1``.
+    """
+    row_indices = torch.arange(n, device=device)
+    if n <= 1:
+        return row_indices
+
+    cycle = torch.randperm(n, device=device)
+    permutation = torch.empty_like(cycle)
+    permutation[cycle] = torch.roll(cycle, shifts=-1)
+    return permutation
+
+
 def parse_version(v: str) -> Tuple[int, ...]:
     """
     Parse version strings into tuples for comparison.
@@ -117,6 +135,7 @@ def _validate_input(
     inputs: Tuple[Tensor, ...],
     baselines: Tuple[Union[Tensor, int, float], ...],
     draw_baseline_from_distrib: bool = False,
+    allow_broadcastable_baselines: bool = False,
 ) -> None:
     assert len(inputs) == len(baselines), (
         "Input and baseline must have the same "
@@ -137,10 +156,24 @@ def _validate_input(
                 " Found baseline: {} and input: {} ".format(baseline, input)
             )
         else:
+            baseline_is_broadcastable = False
+            if allow_broadcastable_baselines and isinstance(baseline, Tensor):
+                try:
+                    baseline_is_broadcastable = (
+                        torch.broadcast_shapes(input.shape, baseline.shape)
+                        == input.shape
+                    )
+                except RuntimeError:
+                    pass
             assert (
                 isinstance(baseline, (int, float))
                 or input.shape == baseline.shape
-                or baseline.shape[0] == 1
+                or (
+                    baseline.dim() > 0
+                    and baseline.shape[0] == 1
+                    and input.shape[1:] == baseline.shape[1:]
+                )
+                or baseline_is_broadcastable
             ), (
                 "Baseline can be provided as a tensor for just one input and"
                 " broadcasted to the batch or input and baseline must have the"
@@ -226,6 +259,10 @@ def _format_feature_mask(
 
     else:
         formatted_mask = _format_tensor_into_tuples(feature_mask)
+        assert len(formatted_mask) == len(inputs), (
+            "Input and feature mask must have the same number of tensors, "
+            f"but input has {len(inputs)} and feature mask has {len(formatted_mask)}."
+        )
 
     return formatted_mask
 

captum/attr/_core/feature_ablation.py

@@ -22,6 +22,7 @@
     _is_tuple,
     _maybe_expand_parameters,
     _run_forward,
+    _validate_input,
 )
 from captum._utils.exceptions import FeatureAblationFutureError
 from captum._utils.progress import NullProgress, progress, Progress
@@ -501,6 +502,7 @@ def attribute(
         is_inputs_tuple = _is_tuple(inputs)
 
         formatted_inputs, baselines = _format_input_baseline(inputs, baselines)
+        _validate_input(formatted_inputs, baselines, allow_broadcastable_baselines=True)
         formatted_additional_forward_args = _format_additional_forward_args(
             additional_forward_args
         )
@@ -782,9 +784,9 @@ def _construct_ablated_input_across_tensors(
                 tensor_mask.append(mask)
 
                 assert baseline is not None, "baseline must be provided"
-                ablated_feature = input_tensor[start_idx:end_idx] * (1 - mask).to(
-                    input_tensor.dtype
-                ) + (baseline * mask.to(input_tensor.dtype))
+                ablated_feature = torch.where(
+                    mask.bool(), baseline, input_tensor[start_idx:end_idx]
+                )
                 ablated_input = ablated_input.to(ablated_feature.dtype)
                 ablated_input[start_idx:end_idx] = ablated_feature
             current_masks.append(torch.stack(tensor_mask, dim=0))
@@ -834,6 +836,7 @@ def attribute_future(
         # converting it into a tuple.
         is_inputs_tuple = _is_tuple(inputs)
         formatted_inputs, baselines = _format_input_baseline(inputs, baselines)
+        _validate_input(formatted_inputs, baselines, allow_broadcastable_baselines=True)
         formatted_additional_forward_args = _format_additional_forward_args(
             additional_forward_args
         )
@@ -1208,6 +1211,10 @@ def _process_ablated_out_full(
                 eval_diff_shape + (inputs[i].dim() - 1) * (1,)
             )
             eval_diff = eval_diff.to(total_attrib[i].device)
-            total_attrib[i] += (eval_diff * mask.to(attrib_type)).sum(dim=0)
+            total_attrib[i] += torch.where(
+                mask.to(device=eval_diff.device, dtype=torch.bool),
+                eval_diff,
+                torch.zeros_like(eval_diff),
+            ).sum(dim=0)
 
         return total_attrib, weights

captum/attr/_core/feature_permutation.py

@@ -9,7 +9,11 @@
 from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union
 
 import torch
-from captum._utils.common import _format_output, _format_tensor_into_tuples
+from captum._utils.common import (
+    _format_output,
+    _format_tensor_into_tuples,
+    _generate_derangement,
+)
 from captum._utils.typing import BaselineType, TargetType, TensorOrTupleOfTensorsGeneric
 from captum.attr._core.feature_ablation import FeatureAblation
 from captum.log import log_usage
@@ -26,9 +30,16 @@ def _permute_feature(x: Tensor, feature_mask: Tensor) -> Tensor:
     while (perm == no_perm).all():
         perm = torch.randperm(n)
 
-    return (x[perm] * feature_mask.to(dtype=x.dtype)) + (
-        x * feature_mask.bitwise_not().to(dtype=x.dtype)
-    )
+    return torch.where(feature_mask.to(device=x.device, dtype=torch.bool), x[perm], x)
+
+
+def _permute_feature_without_self_donors(x: Tensor, feature_mask: Tensor) -> Tensor:
+    n = x.size(0)
+    assert n > 1, "cannot permute features with batch_size = 1"
+
+    perm = _generate_derangement(n, x.device)
+
+    return torch.where(feature_mask.to(device=x.device, dtype=torch.bool), x[perm], x)
 
 
 class FeaturePermutation(FeatureAblation):
@@ -60,6 +71,11 @@ class FeaturePermutation(FeatureAblation):
     This method, unlike other attribution methods, requires a batch
     of examples to compute attributions and cannot be performed on a single example.
 
+    The default permutation may leave individual rows fixed. Set
+    ``exclude_self_donors=True`` to require every selected value to come from a
+    different row. Distinct donor rows may still contain equal feature values;
+    those are valid samples from the empirical feature distribution.
+
     By default, each scalar value within
     each input tensor is taken as a feature and shuffled independently. Passing
     a feature mask allows grouping features to be shuffled together (including
@@ -82,7 +98,8 @@ class FeaturePermutation(FeatureAblation):
     def __init__(
         self,
         forward_func: Callable[..., Union[int, float, Tensor, Future[Tensor]]],
-        perm_func: Callable[[Tensor, Tensor], Tensor] = _permute_feature,
+        perm_func: Callable[[Tensor, Tensor], Tensor] | None = None,
+        exclude_self_donors: bool = False,
     ) -> None:
         r"""
         Args:
@@ -95,9 +112,20 @@ def __init__(
                 which applies a random permutation, this argument only needs
                 to be provided if a custom permutation behavior is desired.
                 Default: `_permute_feature`
+            exclude_self_donors (bool, optional): If True, every selected row
+                receives its value from a different donor row. This cannot be
+                combined with a custom ``perm_func``. Default: False
         """
         FeatureAblation.__init__(self, forward_func=forward_func)
-        self.perm_func = perm_func
+        if exclude_self_donors and perm_func is not None:
+            raise ValueError(
+                "exclude_self_donors cannot be combined with a custom perm_func"
+            )
+        self.perm_func = (
+            _permute_feature_without_self_donors
+            if exclude_self_donors
+            else perm_func if perm_func is not None else _permute_feature
+        )
         # Considering the case when we permute multiple input tensors at once
         # through `feature_mask`, we disregard the feature group if the 0th
         # dim of *any* input tensor in the group is less than

captum/attr/_core/shapley_value.py

@@ -37,6 +37,7 @@
     _is_mask_valid,
     _is_tuple,
     _run_forward,
+    _validate_input,
 )
 from captum._utils.exceptions import ShapleyValueFutureError
 from captum._utils.progress import progress
@@ -83,6 +84,21 @@ def _shape_feature_mask(
     return tuple(mask_list)
 
 
+def _validate_shapley_feature_mask(feature_mask: Tuple[Tensor, ...]) -> None:
+    for mask in feature_mask:
+        if mask.is_complex():
+            values_are_integral = False
+        elif mask.is_floating_point():
+            values_are_integral = bool(
+                torch.isfinite(mask).all() and (mask == mask.round()).all()
+            )
+        else:
+            values_are_integral = True
+        assert values_are_integral and (
+            mask.numel() == 0 or bool((mask >= 0).all())
+        ), "Feature mask values must be non-negative integers."
+
+
 class ShapleyValueSampling(PerturbationAttribution):
     """
     A perturbation based approach to compute attribution, based on the concept
@@ -320,10 +336,12 @@ def attribute(
         # converting it into a tuple.
         is_inputs_tuple = _is_tuple(inputs)
         inputs_tuple, baselines = _format_input_baseline(inputs, baselines)
+        _validate_input(inputs_tuple, baselines)
         additional_forward_args = _format_additional_forward_args(
             additional_forward_args
         )
         formatted_feature_mask = _format_feature_mask(feature_mask, inputs_tuple)
+        _validate_shapley_feature_mask(formatted_feature_mask)
         reshaped_feature_mask = _shape_feature_mask(
             formatted_feature_mask, inputs_tuple
         )
@@ -487,10 +505,12 @@ def attribute_future(
     ) -> Future[TensorOrTupleOfTensorsGeneric]:
         is_inputs_tuple = _is_tuple(inputs)
         inputs_tuple, baselines = _format_input_baseline(inputs, baselines)
+        _validate_input(inputs_tuple, baselines)
         additional_forward_args = _format_additional_forward_args(
             additional_forward_args
         )
         formatted_feature_mask = _format_feature_mask(feature_mask, inputs_tuple)
+        _validate_shapley_feature_mask(formatted_feature_mask)
         reshaped_feature_mask = _shape_feature_mask(
             formatted_feature_mask, inputs_tuple
         )
@@ -754,7 +774,11 @@ def _eval_fut_to_prev_results_tuple(
             )
 
             # aggregate n_perturb
-            cur_attr = (formatted_eval_diff * cur_mask.float()).sum(dim=0)
+            cur_attr = torch.where(
+                cur_mask.to(device=formatted_eval_diff.device, dtype=torch.bool),
+                formatted_eval_diff,
+                torch.zeros_like(formatted_eval_diff),
+            ).sum(dim=0)
             # (*output_shape, *input_feature_shape)
             total_attrib[j] += cur_attr
 
@@ -805,10 +829,11 @@ def _update_current_tensors(
         for i in range(len(current_tensors)):
             if i in feat_list:
                 output_tensors.append(
-                    current_tensors[i]
-                    * (~(mask[i] == feature_index)).to(current_tensors[i].dtype)
-                    + input_tensors[i]
-                    * (mask[i] == feature_index).to(input_tensors[i].dtype)
+                    torch.where(
+                        (mask[i] == feature_index).to(current_tensors[i].device),
+                        input_tensors[i],
+                        current_tensors[i],
+                    )
                 )
 
             else:

captum/attr/_utils/common.py

@@ -281,7 +281,16 @@ def _tensorize_baseline(
     # pyre-fixme[2]: Parameter must be annotated.
     def _tensorize_single_baseline(baseline, input):
         if isinstance(baseline, (int, float)):
-            return torch.full_like(input, baseline)
+            baseline_dtype = (
+                input.dtype
+                if input.dtype == torch.bool
+                else torch.result_type(input, baseline)
+            )
+            return torch.full_like(
+                input,
+                baseline,
+                dtype=baseline_dtype,
+            )
         if input.shape[0] > baseline.shape[0] and baseline.shape[0] == 1:
             return torch.cat([baseline] * input.shape[0])
         return baseline
@@ -292,6 +301,10 @@ def _tensorize_single_baseline(baseline, input):
             type(baselines), type(inputs)
         )
     )
+    assert len(inputs) == len(baselines), (
+        "Input and baseline must have the same dimensions, baseline has "
+        f"{len(baselines)} features whereas input has {len(inputs)}."
+    )
     return tuple(
         _tensorize_single_baseline(baseline, input)
         for baseline, input in zip(baselines, inputs)