Add attention extraction callback for gene-regulation analysis - #63
Open
sayangsep wants to merge 8 commits into
Open
Add attention extraction callback for gene-regulation analysis#63sayangsep wants to merge 8 commits into
sayangsep wants to merge 8 commits into
Conversation
Port the LogAttention callback and per-layer attention recompute hook from the DNA2Cell gene_regulation pipeline so that attention matrices from the gene/CRE modulators can be captured during a regular forward pass of VariantFormer. * seq2gene/modules/layers.py: extend FlashAttLayer with `log_attn_matrix` /`attn_matrix` state and `calculate_attention_matrix`, a manual PyTorch implementation of softmax(QKᵀ/√d + ALiBi)·V that reproduces FlashAttention's output (verified within fp16 tolerance) so the attention can be inspected without changing the forward output. * seq2gene/attn_log_callback.py: Lightning callback / `record_attention` context manager that toggles `log_attn_matrix` on the right encoder layers in either the dual- or combined-modulator model variant and collects matrices keyed by `<modulator>_<mha>_<layer_idx>`. Supports `keep_heads=True` to retain the per-head (H, Q, K) tensor instead of the default head-averaged (Q, K). * tests/test_attn_log_callback.py: 31 unit tests, including a TestManualAttentionMatchesFlashAttention class that confirms the manual computation matches `layer.MHA(src)` for self/cross attention with and without ALiBi, plus a torch SDPA cross-check. * notebooks/explore_attention.ipynb: end-to-end demo on APOE that loads the v4_ag checkpoint, runs a forward pass through the example HG00096 VCF with `record_attention(keep_heads=True)`, and visualizes the CLS → CRE cross-attention vs. distance-to-TSS, the ranked top-CREs table, the full gene-token × CRE heatmap and per-head heatmaps. The notebook also surfaces both the original ENCODE cCRE coordinates and the model-visible window (padded by `cre_neighbour_hood = 50` bp on each side) so attention positions are interpreted unambiguously. Co-authored-by: Cursor <cursoragent@cursor.com>
Adjacent CREs can be closer than 100 bp, so showing the model-visible window ([start - 50, end + 50]) caused bars to overlap in the per-layer attention plot and made the ranked-CRE table needlessly wide. The notebook now uses the original ENCODE cCRE coordinates throughout (midpoints and distance-to-TSS are unaffected by the symmetric padding, which is documented in the markdown but not materialized in `cre_df`). Co-authored-by: Cursor <cursoragent@cursor.com>
The LogAttention callback previously imported `torch.utils.tensorboard. SummaryWriter` at module top level and exposed a `log_heatmaps=True` mode that wrote attention heatmaps to a TensorBoard run directory. No caller in this repo uses that path (the gene_regulation pipeline constructs `LogAttention(log_heatmaps=False, ...)`), so the only effect of keeping it was forcing every importer of `seq2gene.attn_log_callback` to install `tensorboard`. Remove the heatmap-to-TensorBoard branch entirely, including the `log_heatmaps` constructor flag, the `self.writer` attribute, and the `write_figure_to_tensorboard` helper. The matplotlib + seaborn helpers (`create_heatmap`, `create_heatmap_from_matrix`) stay; users who want figures can call them and plot or save themselves. Also rename the first query token of the gene cross-attention from "CLS" to "gene-tissue CLS" throughout the notebook to reflect that under `gene_pooling='multi_registry'` it's a per-tissue learnable embedding (`MultiRegistry.registry_tokens[tissue_id]`) prepended to the gene tokens, not a vanilla CLS. Variable: `gene_tissue_cls_attention`. Tests updated to drop the now-removed `log_heatmaps=False` kwarg; all 31 still pass without `tensorboard` installed. Co-authored-by: Cursor <cursoragent@cursor.com>
Use the more explicit `gene_start_site` terminology instead of the abbreviation `TSS` throughout the notebook (variable names, column names, plot labels, and accompanying markdown). Co-authored-by: Cursor <cursoragent@cursor.com>
Previously the notebook re-derived CRE coordinates from the manifest and mirrored process_subject's sort + minus-strand reversal, then used ``min(K_attn, len(cre_df))`` to reconcile. That assumed any drops introduced by ``ExtractSeqFromBed.process_subject`` happened only at the tail; a drop in the middle would silently misalign every CRE after the gap with the cross-attention key axis. Mirror the DNA2Cell/gene_regulation pipeline by exposing the dataset's own post-filter positions: * Factor the manifest → process_subject → strand-flip path out of ``VCFDataset._get_cres`` into a private ``_resolve_cres`` helper. * Add a public ``VCFDataset.get_cre_positions(gene_id, vcf_path)`` that reuses the helper and returns the per-CRE positions in model order (drops anywhere in the manifest honored by construction). * Update ``notebooks/explore_attention.ipynb`` to call ``get_cre_positions`` directly. The K-axis reconciliation now only has to handle attention-side padding, with an assert guarding against the unexpected ``K_attn < len(cre_df)`` case. Co-authored-by: Cursor <cursoragent@cursor.com>
Adds a "Repeat on the reference genome" section after the per-head heatmap that swaps in vcf_path=None, runs the same gene/tissue/layer configuration through a fresh LogAttention callback, and overlays the gene-tissue CLS attention from the reference and HG00096 VCF runs at the last logged layer. Useful as a baseline for variant-effect comparisons; intro and recap are updated to point at the new section. Co-authored-by: Cursor <cursoragent@cursor.com>
With a single-sample batch the K-axis equals len(cre_df) exactly, so the soft trim was a no-op in practice. Replace it with a hard ``assert K_attn == len(cre_df)`` so any future change that violates the 1-to-1 alignment (e.g. multi-sample batches with mismatched CRE counts) fails loudly instead of silently re-shaping attention. Co-authored-by: Cursor <cursoragent@cursor.com>
Drop cached cell outputs (figures and prints) from ``explore_attention.ipynb`` so the committed notebook stays small and diff-friendly; users re-run it locally to regenerate plots. Co-authored-by: Cursor <cursoragent@cursor.com>
barhomi
approved these changes
May 7, 2026
barhomi
left a comment
barhomi
left a comment
Collaborator
There was a problem hiding this comment.
LGTM as long as we pass the Attention matching test with the latest models
|
|
||
|
|
||
| @pytest.mark.skipif(not CUDA_AVAILABLE, reason=SKIP_CUDA_REASON) | ||
| class TestRecordAttentionEndToEnd(unittest.TestCase): |
Collaborator
There was a problem hiding this comment.
@sayangsep was this run on the current big models?
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
2 participants
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.
Port the LogAttention callback and per-layer attention recompute hook from the DNA2Cell gene_regulation pipeline so that attention matrices from the gene/CRE modulators can be captured during a regular forward pass of VariantFormer.
log_attn_matrix/attn_matrixstate andcalculate_attention_matrix, a manual PyTorch implementation of softmax(QKᵀ/√d + ALiBi)·V that reproduces FlashAttention's output (verified within fp16 tolerance) so the attention can be inspected without changing the forward output.record_attentioncontext manager that toggleslog_attn_matrixon the right encoder layers in either the dual- or combined-modulator model variant and collects matrices keyed by<modulator>_<mha>_<layer_idx>. Supportskeep_heads=Trueto retain the per-head (H, Q, K) tensor instead of the default head-averaged (Q, K).layer.MHA(src)for self/cross attention with and without ALiBi, plus a torch SDPA cross-check.record_attention(keep_heads=True), and visualizes the CLS → CRE cross-attention vs. distance-to-TSS, the ranked top-CREs table, the full gene-token × CRE heatmap and per-head heatmaps. The notebook also surfaces both the original ENCODE cCRE coordinates and the model-visible window (padded bycre_neighbour_hood = 50bp on each side) so attention positions are interpreted unambiguously.