Sparse Autoencoders Find Causal, Lineage-Specific Context Features in Chromatin Foundation Models

Nicole Ching, Ayushi Mehrotra

Keywords: Sparse Autoencoders, Steering, Causal interventions

Other Keywords: gene expression, chromatin, foundation models

TL;DR: SAE methodology transfers cleanly from language to genomics, and that our proposed metric provides a general primitive for identifying contrastive concepts in any probed transformer.

Abstract:

Sparse autoencoders (SAEs) have produced important insights in language model interpretability, but their utility on transformers trained on scientific data remains underexplored. We extend the SAE-plus-causal-intervention toolkit to an epigenomics foundation model, EpiBERT, and ask whether it internally encodes a biologically meaningful contrast: in vitro (cell line) vs. in vivo (primary tissue) chromatin context. We train layer-wise Sparse Autoencoders (SAEs) with BatchTopK activations across six matched ATAC-seq conditions spanning blood, liver, and lymph lineages, introduce the Context Divergence Score (CDS)—a contrastive t-statistic applicable to any probed transformer—to identify context-specific features, and validate them through causal ablation, linear context-steering, and three-level biological annotation (ChromHMM, HOMER, GO:BP). We find a depth-stratified context representation: context-specific features
grow 3.8-fold from early to late layer (57 → 215 Bonferroni-significant), mirroring the late-layer concentration of high-level features in language model SAEs. Causal ablation of CDS-selected features yields a large effect, context-steering closes 11.2% of the prediction gap at 4.5× above random, and biological annotation grounds the discovered features in lineage-defining transcription factors. These results demonstrate that the SAE methodology transfers cleanly from language to genomics, and that CDS provides a general primitive for identifying contrastive concepts in any probed transformer. Code is available at https://anonymous.4open.science/r/in_vivo_vs_in_vitro_chromatin_contexts/.

Submission Number: 77