Skip to content

Harden core attribution edge cases (#1927) - #1927

Open
craymichael wants to merge 8 commits into
meta-pytorch:masterfrom
craymichael:export-D117635569
Open

Harden core attribution edge cases (#1927)#1927
craymichael wants to merge 8 commits into
meta-pytorch:masterfrom
craymichael:export-D117635569

Conversation

@craymichael

@craymichael craymichael commented Aug 28, 2026 •
edited by meta-codesync Bot
Loading

Copy link
Copy Markdown
Contributor

Summary:

Problem

  • DataLoader tuple handling and input_roles mapping disagreed with the public contract.
  • Arithmetic masking leaked NaN and Inf through masked positions.
  • LIME, DataLoader attribution, and Occlusion could combine tensor baselines with inputs on another device.
  • LIME always enabled pinned memory even when its surrogate-training tensors already lived on an accelerator.
  • An explicitly supplied feature-mask tuple containing only empty tensors reached min() with no elements.
  • KernelSHAP one-feature, sample-count, empty-pool, dtype, and boolean-baseline edge cases were unsafe.
  • Shapley masks and outputs could be allocated on the wrong device.
  • SVS-P indexed a full-shape 1-D mask as two-dimensional.
  • The ordinary permutation helper retried forever for a zero-row batch because the empty permutation equals the empty identity.

Counterexamples

  • A tuple batch was treated as one input, and roles 1 and 2 behaved in reverse.
  • 0 * NaN remained NaN instead of selecting the baseline.
  • A CPU tensor baseline combined with a meta/CUDA input failed in torch.where.
  • A meta/CUDA LIME sample produced a DataLoader with pin_memory=True, which cannot pin accelerator tensors.
  • An all-empty explicit feature mask raised ValueError: min() iterable argument is empty.
  • Float baseline 1.0 for a boolean input was coerced to bool because 1.0 == 1.
  • One-feature KernelSHAP divided by zero.
  • Mixed-device Shapley used input 0 device for every tensor.
  • Input and feature mask shaped [3] raised IndexError: too many indices for tensor of dimension 1 in both main and UFO SVS-P.
  • _generate_permutation(0) repeatedly called torch.randperm(0) and could never escape its identity-rejection loop.

Fix

  • Correct tuple/role handling and reject unknown roles.
  • Use boolean selection in DataLoader attribution, LIME, and Occlusion.
  • Move tensor baselines to each input device before selection.
  • Pin LIME surrogate data only when every materialized training tensor is host-resident.
  • Reject all-empty explicit feature masks with a descriptive validation error.
  • Validate LIME masks and sampling, reject empty generators/pools, handle one-feature KernelSHAP, and preserve promoted baseline dtype.
  • Restrict boolean preservation to actual bool scalar baselines.
  • Allocate Shapley state per input device and preserve output precision.
  • Compare full-shape masks against mask[0:1], which works for any positive tensor rank.
  • Return the identity index tensor immediately for permutation sizes zero and one.

Differential Revision: D117635569

@meta-cla meta-cla Bot added the cla signed label Aug 28, 2026
@meta-codesync

meta-codesync Bot commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

@craymichael has exported this pull request. If you are a Meta employee, you can view the originating Diff in D117635569.

craymichael added a commit to craymichael/captum that referenced this pull request Aug 28, 2026
@craymichael @meta-codesync
Harden core attribution edge cases (meta-pytorch#1927)
54b0a47 
Summary:
Pull Request resolved: meta-pytorch#1927

Problem

- DataLoader tuple handling and `input_roles` mapping disagreed with the public contract.
- Arithmetic masking leaked NaN and Inf through masked positions.
- LIME, DataLoader attribution, and Occlusion could combine tensor baselines with inputs on another device.
- LIME always enabled pinned memory even when its surrogate-training tensors already lived on an accelerator.
- An explicitly supplied feature-mask tuple containing only empty tensors reached `min()` with no elements.
- KernelSHAP one-feature, sample-count, empty-pool, dtype, and boolean-baseline edge cases were unsafe.
- Shapley masks and outputs could be allocated on the wrong device.
- SVS-P indexed a full-shape 1-D mask as two-dimensional.
- The ordinary permutation helper retried forever for a zero-row batch because the empty permutation equals the empty identity.

Counterexamples

- A tuple batch was treated as one input, and roles 1 and 2 behaved in reverse.
- `0 * NaN` remained NaN instead of selecting the baseline.
- A CPU tensor baseline combined with a meta/CUDA input failed in `torch.where`.
- A meta/CUDA LIME sample produced a DataLoader with `pin_memory=True`, which cannot pin accelerator tensors.
- An all-empty explicit feature mask raised `ValueError: min() iterable argument is empty`.
- Float baseline `1.0` for a boolean input was coerced to bool because `1.0 == 1`.
- One-feature KernelSHAP divided by zero.
- Mixed-device Shapley used input 0 device for every tensor.
- Input and feature mask shaped `[3]` raised `IndexError: too many indices for tensor of dimension 1` in both main and UFO SVS-P.
- `_generate_permutation(0)` repeatedly called `torch.randperm(0)` and could never escape its identity-rejection loop.

Fix

- Correct tuple/role handling and reject unknown roles.
- Use boolean selection in DataLoader attribution, LIME, and Occlusion.
- Move tensor baselines to each input device before selection.
- Pin LIME surrogate data only when every materialized training tensor is host-resident.
- Reject all-empty explicit feature masks with a descriptive validation error.
- Validate LIME masks and sampling, reject empty generators/pools, handle one-feature KernelSHAP, and preserve promoted baseline dtype.
- Restrict boolean preservation to actual bool scalar baselines.
- Allocate Shapley state per input device and preserve output precision.
- Compare full-shape masks against `mask[0:1]`, which works for any positive tensor rank.
- Return the identity index tensor immediately for permutation sizes zero and one.

Differential Revision: D117635569
@meta-codesync meta-codesync Bot changed the title Harden core attribution edge cases Harden core attribution edge cases (#1927) Aug 28, 2026
@craymichael @meta-codesync
Validate perturbation baselines before attribution
50345ff 
Summary:
Summary

Feature Ablation and Shapley formatted baselines but did not validate tuple arity or tensor shape before running the model. `_tensorize_baseline` also paired inputs and baselines with `zip`, silently discarding extra baselines.

Problem

For two inputs, a one-element baseline tuple reached the model with a missing argument, while a three-element tuple silently ignored its final baseline. A baseline pool with shape `[2, 1]` for a three-row input was neither a per-example baseline nor a singleton baseline and failed later through opaque broadcasting errors.

Fix

* Validate baselines in both synchronous and future Feature Ablation and Shapley entry points.
* Defensively reject arity mismatches in `_tensorize_baseline`.
* Require Shapley tensor baselines to match the input or use a singleton leading dimension with matching trailing dimensions.
* Preserve Feature Ablation’s documented support for any tensor shape that broadcasts exactly to the input, including `[F]` and 0-D tensors.

Test Plan

Before: the new regressions failed through late `IndexError`, missing-forward-argument, and broadcast errors; the extra-baseline case did not raise at all.

After:
* `buck test fbcode//pytorch/captum/tests/attr:test_common fbcode//pytorch/captum/tests/attr:test_feature_ablation fbcode//pytorch/captum/tests/attr:test_shapley` — Pass 131, Fail 0.
* `arc lint -a` on changed Python files — no source lint issues; focused autodeps updates applied.
* `arc lint -a --engine extra --take CITRINEAGENT` on changed Captum implementation files — no issues.
* `arc pyre check-owning-targets` on changed files — no type errors.

Differential Revision: D117601314
@craymichael @meta-codesync
Select perturbed values without arithmetic masking
9b7cb72 
Summary:
Summary

Perturbation paths combined original and replacement values as `old * (1 - mask) + new * mask`. IEEE arithmetic makes `NaN * 0` and `Inf * 0` non-finite, so values outside the selected feature could corrupt model inputs and attributions.

Counterexample

For input `[1, 2]`, baseline `[0, NaN]`, and mask `[0, 1]`, ablating feature 0 should evaluate `[0, 2]`. Arithmetic masking instead produced `[0, NaN]`, so feature 0 incorrectly received a `NaN` attribution.

Fix

* Use `torch.where` for Feature Ablation, Feature Permutation, Shapley feature updates, and attribution-mask accumulation.
* Apply the same selection semantics to within-group baseline construction, add-back, and permutation helpers.
* Mirror the fix in the legacy UFO implementations so alternate call paths cannot retain the corruption.
* Keep masks and donor indices on the destination tensor device.

Test Plan

Before: six focused regressions failed across core and within-group paths (`Pass 159, Fail 6`).

After:
* `buck test fbcode//pytorch/captum/tests/attr:test_feature_ablation fbcode//pytorch/captum/tests/attr:test_feature_permutation fbcode//pytorch/captum/tests/attr:test_shapley fbcode//pytorch/captum/tests/attr/fb:test_within_groups_utils fbcode//pytorch/captum/tests/attr/fb:test_shapley_value_permutation` — Pass 178, Fail 0.
* Regressions cover `NaN`, `+Inf`, and `-Inf` in selected and inactive positions.
* `arc lint -a` on all changed Python files — no new lint issues; existing legacy UFO line-length advice remains unchanged.
* `arc lint -a --engine extra --take CITRINEAGENT` on implementation files — no issues.
* `arc pyre check-owning-targets` on changed files — no type errors.

Differential Revision: D117601315
@craymichael @meta-codesync
Validate perturbation feature masks before attribution
37146b8 
Summary:
Summary

Perturbation methods accepted feature-mask tuples with the wrong arity. Feature Ablation could then return zero attribution for an omitted input, while Shapley failed later or enumerated an invalid feature space. Shapley also accepted fractional, negative, non-finite, and complex group IDs even though its permutation loop requires non-negative integer IDs.

Counterexample

With two model inputs and a one-element feature-mask tuple, Feature Ablation perturbed only the first input and silently left the second input’s attribution at zero. With Shapley mask `[0, 0.5]`, the integer feature loop omitted group `0.5`.

Fix

* Require explicit feature-mask tuples to have one tensor per input in the shared formatter.
* Validate Shapley group IDs before feature enumeration.
* Accept integral-valued float masks and bool masks for compatibility.
* Keep value-domain validation out of Feature Ablation/Permutation so Greedy Feature Selection’s internal `-1` sentinel and existing float-typed masks continue to work.

Test Plan

Before: four focused regression methods failed (`Pass 132, Fail 4`), including silent acceptance of short and long mask tuples and invalid Shapley IDs.

After:
* Broad perturbation and wrapper suite — Pass 422, Fail 0, covering Feature Ablation, Feature Permutation, Shapley, DataLoaderAttribution, internal wrappers, WithinGroupSVS, SVS-P, AddOneBack, MarginalWithinGroups, Greedy Feature Selection, and shared mask utilities.
* `arc lint -a` on changed files — no issues.
* `arc lint -a --engine extra --take CITRINEAGENT` on implementation files — no issues.
* `arc pyre check-owning-targets` on changed files — no type errors.

Differential Revision: D117601317
@craymichael @meta-codesync
Preserve fractional scalar baselines for integer inputs
967917d 
Summary:
Summary

`_tensorize_baseline` used `torch.full_like(input, baseline)`, which silently cast a floating-point scalar baseline to the input dtype. Integer inputs therefore lost fractional baseline values before Shapley attribution.

Counterexample

For integer input `[1, 2]` and scalar baseline `0.5`, additive Shapley attribution should be `[0.5, 1.5]`. The old tensorization converted the baseline to `[0, 0]` and returned `[1, 2]`.

Fix

Use PyTorch result-type promotion when materializing scalar baselines, while retaining `full_like` so device, layout, and memory format are preserved. Boolean inputs keep their historical boolean baseline dtype.

Test Plan

Before: the new sync regression failed with both elements off by `0.5`.

After:
* `buck test` across Shapley, WithinGroupSVS, DeepLift, LayerDeepLift, LayerIntegratedGradients, IntegratedGradients, and LayerConductance consumers — Pass 196, Fail 0.
* Focused common/Shapley/WithinGroupSVS rerun — Pass 97, Fail 0.
* Unit coverage verifies fractional integer promotion, boolean dtype preservation, and channels-last layout preservation.
* `arc lint -a` — no issues.
* `arc lint -a --engine extra --take CITRINEAGENT` — no issues.
* `arc pyre check-owning-targets` — no type errors.

Differential Revision: D117601321
@craymichael @meta-codesync
Add an opt-in no-self donor policy
c1f59c8 
Summary:
Problem
- Ordinary permutation importance may legally assign a row to itself.
- That makes an explicitly requested leave-one-out donor interpretation false.
- Equal-valued replacements are not the same bug: a distinct donor may validly carry the same categorical or constant value.
- The first implementation also changed scripted-model defaults, exposed a private callback as the API, duplicated derangement logic, and forwarded arbitrary future config fields into Captum.
- Rank Suggest's CLI temporarily hard-coded `--n-samples` support to SVS-P, regressing plain SVS even though its typed args also define `n_samples`.

Simple counterexample
For rows `[A, B, C]`, permutation `[A, C, B]` leaves row 0 unchanged. That is valid standard permutation importance, but invalid when the caller asks that every recipient use a different donor. Conversely, two distinct rows may both contain `A`; rejecting that replacement would bias the empirical donor distribution.

Fix
- Preserve ordinary permutation behavior by default.
- Add public `exclude_self_donors` constructor/config options for FeaturePermutation and both SVS-P implementations.
- Generate a no-fixed-point random cycle with one device-side `randperm`, avoiding retry loops and host synchronization.
- Reuse one donor ordering across grouped and nested tensors so one synthetic row always comes from one donor row.
- Reject ambiguous combinations with custom `perm_func` or `baseline_generator` callbacks.
- Route only explicitly supported algorithm kwargs instead of forwarding every dataclass field.
- Raise typed validation errors for unsupported algorithm/config combinations.
- Validate `--n-samples` by the selected args dataclass's actual fields, preserving plain SVS and Kernel SHAP while continuing to reject shuffling.

Differential Revision: D117601323
@craymichael @meta-codesync
Weight distributed attribution aggregation by samples
12556a5 
Summary:
- Weight rank/batch attribution means and coverage by their actual example counts.
- Infer the local count from labels first, allowing broadcast weights without confusing their leading dimension for the batch size; fall back to weights only when labels are unavailable.
- Gather sample counts on the same device as attribution collectives, including `device=None` callers using NCCL.
- Gather counts alongside results and skip zero-sample ranks without skipping collectives.
- Scope every barrier, gather, and reduction to the supplied process group and translate subgroup rank 0 to its global destination rank.
- Extend Captum online mean, variance, standard deviation, and sum statistics with backward-compatible frequency weights.

Counterexamples:
- Rank 0 reports mean attribution 1 from one example. Rank 1 reports mean attribution 3 from three examples. Equal rank weighting returns 2.0; the correct example-weighted global mean is `(1 * 1 + 3 * 3) / 4 = 2.5`.
- A four-row label with a broadcast weight shaped `[1]` was rejected as inconsistent, even though the label establishes the true batch size and the weight is validly broadcastable.
- With `device=None`, accelerator attribution tensors used an accelerator collective while the new sample-count tensor was created on CPU, which is invalid for NCCL.

Implementation:
The evaluator derives each local batch count from internally consistent label tensors, or from weight tensors only when labels do not expose a count. The publisher gathers those counts on the attribution collective device and in the same process group as attribution tensors. The singleton summarizer applies them as frequency weights. Existing unweighted callers retain weight 1 and custom Stat implementations still receive their legacy one-argument `update` call.

Differential Revision: D117601324
@craymichael @meta-codesync
Make async and grouped perturbation attribution correct
98b8be8 
Summary:
Problem
- Async ablation mutated shared accumulators from completion callbacks.
- Output-shape validation leaked across calls.
- Async permutation ignored sampling and skipped-forward semantics.
- Cross-tensor feature groups could use incoherent donor rows.
- Permutation indices were created without the input device, and singleton batches could enter invalid or nonterminating paths.

Counterexamples
- Futures completing out of order could lose a contribution.
- Correlated tensors such as (user_id, user_age) could be sourced from different rows.
- A second call with a different output shape skipped validation.
- Batch size 1 could loop while searching for a non-identity permutation.
- Accelerator inputs could receive CPU permutation indices.

Fix
- Reduce async contributions independently and define zero-work results.
- Validate shape per call.
- Honor n_samples and run_forward_on_skip asynchronously.
- Share one donor mapping across a global feature group and require atomic grouped callbacks.
- Create permutation indices on the input device and make singleton permutation the identity.
- Add deterministic noncontiguous-ID and multi-block grouping coverage.

Differential Revision: D117635564
@craymichael @meta-codesync
Harden core attribution edge cases (meta-pytorch#1927)
bcd3752 
Summary:
Pull Request resolved: meta-pytorch#1927

Problem

- DataLoader tuple handling and `input_roles` mapping disagreed with the public contract.
- Arithmetic masking leaked NaN and Inf through masked positions.
- LIME, DataLoader attribution, and Occlusion could combine tensor baselines with inputs on another device.
- LIME always enabled pinned memory even when its surrogate-training tensors already lived on an accelerator.
- An explicitly supplied feature-mask tuple containing only empty tensors reached `min()` with no elements.
- KernelSHAP one-feature, sample-count, empty-pool, dtype, and boolean-baseline edge cases were unsafe.
- Shapley masks and outputs could be allocated on the wrong device.
- SVS-P indexed a full-shape 1-D mask as two-dimensional.
- The ordinary permutation helper retried forever for a zero-row batch because the empty permutation equals the empty identity.

Counterexamples

- A tuple batch was treated as one input, and roles 1 and 2 behaved in reverse.
- `0 * NaN` remained NaN instead of selecting the baseline.
- A CPU tensor baseline combined with a meta/CUDA input failed in `torch.where`.
- A meta/CUDA LIME sample produced a DataLoader with `pin_memory=True`, which cannot pin accelerator tensors.
- An all-empty explicit feature mask raised `ValueError: min() iterable argument is empty`.
- Float baseline `1.0` for a boolean input was coerced to bool because `1.0 == 1`.
- One-feature KernelSHAP divided by zero.
- Mixed-device Shapley used input 0 device for every tensor.
- Input and feature mask shaped `[3]` raised `IndexError: too many indices for tensor of dimension 1` in both main and UFO SVS-P.
- `_generate_permutation(0)` repeatedly called `torch.randperm(0)` and could never escape its identity-rejection loop.

Fix

- Correct tuple/role handling and reject unknown roles.
- Use boolean selection in DataLoader attribution, LIME, and Occlusion.
- Move tensor baselines to each input device before selection.
- Pin LIME surrogate data only when every materialized training tensor is host-resident.
- Reject all-empty explicit feature masks with a descriptive validation error.
- Validate LIME masks and sampling, reject empty generators/pools, handle one-feature KernelSHAP, and preserve promoted baseline dtype.
- Restrict boolean preservation to actual bool scalar baselines.
- Allocate Shapley state per input device and preserve output precision.
- Compare full-shape masks against `mask[0:1]`, which works for any positive tensor rank.
- Return the identity index tensor immediately for permutation sizes zero and one.

Differential Revision: D117635569
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

cla signed meta-exported

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@craymichael