Repository navigation
Enable DeepSHAP for Enformer models - #201
Open
avantikalal wants to merge 4 commits into
Open
avantikalal wants to merge 4 commits into
avantikalal wants to merge 4 commits into
Conversation
tangermeme 1.5.0 added integrated_gradients_op, which builds a per-layer integrated-gradients hook for modules that DeepLIFT/SHAP cannot register directly (see jmschrei/tangermeme#3). Enformer's Attention and AttentionPool layers, imported verbatim from enformer_pytorch, are exactly this case: they use functional softmax/einsum calls instead of registered submodules. Register integrated_gradients_op for these two layers when running deepshap on EnformerModel/EnformerPretrainedModel instead of raising NotImplementedError, and bump the tangermeme dependency accordingly. Fixes #56 Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
enformer_pytorch's GELU is a custom sigmoid-based approximation (torch.sigmoid(1.702 * x) * x), not torch.nn.GELU, so it is not covered by tangermeme's default nonlinear-op table any more than Attention/AttentionPool are. Register it with tangermeme's exact elementwise rule (_nonlinear), as jmschrei recommended in jmschrei/tangermeme#3. Verified on a real trained EnformerPretrainedModel checkpoint (Genentech/human-atac-catlas-model): adding this hook substantially reduced both the attribution scale and DeepLIFT/SHAP's internal convergence-delta warnings compared to registering only Attention/AttentionPool. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
gReLU models pooled to a single value per task return shape (N, tasks, 1), not (N, tasks). tangermeme.deep_lift_shap indexes a single task with `y[:, target]`, which only consumes the tasks axis and leaves the trailing 1 in place. The resulting shape mismatch between its output-difference (N/2, 1) and attribution-sum (N/2,) tensors is then silently broadcast into a meaningless (N/2, N/2) matrix instead of raising a shape error, corrupting its internal convergence-delta check and producing spurious/inflated "Convergence deltas too high" warnings. Wrap the model passed to deep_lift_shap in a small adapter (_SqueezeOutput) that drops a trailing length-1 axis, restoring the flat (N, tasks) shape deep_lift_shap's convergence check expects. Models with genuine multi-position outputs (shape (N, tasks, L>1)) are left untouched. Verified against a real trained Enformer-based checkpoint (Genentech/human-atac-catlas-model): the same convergence deltas that previously appeared as a bogus (N, N) matrix with values up to ~10 now come out correctly shaped and an order of magnitude smaller. Fixes #202 Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
3 tasks done
Wrap the model in _SqueezeOutput for the Enformer deepshap path too, so its convergence-delta check isn't corrupted by the same shape mismatch fixed for the general case in #202/#203. Without this, the Attention/AttentionPool/GELU hooks added here couldn't be evaluated against a trustworthy convergence signal. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
tangermemedependency to>= 1.5.0, which addsintegrated_gradients_op(see jmschrei/tangermeme#3)NotImplementedErrorguard inget_attributionsthat blockedmethod="deepshap"forEnformerModel/EnformerPretrainedModelintegrated_gradients_op(K=8)as the DeepLIFT/SHAP hook for Enformer'sAttentionandAttentionPoollayers (imported fromenformer_pytorch), since both use functionalsoftmax/einsumcalls that can't be hooked directly by the default rescale ruleadditional_nonlinear_opspassed via**kwargsare merged with (and can override) these defaults rather than being silently droppedFixes #56
Test plan
setup.cfgandsrc/grelu/interpret/score.pyparse/import cleanlyget_attributions(..., method="deepshap")against a trainedEnformerPretrainedModeland confirm convergence deltas (print_convergence_deltas=True) are acceptably small on real weights — jmschrei's tangermeme fix was only validated against a toy re-implementation of these layers, not a real Enformer checkpointtests/test_interpret.pydocs/tutorials/4_design.ipynband5_variant.ipynbwhen this was originally disabled🤖 Generated with Claude Code