Skip to content

Separate attention-map capture from gradient hooks - #1232

Draft
LIU-kaiyu wants to merge 1 commit into
sunlabuiuc:masterfrom
LIU-kaiyu:liu-kaiyu/1166-attention-capture
Draft

Separate attention-map capture from gradient hooks#1232
LIU-kaiyu wants to merge 1 commit into
sunlabuiuc:masterfrom
LIU-kaiyu:liu-kaiyu/1166-attention-capture

Conversation

@LIU-kaiyu

Copy link
Copy Markdown

Fixes #1166.

AttentionRollout previously enabled backward hooks despite using only attention maps, causing it to fail inside torch.no_grad(). This change separates map capture from gradient capture in Transformer and StageAttentionNet.

Introduces AttentionInterpretable and its subclass GradientInterpretable, retaining CheferInterpretable as an alias. Rollout explicitly requests forward-only capture and runs without autograd; Chefer explicitly requests gradients. Attribution calculations remain unchanged.

Existing positional register_hook arguments and concrete models’ one-argument hook behavior are preserved. Disabling capture stops both modes while retaining the latest results; the next attention forward clears stale gradients and replaces or clears the map. Chefer now raises clear errors for empty layers, missing tensors, and mismatched map/gradient shapes.

Compatibility: custom models must explicitly inherit the appropriate interface and accept the capture_gradients keyword. The alias does not adapt obsolete method signatures. Overlapping captured forwards/backwards on one model remain unsupported and are documented.

Validation: 36 focused tests passed in Pixi, including 11 new tests covering no-grad execution, unchanged parameter gradients, interpreter switching, cache lifecycle, positional compatibility, migration errors, and failure cleanup. All changed executable statements were covered. The synthetic example, new docstring examples, and #1166-specific contribution checks passed.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Refactor: split CheferInterpretable into AttentionInterpretable and GradientInterpretable

1 participant