Towards scientific discovery with dictionary learning: Extracting biological concepts from microscopy foundation models
Authors: Konstantin Donhauser
Kristina Ulicna
Gemma Elyse Moran
Aditya Ravuri
Kian Kenyon-Dean
Cian Eastwood
Jason Hartford
Venue: ICML 2025 poster (2025)
Keywords: mechanistic interpretability
scientific discovery
masked auto-encoders
Primary area: deep_learning->foundation_models
TL;DR: A first step towards the use of dictionary learning (a popular tool from mechanistic interpretability) for the scientific discovery of biological concepts, using microscopy foundation models

Abstract: Sparse dictionary learning (DL) has emerged as a powerful approach to extract semantically meaningful concepts from the internals of large language models (LLMs) trained mainly in the text domain. In this work, we explore whether DL can extract meaningful concepts from less human-interpretable scientific data, such as vision foundation models trained on cell microscopy images, where limited prior knowledge exists about which high-level concepts should arise. We propose a novel combination of a sparse DL algorithm, Iterative Codebook Feature Learning (ICFL), with a PCA whitening pre-processing step derived from control data. Using this combined approach, we successfully retrieve biologically meaningful concepts, such as cell types and genetic perturbations. Moreover, we demonstrate how our method reveals subtle morphological changes arising from human-interpretable interventions, offering a promising new direction for scientific discovery via mechanistic interpretability in bioimaging.

Towards scientific discovery with dictionary learning: Extracting biological concepts from microscopy foundation models
Konstantin Donhauser"
Kristina Uliena 22 Gemma Elyse Moran Aditya Ravuri +5 Kian Kenyon-Dean
Cian Eastwood 23
Jason Hartford 236
Abstract
encode many more concepts than they have neurons, and as a result, one cannot understand the model by inspecting individual neurons. One hypothesis for how neurons encode multiple concepts at once is that they are low-dimensional projections of some high-dimensional, sparse feature space. Quite surprisingly, there is now a large body of empirical evidence that supports this hypothesis in language models (Mikolov et al. 2013 Elhage et al., 2022] Park et al. 2023), games (Nanda al. 2023) and multimodal vision models (Rao et al. 2024). by showing that high-level features are typically predictable via linear probing. Further, recent work has shown that model representations can be decom- posed into learning models, estimated via sparse autoencoders (SAEs, Templeton 2024 Rajamanoharan et al. 2024b a Gao et al.
Sparse dictionary learning (DL) has emerged as a powerful approach to extract semantically mean- ingful concepts from the internals of large lan- guage models (LLMs) trained mainly in the text domain. In this work, we explore whether DL can extract meaningful concepts from less human- interpretable scientific data, such as vision foun- dation models trained on cell microscopy images, where limited prior knowledge exists about which high-level concepts should arise. We propose a novel combination of sparse DL algorithm, Iter- ative Codebook Feature Learning (ICFL), with a PCA whitening pre-processing step derived from control data. Using this combined approach. we successfully retrieve biologically meaning- ful concepts, such as cell types and genetic per- turbations. Moreover, we demonstrate how our method reveals subtle morphological changes aris- ing from human- comman-nterpreetbble interventions, offer- ing promising new direction for scientific discov- ery via mechanistic interpretability in bioimaging.
interpretable concepts using dictionary
2024).
However, these successes rely on some form of text su- pervision, either directly through next-token prediction indirectly via contrastive objectives like CL.IP (Radford et al. 2021). which align text and image representations. Further, these successes appear in domains which are natu- rally numan-interpretablee (f.e. text, games, natural images). As a result, one may worry that high-level features can be extracted only in settings that we already understand. This raises natural question: can extract similarly mean- ingful high-level features from completely unsupervised models in domains where we lack strong prior knowledge?
1. Introduction
Large scale machine learning systems are extremely effec- tive at generating realistic text and images. However, these models remain black boxes: it is difficult to understand how they produce such detailed reconstructions, and what extent they encode semantic information about the target domain in their internal representations. Investigation of how models encode and use high-level, numan-interpretalee concepts is challenging due to the "superposition hypothe- sis" [Bricken et al. 2023). which states that neural networks
For example, in computational biology, masked autoen- coders (MAE) trained on cellular microscopy images have been shown to be very effective learning representations that recover known biological relationships (Kraus al.) 2024 However, it is not known whether analogous high- level features can be extracted from large MAE-based foun- dation models of cell microscopy data. These settings are precisely where extracting high-level features could be most valuable: given that models can detect subtle differences in images (even those that are very challenging for human experts to interpret), we might hope that the mechanistic in- terpretability techniques would enable better understanding of the ways in which these subtle differences arise.
Equal contribution Work done while interning at Va- Ience Labs. 'ETH Zürich 'Valence Labs 'Recursion 'Rttges University 'University of Cambridge Uniivessity of Manch- ester. Correspondence to: Kristina Uliena, Jason Hartford <kristina@ valencelabs.com com,, @ason@valcncelabss valencclabs.com com..
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
In this work, we study the extraction of high-level features

---

Towards scientific discovery with dictionary learning
(A)
original imssge
MAE
reconsizucted imagge
(D)
feature #1
CDC
dec
intermediate tokens
biological concepts
Sesturre
(c) feature
Herration: O
directions
tolen
rosidual
Heration
fecture
xif rosidaal
aph
token
Seature
4
Iberation: :
MAX
concept direction correlation
min
Figure 1. Graphical abstract illustrating our algorithmic approach to biological concept interpretability from cell image data, consisting of (A) token-wise emheddings of cell images leamedd by a masked autoencoder (MAF), (B) reconstruction the averaged token 2, through large codebook W and sparse feature vector of sparsity K, which both get updated through (C) iterative codebook feature learning (ICFL) adapted from the matching pursuit algorithm for J iterations. The learned codebook columns represent feature directions (orange) which correlate with (D) biological concepts visualized as cell images ranked according to the correlation strength with three selected features. Each feature captures distinct cellular morphologies: Feature #1 activates for cells with an elongated, spindle-like shape (left) and anti-correlates for sparser and clumped cells (right); Feature #2 activates for cells that are densely packed with closely arranged nuclei (left) and dcactivates when cell density drops (right); and Feature #3 activates for compact, round, bright cells without cell-cell contacts almost entirely made up from just nuclei (left), in contrast to nulli--nucletedd cells which occupy larger areas (right).
from large-scale MAEs trained on microscopy images from a cell screening campaign containing cells that have been modulated with genetic and/or small molecule perturbations (Fay et al. 2023 Understanding the morphological changes induced by such perturbations is an inherently difficult and fundamental problem that plays a crucial role in drug dis- covery (Celik et al., 2024). Recent progress in this field has been made by using machine learning to learn represen- tations of perturbations based on their corresponding cell images (Figure A). One direction to analyze their meaning- fulness is to build relationship maps, such as by computing cosine similarities of prot-prrcessed MAE representations
($ and is combined with PCA whitening on control dataset which acts as a form of weak supervision ($ 5 to
further improve selectivity of the extracted features ($ 6
Second, given that MAEs are not supervised with text or any other labels or annotations, do we learn interesting concepts at all? We find that the DL approaches extract a wide variety of semantically meaningful features. Some relatively simple, like light leaks or dead cells, while others capture abologically-meaningguu concepts (Figures [] & [
Finally, we ask how do we evaluate the biological content of the sparse features? We achieve such interpretations by directly involving domain experts to assess feature quality and biological relevance (Figures IP & 6 and we intro- duce carefully curated linear probing benchmark ($ [85] $ 6.2 to show that ICFL extracts features of a compara- ble interpretability level to biology-informed hand-crafted features.
of cell images (Kraus et al. 2024 Celik et al. 2024] Lazar jet al., 2024; Kenyon-Deea et al. 2025). However, these deep learning- based methods provide limited insights about the morphological changes arising from the perturbations: we can tell whether two perturbations are similar (or dis- similar) via cosine similarity, but we cannot tell why (or the ways in which) they are different. That is, we collapse the relationships between multidimensional representations down single score.
In summary, we show how dictionary learning can be used to extract features from a large scale masked autoencoder trained on microscopy images (Kraus et al.. 2024). We find features correlated with meaningful concepts, such as individual cell types or genetic perturbations. Moreover, via linear probing, we show that the learned features preserve significant amounts of biologically meaningful information.
If we are to use dictionary learning in this context, need to address three questions. First, how do we extract sparse features from MAEs? While it is possible to learn some fea- tures using TopK SAES (Makhzani & Frey/ 2013 Gao et al. 2024). we find that we achieve both better reconstruction and more selective features using an efficient variant of the matching pursuit algorithm Maalatt& & Zhang 1993). Our new proposed algorithm, Iterative Codebook Feature Leam- ing (ICFL, Figure D -C, naturally avoids "dead" features
2. Related work
Dictionary Learning The dictionary learning (DL) prob- lem became popular in the 1990's (Mallat & Zhang [993
2

---

Towards scientific discovery with dictionary learning
3. Background
w/o PCA whitening
w/ PCA whitening
ICFL TopK
55
341
The superposition hypothesis Let E IR denote a repre- sentation obtained from a transformer model; as an example, x may be the (aggregated) embedding of token (resp. to- kens) after the first half of the transformer layers. Bricken et al. (203)) hypothesize that (i) such token representations x E R are linear combinations of concepts; (ii) the number of available concepts M significantly exceed the dimension of the representation d; and (iii) each token representation is the sum of a sparse set of concepts. These desiderata are satisfied by the following model that widely studied in compressed sensing and dictionary learning:
7640
8026
Table 1. The number of "dead features" (out of 8192) activated less than a fraction of 10- during the last 1000 training steps, for both TopK and ICFL with and without PCA whitening ($
Olshausen & Field 1996) and has since been extensively studied in the literature (e.g.. (Aharon et al. 2006; Donoho ict al.,
2001:
Spielman et al., 2012)). Recently, (Bricken
ett al.
2023 proposed to use DL to extract features from
internal representations of LLMs, with several follow-up up
works (e.g., (Rajamanoharan al., 2024b a Gao et al. 2024)), )) including different modalities such as vision (e.g.. (Gorton 2024| Rao et al., 2.24). These works build upon the assumption that large-scale transformer models encode "concepts" as linear directions (as reviewed in (Jiang et al. (0244), commonly referred to as the "linear representation hypothesis". Raaeendran et al. (2024) provides the theoret- ical conditions under which "concepts" can be provably recovered from data. Finally, similar approaches have been previously employed directly on raw data, such as for cor- recting unwanted variation in gene expression data (Jacob) et al. 2016) or for detecting and counting objects by learn- ing templates holographic images (Yellin et al., 2017).
M
x R Wz =
Σ ===1
2m W
where 1 < d (1)
where W' IR is latent dictionary matrix and R is sparse latent vector. In this paper, we will refer to the columns Wi, as "feature directions" and as "features".
Feature learning using Topk SAES. Given a set of rep- resentations { learning both W and { dic- tionary learning or sparse coding problem (Olshausen & Field 1997), with long history of works proposing effi- cient algorithms with provable guarantees (Aharon et al., 2006 Arora et al. 2014 2015). In the context of mechanis- tic interpretability, the dominant choice for learning these parameters are two-layer sparse autoencoders (SAFs). In this paper, we compare to the state-of-the-art method called TopK SAE, originally proposed by (Makhzani & Frey 2013)
Causal representation learning and explainability The disentanglement and causal representation literature (CRL) share the goal of learning high-level, interpretable concepts from data (Bengio al. 2013 Kulkarni al., 2015 Hig-
and recently studied by (Gao et al. 2024). Following their notation, the model for tokens i E N] is:
gins et al. 2017 Chen et al.
2016 Eastwood & Williams
Two key differences with the dictionary learning approach are: (i) disentanglemenn/CRR. methods consider low-dimensional representations to cap- ture the factors of variation in data, whereas overcomplete dictionary learning seeks a higher-dimensional representa- tion to capture a large set of sparsely- firing concepts; and (ii) disentanglement(CR1. methods aim to be inherently inter- pretable, whereas this paper considers a post-hoc approach to interpret pre-trained models. Related work on post-hoc explainability also learns "concept vectors" in neural net- work internal states (Kim et al. 2018 Ghorbani et al. 2019); a key difference that these methods use llasled data, whereas this paper uses an unsupervised approach to discover concepts. Another line of feature-visualization works aim to interpret internal states/neurons by finding the data points (or gradient-optimized inputs) that lead to maximal activation (Mordvintsev et al., 2015; Olah et al. 2017: Borowski et al. 2021). In contrast to these local data point centric approaches, dictionary learning methods seek global parameters (dictionaries) to explain models.
2018 Schölkopf al. 2021)
x; = Wd + bpres with 2; = Topee(eenc.r; bpre);
where TopK(.) is an operator that sets all but the K largest elements to zero. The parameters (Waee, Wenes Opre are
Algorithm Batched-OMP
Parameters Wdee; bprei model representation x;; # sparse features L. per iteration and # iterations J
Input:
2: Initialize (1)) x bprc
3: for t 1= TO J do
4: Select top L columns of Waec which maximize (Wdee,m, ((())
5: Solve 2(() = argmin, (rrgming|Ixx(0)) - Waec= 212 with a non-
zero only for L selected columns
Update (t+1)) == zt (tecz(t)
end for
8: Output: Sparse features = Σi=
3

---

Towards scientific discovery with dictionary learning
learned by minimizing the reconstruction loss:
Repeating this process, the final output a is taken to be 2 Σi= z(0) Consequently, 2 is most A: Jp--sssee. The key advantage over SAEs is that early iterations subtract dominant feature directions from 2, allowing the algorithm in later iterations to select a broader set of feature directions that are not as correlated with the main features in z.
N
L(Waecs dees Uprc) ==
Σ
|1a - Till2, where
(2)
T. = Wane TopK (Wone - Upro) + Upre-
A problem with the above optimization is that some feature directions W dec.m are barely used: that is, we have inactive features 2im = 0 for almost all E [NT] This is called the "dead feature" phenomenon. To reduce the amount of dead features, (Gao et al.) 2024) introduce an additional recon- struction error term containing only these feature directions to encourage their usage in the reconstruction (see Table T
Learning the feature directions using gradient descent Given the features 2, we then update the decoder parame- ters { does ). We do so via batched gradient descent minimizing the reconstruction loss:
LICFL(Wdes. b) = Σ ||a Wdec - Dpre 112-
(4)
As 2; is fixed in this gradient step, the algorithm does not need to propagate gradients through = in order to learn an encoder; instead the "encoder" is simply batched-OMP which will use any feature direction Waee.m that aligns with some Ij, and hence as long as some I; aligns with Waec.m that Wdes. will receive gradient updates. We find that this procedure is much more robust against the emergence of "dead" features during training, as shown in Table 1
Orthogonal Matching Pursuit The OMP algorithm (Mal- lat& & Zhang 1993) is widely used in signal processing appli- cations. It assumes that the dictionary, W, is given, and after initializing r²(0) = x it iteratively finds a k-sparse vector å by repeating the following steps:
1. Find the column in W that solves Wy arg max; WW; | and append it to the matrix W of selected columns.
In practice, we initialize Wdec using the standard pytorch initialization for a linear transform and leverage random resets to ensure that the columns of Wdee are not too corre- lated (high correlation is sometimes referred to as feature collapse). Specifically, after every 100 stochastic gradient descent steps, we take every pair of columns of Wee that have cosine-similarity above 0.9 and randomly initialize one of the pairs with vector selected uniformly at random from the hypersphere. One could think of the random resets as playing a role analogous to the projection step in pro- jected gradient descent, in order to enforce the constraint that cosim (W,, W3) < 1 1 d. However, unlike projected gradient, we are not deterministically projecting onto the boundary of the constraint set; instead, the random reset ensures that we are at some (random) point in the interior of the constraint set (with high probability). This trades off the compute costs and algorithmic complexity of a true projection step for a cheap alternative that works well in practice. Before running Algorithm [ we always center the representations x by subtracting the average representa- tion of unperturbed samples from the control distribution, such that the origin represents the unperturbed state. Before ICFL., we normalize the representations to have unit norm.
2. Solve 24)) = arg min, (1x(0) - W
3. Update (((++)) == sl1) W,,
and terminate when either (t++1) is sufficiently small (i.c., the algorithm has converged) or 20)) is k-sparse for some pre-specified k.
4. Iterative codebook feature learning (ICFL)
We can interpret the OMP algorithm as playing the same role as the encoder in TopK SAES in that the algorithm provides a mapping from a representation, x, to z such that is (at most) &- sparsee. But standard OMP is not useful for dictionary learning because (i) the dictionary is pre- specified and fixed for the duration of the algorithm and (ii) its sequential operation is slow for large k. With two simple modifications we address these issues in an algorithm we call Iterative Codebook Feature Learning (ICFL), which computes the features 2 using a batched variant of MP. During training, we then update the feature directions, Waecs using gradient descent.
4.1. PCA whitening using control dataset
Computing the features = using batched-OMP Given the current matrix of feature directions Waec, we first select a batch of the L. columns most aligned with (1) = xr Then, we proceed as before, learning the features 21)) that best reconstruct x se (", using only these columns (i.e. z(1) is L-sparse). Next, to obtain 2(2), we repeat this step, but replace x with the residual (2) = x (1).
As dictionary learners seek to minimize the Euclidean dis- tance between the model representations x and their recon- structions f = Wdecz, the learned features are naturally biased towards capturing the dominant directions in the data (i.c., those that explain the most variance). In biological applications, these directions will typically correspond bio-
4

---

Towards scientific discovery with dictionary learning
logical variation that is present in all experiments (e.g. mor- phological changes induced by the cell cycle), and hence are not the most useful concepts to learn. To avoid this variation dominating the reconstruction loss, we use dataset of con- trol samples as a form of weak supervision, downweighting dominant directions in this control dataset as we know they do not correspond to the biological perturbations of interest. In particular, let X matrix of representations of n samples from the distribution of images of unperturbed/ / control cells, P(X Let T X - W (X- μ)) denote the linear map the maps the control cells to distribution with zero mean and unit covariance, where W is PCA matrix. By applying the linear tranformation, T, to the peturbed cells sampled from, P (for all perturbations 6,) before normaliza- tion, we effectively scale up the differences between pertur- bation representations. For our multi-cell data, unperturbed HUVEC-cell images act as our control dataset. Note that similar PCA whitening on control dataset has been used to improve the quality of the learned multi-cell image rep- resentations (Kraus et al. 2024 Kenyon-Dess Dean et al. 2025).
the tokens (representations)) are 1024 and 1664 for MAB-L and MAE-G, respectively. All the visualizations used Cell Painting microscopy images from the RxRx1 (Sypetkowski et al. 2023) and RxRx3 (Fay et al. 2023 datasets.
We extract the tokens from layers 16 (MAE-L) and 33 (MAE- G), respectively. The motivation for using intermediate instead of final layers is that these tokens are more likely to capture abstract high level concepts that are used by the model internally to solve the self self-supervised learning task (Alkin et al. 2024). We selected this layer by finding the maximized linear probing performance on the functional group task from the original embeddings ($
Preserving linear probing signals To investigate whether the features found by sparse DL retain important informa- tion from the original representation, we define five different classification tasks (Table 2 For each classification task, we use a separate (potentially overlapping) dataset and split into train and test data to distinguish the labels across:
5. Experimental setup
(1) 23 different cell types which are almost perfectly dis- tinguishable via linear classification.
Data source and foundation model evaluated our DL approach on two large-scale masked autoencoders trained on cellular microscopy Cell Painting image data using 256 x 256 x 6 pixel crops as input and a patch size of 8 xx 8, following the same procedures as described in Kraus ct al. (2024); Kenyon-Dean et al. (2025). These models were trained on data from multiple cell types that were perturbed with CRISPR gene knockout perturbations. Both models used the architecture hyperparameters from Kraus et al. (ee2yo/Kenyon-DD Dean et al. (2025), with the smaller of the two using the ViT-L/8 configuration. while the larger model used the ViT-G/8 configuration. We refer to these models as MAE-L and MAE-G, respectively. We obtain a single token per input crop by aggregating all patch tokens (excluding the class token). For both the residual stream and the at- tention output (after the out projerojection,, the dimension d of
(2) 272 different experiment batches. Even in controlled conditions, subtle changes in experimental conditions induce strong batch effects, i.e. changes in experi- mental outcomes due to experiment-specific variations unrelated to the perturbation that is being tested.
(3) 1138 siRNA perturbations from the RxRx1 dataset (Sypetkowski et al.) [023), where the single-gene ex- pression is partially (or completely) silenced using short interfering (si- (NNA,,,,, which targets the gene mRNA for destruction via the RNA interference path- way (Fie al. 1998). As the extent of siRNA knock- downs is hard to quantify and prone to significant but consistent off-target effects, we also evaluated:
(4) 5 single- CRISPR perturbation knockouts which induce strong and consistent morphological profiles across cell types, known as "perturbation signal bench- marks" (Celik et al. 2024). Unlike the SiRNA ap- proach, CRISPR (Jinck et al. 2012) cuts the gene DNA directly, which induces a mutation and represses the gene function. As perturbations to similar genes often result in similar visual effect, also established:
Threshold
Task
Classes
Samples
BTA
0.5
0.2
Cell
23
110,971
97.2%
73
455
type
Batch
272
80,000
87.8%
77
siRNA
1,138
81,224
51.6%
0
141
(5) 39 functional gene groups composed of CRISPR single-gene knockouts categorized by phenotypic rela- tionships between the genes, including major protein complexes, as well as metabolic and signaling path- ways. Each gene group targets similar or related cellu- lar process, which results in inducing morphologically similar changes in the cells (Celik et al., 2024).
CRISPR
5
79,555
94.6%
0
2
Group
39
57,863
32.1%
0
37
Table Five tasks ( classified by linear probes trained on well- level aggregated representations from the residual stream from intermediate layer from MAE-G G (left) and counted for features with average selectivity above threshold for at least one label (right).
5

---

Towards scientific discovery with dictionary learning
0.8
CP ICFL
ICRL
0.7-
CP
average selectivity
0.40
0.6
ICFL selectivity
selectivity
0.5-
0.3
0.35
commetation
0.4-
0.3-
0.2
0.30
0.2-
0.25
0.1-
0.0-
0
15 ordered genes
10
20
25
6
200
400
600
800
1000
0.2
0.3
0.4
0.5
0.6
0.7
sparsity
CP selectivity
Figure 2. Comparison of feature average selectivity scores from CellProfiler (CP) and ICFL for subset of Task (3). (A) Max avg selcctivity scores for cach label in descending order, (B) averaged across labels at different thresholds for CP and sparsity levels for ICFL, as function of the average number of non-zero values. (C) Correlation of max avg selectivity scores for each label between CP and ICFL.
To remove the impact of spurious correlations between per- turbations and batch effects on the test accuracy, we always use mutually exclusive experiments for test and train data, except for Task (2), where the goal is to predict the ex- periment. Except for Task (1), all classification tasks use HUVEC cells. The representations are always well-level ag- gregated, i.e. take the mean over the total of 1,024 (8 x pixel regions) individual 022-diimesio tokens from all 36 256 x 256 non-edge crops from an 2, 048 x 2, 048 pixel image of a given well (Celik et al. 2024). Due to heavy class imbalance (particularly for Task (1)), we train our linear probes using logistic regression on class-balances cross-entropy loss and report balanced test accuracy (BTA).
the MAE. For each feature, we then compute two selectivity scores: the average, or avg selectivity score, which is the % of times that the feature is active given that label occurs minus the % of times the feature is active given any other label. As a stronger notion of correlation, we also use the maximum, or max selectivity score, that subtracts the maximum % for any other label. The selectivity score has been originally proposed in the context of neuroscience (Hubel& & Wiesel 1968) and has also been used by (Madan) et al.
2022 to measure the "monosemanticity" of neurons.
In Table 2 we present the number of features cxhibiting an average selectivity greater than given threshold for at least one label, across all five classification tasks. We observe that dominant concepts, such as cell types, batch effects, and siRNA perturbations that induce strong morphological changes, comprise a substantial portion of features display- ing high selectivity. Moreover, for labels from the functional gene groups in Task (5), we identify more than 100 features with selectivity scores of at least 0.1.
Training the DL models By default, we always choose a sparsity of K = 100 for Topk SAFS and J = 20, T, = 5 (resulting in a max sparsity of 100) for ICFL as described in $ and use a total of M 8192 features. Unless otherwise specified, we always apply the PCA whitening described in 4.1 and use representations from the residual stream. We train the models using 40M tokens (one token per image crop) with a batch size of 8192 for 300k iterations. Our learning rate is 5 x 10-5 for all experiments. Similar to Gao et al. (2024), we observed that changing the learning rate has a limited impact on the outcome. We present an ablation for the learning rate in Appendix C
Separation along feature directions The selectivity score analysis showed that activation patterns of the sparse features can be strongly correlated with genetic perturba- tions. We now consider the feature directions, and show that certain feature directions can distinguish specific genetic perturbations. Specifically, in Figure 3 in each histogram we plot the cosine similarity of feature direction with both (i) representations from a specific siRNA perturbation (blue) and (ii) representations from all other perturbations (orange). The plot shows that certain feature directions effectively sep-
6. Experimental results
In this section, we present our results using features ex- tracted from ICFL combined with PCA whitening (if not further specified), and comparing with Top-K SAES 6.3.
ALG3
DERL2 (F3144)
SPATA2
EUFER Miee
THED2
6.1. Features are correlated with biological concepts
(F7647)
(5118))
Initially, we investigate how strongly correlated the features are with labels from the classification tasks (Table 2)
-000
C0C
5CC
-0.00
C11
020
Selectivity of features for biological concepts For each dataset associated with a classification task, we extract from every image a feature vector using the center crop as input to
Figure 3. Cosine similarity histograms for selected pairs of rep- resentations from Task (3) and features directions, if' associated perturbation is applied (blue) vs. any other perturbation (orange).
6

---

Towards scientific discovery with dictionary learning
(A)
(1)
(2)
(3)
(4)
(5)
(B)
(C)
1.00
similarity
ICFL
0.50
0.7-
cosine
0.6
TopK
0.25
0.4
baseline
0.5-
PCA whitening
0.00
w/
w/
w/
w/o
wi
WO
w
w/o
w/
wo
50
75
100
125
150
sparsity
sparsity
Figure 4. Comparison of ICFL with TopK SAEs. (A) Balanced test accuracy (BTA) of linear probes trained on the original representation
(solid line) and reconstructions from ICFL and TopK SAEs with and without PCA whitening for five tasks from $ (B) Balanced test
accuracy (BTA) as a function of the sparsity (dashed line is the original representation x) for classification Task (5). (C) Cosine similarity of reconstruution and original representations as a function of sparsity for tokens from hold-out validation dataset.
arate the two groups. and as such these feature directions capture important biological information.
ure 2 we illustrate the correlation between the best average selectivity scores from the CP and ICFL features for each label. We found strong correlation between the two sets of features (Pearson coefficient of 0. 0.71), suggesting that ICFL was capable of identifying features that captured patterns similar to those detected by CP.
Linear probing for biological signals Finally, as mea- sure for how much bbooogicalll-relevame information is lost when extracting sparse features, we compare the accu- racy of linear probes on the representations, x, with linear probes on the reconstructions from ICFL, f = Wdec Al- most the entire signal is preserved for simple concepts such as cell types, batch effects and perturbations with strong morphological changes (1,2.4) (Figure AA). For the difficult tasks of distinguishing between many genetic perturbations (3,5), a substantial amount of the linear signal is preserved.
6.3. Comparison of ICFL with TopK SAE
Finally, we found that ICFL, when combined with PCA whitening, retained more biological signal than TopK SAEs. We first evaluated the reconstruction of the respective unsu- pervised DL algorithms by comparing the cosine similarity between the original and reconstructed representations. Fig- ure ure shows that the reconstruction quality of ICFL was much higher than Topk SAE for the same sparsity con- straints when using PCA whitening. We provide further ablations in Appendix
6.2. Comparison with features from CellProfiler
We compared the average selectivity scores of features from ICFL with those from set of 964 handcrafted features gen- erated by CellProfiler (CP) (Carpenter el al., 2006). These features were designed by domain experts and are widely used for microscopy image analysis. In this experiment we consider a subset of images from the public RxRx -dttsstt (Sypetkowski al.) (023). CP extracts features for each cell; for a multi-cell image, we calculate an image-level CP feature by averaging the CP features from each cell in the image. The CP features are not sparse; to make them comparable with the sparse ICFL features, we thresholded the CP features at a and 1 - a quantiles, with a chosen such that the average number of non-zeros was 100. A CP feature was considered be "activated" when its value, under perturbation conditions, exceeded these quantiles.
To evaluate biologically meaningful information retention, Figure 4A compares the linear probing accuracies for the two algorithms. Both TopK SAE and ICFL features yielded a similar linear probing accuracy. while we could see a clear drop if no PCA whitening was used during rre-processing Figure 4 shows that while we found that increasing the number of non-zeros improved the average balanced test accuracy, the effect was not strong as PCA whitening.
Finally, in Figure & we show the selectivity scores for both ICFL features and Topk SAES. We saw that ICFL features consistently achieved higher selectivity scores than TopK SAE features. Moreover, especially for cell types, we ob- served a high max selectivity across almost all cell types, while for more complex features we observed a more modest selectivity score of more than RJ 0.1 across all labels.
In Figure 2 we plotted the highest average selectivity score for each genetic perturbation, a subset of Task (3). We found that the features extracted by ICFL almost matched the selectivity scores of the handcrafted, human-designed features. Additionally, in Figure [BB shows the average score across all labels a function of various thresholding levels for the CP features. We again observed that fea- tures performed comparably to CP features. Interestingly, CP features peaked at high levels of non-zeros (== 300), leaving future work to assess whether this peak selectivity could be matched using deep learning-base based approaches which used significantly fewer non-zero elements. In Fig-
7. Interpretability analysis of cell imaging data
In this section, we illustrate non-trivial patterns captured by selected features, exemplifying how domain experts can interpret and validate biological concepts learned by DL. To understand what parts of each image are most "aligned" with a particular feature, we generate heatmaps for each image by calculating the cosine similarity of the feature direction with each of the 8 x image patch token representations.
7

---

Towards scientific discovery with dictionary learning
spocfit gose
Inverse focus Monitoring of the lack of channel-specific signal is plausible explanation in other images where the tokens are not always co-localized with regions occu- pied by cells. Several genes appear to follow an "inverse" trend, including PIK-1, which enables cell cycle progres- sion through mitosis, and TMED-2, which helps to regulate intracellular protein transport. While both of these gene per- turbations render the cells in characteristic affected state (small, clumped cells struggling divide vs. large, spread out, actively dividing cells), their most aligned tokens cor- respond to areas not covered by cells, confirmed by highly negative correlations across all channels Frigure5D-E)
750-3
TMFD-2
tolem bastur
ortppaiie Lesiue
orppoadd chaiadd
DNSA
-0.00
actin
40.64
40.6
JINA
-0.00
.0.77
Figure 5. Channel-specifie features visualization from selected single-gene perturbations. Token heatmaps (top) are plotted above the composite 6-channel images (middle) and channel-specific staining images (bottom) of selected subcellular compartments, along with the channel- Pearson correlation coefficients,
7.2. Interpretability analysis on single-cell level
We continue by focusing on single feature which demon- strated a clear biological relationship. This feature is strongly correlated with gene knockouts from the adherens junctions pathway, a label from the functional gene pertur- bation group from Task (5) ($.). The adherens junctions are simple to visually interpret as their function is to connect cell membranes to cytoskeletal elements and form cell-cell adhesions; they can be thought of as "glue proteins" that stick cells together to maintain tissue integrity.
7.1. Channel-specifie perturbation signal
We begin by examining the extent to which we can recover channel-specific signal associated with single-gene pertur- bations. For this exercise, we queried specific single gene perturbations: (i) OPA-1, which contributes to the mainte- nance of correct shape of mitochondria, (ii) ALG-3, which aids in the modification of proteins and lipids in the en- doplasmic reticulum (ER) after synthesis, and (iii) TSC-2, which contributes to the control of the cell size (Figure 5)
The images most correlated with this feature direction re- flect this disrupted cellular morphology; the images mostly comprise of small, bright and isolated cells (Figure (Figure[oA-B) 6oA-)). Note that CRISPR genetic perturbation affect most but not all cells, leaving some cells in unper- turbed, control- state. In the case of adherens junctions, the perturbation-affected cells would appear round, bright and isolated, while the cells which seem to have escaped the perturbation would appear flat, large and wide-spread, attempting to establish connections with their neighbors, which we successfully confirm (Figure 6 white arrows).
not perfect and
OPA-1 The mitochondrial channel shows that most corre- lated tokens are overlaid with distant regions where enlarged mitochondria present (Figure 5AA white arrows). Qual- itative channel examination highlights this delicate detail which is not obvious from the composite images, confirming that our approach identifies thannel-specifii level of detail.
Strikingly, token-level heatmaps appear to have recog- nized the difference between these two cell populations, as the tokens corresponding to control-like cells are not aligned with the feature direction of the image (Figureloc-D, dark regions). These "dark" areas consistently belong to cells whose actin meshwork extensively protrudes away from the cell center, representative of control-like cells. It appears that these dark regions are misaligned with the concept di- rection because these cells do not actively contribute to the representation of the perturbation; instead, they represent noisy signal which is present across almost every image. On the other hand, the regions most correlated with the feature direction belong to areas surrounding the perturbed, com- pact cells and appear to form ring-like pattern around them. As the perturbation results in smaller and rounder cells, this suggests that the concept corresponds to the expected but missing actin (Figure 6B-D, white circles).
ALG-3 The most aligned tokens appear specific to regions of endoplasmic reticulum (ER) and RNA with which AL.G 3 co-localizes. This dense image suggests that our token heatmap is prevalently focused on ER-specific information (Figure 5B., which is consistent with ALG-3 function in aiding the attachment of sugar-like groups to proteins.
TSC-2 We examine the plasma membrane- and Golgi apparatus-specifie channel to relate perturbed cell size con- trol to the token alignment. We confirm that this channel correlates most strongly with the queried concept direc- tion, but this time in negative direction. As the plasma membrane- and, hence, cytoplasmic area are the most extensive from the cell center, the mostly aligned tokens ap- pear to focus specifically on regions which are not covered by the cell membrane Figure 5C). Although this behavior is harder to interpret, it is suggestive of that the salient feature for these perturbations is the lack of cell density in a well.
Next, we aimed to quantitatively confirm whether the token-
8

---

Towards scientific discovery with dictionary learning
level heatmaps enable a consistent separation between the control-like cells and correctly perturbated cells. To do so, we selected 5 cell images from the adherens junctions pathway (further shown in Figure [O), in which a domain expert first manually counted the cells, classified each cell instance as (un-) -un-)pperurbed, and manually segmented them. The corresponding token-level heatmaps were then exam- ined to annotate cell instances by whether they align with the overall feature direction ("bright') or not ('dark"). ).
(A) ocspooiie
(B)
specifie chattee
(C)
token
olgumess cell llurity
acrumbetion
algrreeet beatmap
CP Hee
413
Out of 31 cells which were scored as control-like by the hu- man annotator, the token level heatmap areas corresponding to these single cells were "dark" in 27 instances (Table 3 The darker cell areas correspond to tokens with lower align- ment to the general concept direction of the image, and hence are indicative of lower importance of these areas in the overall image for the functional group classification task. Overall, we report that the token heatmaps are capable of recalling 87% of the manually labelled cells by an expert annotator (Table 3) hinting at the potential of our approach to separate cell subpopulations in a supervision free manner.
Figure 6. Interpretation of single-cell identities from cell images through (A) composite images and (B) their actin-specific chan- nel, which strongly correlates with a feature from the adherens junctions gene group. (C) Segmented single cells classified based morphology and interactions with neighbors (white arrows), Token-level aligament as (D) heatmaps highlighting cells with ring-like boundaries (white circles) and (E) category histograms with compared distributions (means as solid lines, onc-sided Mann Whimey U test; ns, no significant difference; ****, p < 0.0001).
To remove the human subjectiveness of what constitutes a 'dark' vs. "bright' cell in the token heatmap, we continued to evaluate this cell identity scoring approach in unbiased yet quantitative way. To this end, we generated pixel-level segmentation masks of 3 categories: background, perturbed cells and unaffected cells (Figure 6 GC). The histograms of the relative token alignment per cell identity confirm a sta- tistically significant difference between the control-like and perturbed single-cell instances, which are equally informa- tive of the perturbation the background (Figure E).. As typical control images are more dense (as in Figure [D. feature #2 left), the exposure of image background is an im- portant characteristic of the adherens junctions phenotype.
Tools capable of accurately and effectively disentangling these mixtures with single-cell precision in an unsupervised manner therefore remain scarce. Our token-level analysis presents a potential avenue for enabling modern approaches to scientific discovery, leveraging methods from mechanistic interpretability uncover previously unknown biological concepts at the resolution of individual cells.
8. Conclusion
In this paper, we have explored the extent to which dic- tionary learning (DL) can be used to extract biologically meaningful concepts from microscopy foundation models. The results are encouraging: with the right approach, we were able to extract sparse features associated with distinct and
In biological systems, complex mixtures of cell morpholo- gies are frequently observed: however, identifying which ex- act cell instance belongs to which population would require deep expertise and time-consuming manual labeling at scale.
morphological traits. That said, these sparse features are clearly incomplete: their lin- ear probing performance significantly drops on tasks that involve more subtle changes in cell morphology. It is not clear what extent this limitation stems from our current DL techniques, the scale of our models, or whether these more subtle changes are simply not represented linearly in embedding space. Nonetheless, it is evident that the choice of the DL. algorithm matters to extract meaningful features.
Image
Manual Labeling
ICFL Recall
Cells:
Total
Perturbed
Control
'Control'
[%]
A
28
22
6
6
100.0
34
26
8
6
75.0
20
17
3
100.0
We also proposed a new DL algorithm, Iterative Codebook Feature Learning (ICFL), and the use of PCA whitening on a control dataset as a form of weak supervision for the feature extraction. In our experiments, we found that both ICFL and PCA significantly improve the selectivity of ex- tracted features, compared to TopK sparse autoencoders. We hope that future work further explores the use of DL and other mechanistic interpretability techniques for scientific discovery, and the use of ICFL for other modalities like text.
D
23
16
7
6
85.7
E
19
12
7
6
85.7
Total:
124
93
31
27
87.1
Table Tokens align based on single-cell identities in multi-cell images. 'Control' recall is calculated as percentage of control-like ('dark') cells in heatmaps over all human-labeled cells. Note that "perturbation" recall is 100% in all images from Figure 10

---

Towards scientific discovery with dictionary learning
Impact Statement
Celik, S., Hütter, J.-C., Melo Carlos, S., Lazar, N. H., Mo- han, R., Tillinghast, C., Biancalani, T., Fay, M. M., Earn- shaw, B. A., and Haque, I. S. Building, benchmark- ing, and exploring perturbative maps of transcriptional and morphological data. PLOS Computational Biol- ogy, 20110)x 2024. doi: 10.u33/iouurnal..bbbi 1012463.
This is an exploratory paper in the field of AI-driven sci- entific discovery. This paper primarily motivates research on the use of mechanistic interpretability for scientific dis- covery. The findings support the observation that deep net- works are able to learn interesting semantically meaningful features. If future iterations of this work could reliably eluci- date the mode of action of perturbations on cell, that would enable a dramatic improvement of our understanding of cell biology; but to get there, we need more reliable mechanis- tic interpretability methods of disentangling representations into biologically meaningful features.
Chen, X., Duan, Y., Houthooft, R., Schulman, J., Sutskever, I., and Abbeel, P. InfoGAN: Interpretable representation learning by information maximizing generative adversar- ial nets, In Advances in Neural Information Processing Systems, volume 29, pp. 2180-2188 2016.
Darcet, T., Oquab, M., Mairal, J., and Bojanowski, P. Vision transformers need registers. arXiv preprint arXiv:2309. 16588, 2023.
References
Aharon, M., Elad, M., and Bruckstein, A. K-svd: An algo- rithm for designing overcomplete dictionaries for sparse representation. IEEE Transactions on signal processing, 54(11):431 1-422,, 2006.
Donoho, D. L.,, Huo, X., et al. Uncertainty principles and ideal atomic decomposition. JEEE transactions on infor- mation theory, 4777):244--2622, 2001.
Alkin, B., Miklautz, L., Hochreiter, S., and Brandstetter, J. Mim-refiner: A contrastive learning boost from in- termediate pre-trained representations. arXiv preprint arXiv:2402.1009 2024.
Eastwood, C. and Williams, C. K. I. A framework for the quantitative evaluation of disentangled representations. In International Conference on Learning Representations, 2018.
Arora, S., Ge, R., and Moitra, A. New algorithms for learning incoherent and overcomplete dictionaries. In Conference on Learning Theory, pp. 779-806. PMLR, 2014.
Elhage, N., Hume, T. Olsson, C.,, Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds. Z., Lasenby, R., Drain, D., Chen, C., et al. Toy models of superposition. arXiv preprint arXiv:2209.16652 2022.
Arora, S., Ge, Ma, T., and Moitra, A. Simple, efficient, and neural algorithms for sparse coding. In Conference on learning theory, pp. 113-149, PMLR, 2015.
Fay, M. M., Kraus, O., Vietors, M., Arumugam, L., Vuggu- mudi, K., Urbanik, J., Hansen, K., Celik, S., Cernek, N., Jagannathan, G., et al. Rxrx3: Phenomics map of biology. bioRxiv, pp. 2023-02, 2023.
Bengio, Y., Courville, A., and Vincent, P. Representation learning: A review and new perspectives. IEEE Transac- tions on Pattern Analysis and Machine Intelligence, 35 (8):1798 18289 2013.
Ferrando, J., Sarti, G., Bisazza, A., and Costa-jusså, M. R. A primer the inner workings of transformer-based lan- guage models. arXiv preprint arXiv:2405.00200 2024.
Borowski, J., Zimmermann, R. S., Schepers, J., Geirhos, R., Wallis, T. S. A., Bethge, M., and Brendel, W. Exem- plary natural images explain (cnn) activations better than state-of-the-art feature visualization. In International Conference on Learning Representations, 2021.
Fire, A., Xu, S., Montgomery, M. K., Kostas, S. A., Driver, S. E., and Mello, C. C. Potent and specific genetic inter- ference by double-stranded ma in caenorhabditis elegans. Nature, 39 (66666)) (806699)800-11... 1998. doi: 10. 1031038/3588..
Gao, L., Tour, T. D., Tillman, H., Goh, G., Troll, R., Radford, A., Sutskever, I., Leike, J., and Wu, J. Scal- ing and evaluating sparse autoencoders. arXiv preprint arXiv:2406.00 04090, 2024.
Bricken, T., Templeton, A., Batson, J., Chen, B., Jermyn, A., Conerly. T., Turner, N., Anil, C., Denison, C., Askell, A., et al. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2, 2023.
Ghorbani, A., Wexler, J., Zou, J. Y., and Kim, B. Towards automatic concept- based explanations. In Neural Infor- mation Processing Systems, 2019.
Carpenter, A. E., Jones, T. R., Lamprecht, M. R., Clarke, C., Kang, I. H., Friman, O., Guertin, D. A., Chang, J. H., Lindquist, R. A., Moffat, J., et al. Cellprofiler: image analysis software for identifying and quantifying cell phenotypes. Genome biology, 7:1-11, 2006.
Gorton, L.
The missing curve detectors of inceptionvl: Applying sparse autoencoders to inceptionv! I early vision. arXiv preprint arXiv:2406. 03662, 2024.
10

---

Towards scientific discovery with dictionary learning
Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A. (3VAE: Learning basic visual concepts with a constrained varia- tional framework. In International Conference on Learn- ing Representations, 2017.
Lazar, N. H., Celik, S., Chen, L., Fay, M. M., Irish, J. C., Jensen, J., Tillinghast, C. A., Urbanik, J., Bone, W. P., Gibson, C. C., and Haque, 1. S. High-resolution genome- wide mapping of fhromosome-arm-- duced by crispr-cas9 editing. Nature Genetics, 56(7): 1482-1493, 2024, doi: 10. 17838/44588-022-011111
truncations in-
Huang, P.-Y., Sharma, V., Xu, H., Ryali, C., Fan, H., Li, Y., Li, S.-W., Ghosh, G., Malik, J., and Feichtenhofer, C. MAViL: Masked Audio-Videe Learners, arXiv, 2022. doi: D0.48550/axii.2222..555
Madan, S., Henry. T., Dozier, J.. Ho, H.. Bhandari, N., Sasaki, T., Durand, F., Pfister, H., and Boix, X. When and how convolutional neural networks generalize to out- of-distribution eeregorr-viewpoing combinations. Nature Machine Intelligence, 4(2): (2):146-153, 2022.
Hubel, D. H. and Wiesel, T. N. Receptive fields and fune- tional architecture of monkey striate cortex. The Journal of physiology, 195(1):215-243, 1968.
Makhzani, A. and Frey, B. K-sparse autoencoders. arXiv 2013.
preprint
Jacob, L., Gagnon-Bartsch J. A., and Speed, T. P. Cor- recting gene expression data when neither the unwanted variation nor the factor of interest are observed. Biostatis- tics, 17(1):16 28, 2016.
Mallat, S. G. and Zhang, Z. Matching pursuits with time- frequency dictionaries. IEEE Transactions on signal processing, 41(12):3397-3415 1993.
Mikolov, T., Yih, W.-L, and and Zweig, G. Linguistic regularities in continuous space word representations. In Proceedings of the 2013 Conference of the North American chapter of the association for computational linguistics: Human language technologies, pp. 746-751, 2013.
Jiang, Y., Rajendran, G., Ravikumar, P., Aragam, B., and Veitch, V. On the origins of linear representations in large language models. arXiv preprint arXiv:2403.03867 2024.
Jinek, M., Chylinski, K., Fonfara, I., Hauer, M., Doudna, J. A., and Charpentier, E. A programmable dual- ma-guidee dna endonuclease in adaptive bacterial im- munity. Science, 337(6096):81682)) 2012. doi: 10. 122589..
Mordvintsev, A., Olah, C., and Tyka, M. Inceptionism: Going deeper into neural networks. Google research blog, 20(14):5, 2015.
1126/science. 1225829.
Nanda, N., Lee, A., and Wattenberg. M. Emergent linear rep- resentations in world models of self- supe-supe sequence models. arXiv preprint arXiv:2309.00944 2023.
Kenyon-Dean, K., Wang, Z. J., Urbanik, J., Donhauser, K., Hartford, J., Saberian, S., Sahin, N., Bendidi, I,, Celik, S., Fay, M., Rodríguez Vera, J. S., Haque, I. S., and Kraus, O. Vitally consistent: Scaling biological representation learning for cell microscopy. In Proceedings of the 42nd International Conference on Machine Learning (ICML), 2025.
Olah, C.,, Mordvintsev, A., and Schubert, L. Feature visual- ization. Distill, 2(11):e7, 2017.
Olshausen, B. A. and Field, D. Emergence of simple-cell receptive field properties by learning sparse code for natural images. Nature, 38 (6583))607-609) 1996.
Kim, B., Wattenberg, M., Gilmer, J., Cai, C., Wexler, J., and Viegas, F. Interpretability beyond feature attribution: Quantitative testing with concept activation vectors (tcav). In International Conference on Machine Learning, 2018.
Olshausen, B. A. and Field, D. J. Sparse coding with an overcomplete basis set: A strategy employed by v1? Vi- sion research, (1(23):3311-3325 1997.
Park, K., Choe, Y. J., and Veitch, V. The linear represen- tation hypothesis and the geometry of large language models. arXiv preprint 1rXiv:2311.03655 2023.
Kraus, O., Kenyon-Dean, K., Saberian, S., Fallah, M., McLean, P., Leung. J., Sharma, V., Khan, A., Balakr- ishnan, J., Celik, S., Beaini, D., Sypetkowski, M., Cheng, C., Morse, K., Makes, M., Mabey, B., and Earnshaw, B. Masked autoencoders for microscopy are scalable learn- ers of cellular biology. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11757-11768, 2024.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G.,
Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PMLR, 2021.
Kulkarni, T. D., Whitney, W. F., Kohli, P., and Tenenbaum, Deep convolutional inverse graphics network. In Ad- vances in Neural Information Processing Systems, vol- ume 28, pp. 2539 2539,2547, 2015.
Rajamanoharan, S., Conmy, A., Smith, L., Lieberum, T. Varma, V., Kramár, J., Shah, R., and Nanda, N. Improving dictionary learning with gated sparse autoencoders. arXiv preprint orXiv:2404. 16014, 2024a.

---

Towards scientific discovery with dictionary learning
Rajamanoharan, S., Lieberum, T., Sonnerat, N., Conmy, A., Varma, V., Kramár, J., and Nanda, N. Jumping ahead: Improving reconstruction fidelity with jumprelu sparse autoencoders, arXiv preprint arXiv:2407. 14435, 2024b.
Rajendran, G., Buchholz, S., Aragam, B., Schölkopf, B., and Ravikumar, P. Learning interpretable concepts: Unifying causal representation learning and foundation models. arXiv preprint arXiv:2402.092.55 2024.
Rao, S., Mahajan, S., Böhle, M., and Schiele, B. Discover-then-name: Task-agnostie concept bottlenecks via automated concept discovery. arXiv preprint arXiv:2407. /4499, 2024.
Schölkopf. B., Locatello, Bauer, S., Ke, N. R., Kalch- brenner, N., Goyal, A., and Bengio, Y. Toward causal representation learning. Proceedings of the IEEE, 109(5):
612-634, 2021.
Spielman, D. A., Wang, H., and Wright, Exact recovery of sparsely-used dictionaries. In Conference on Learn- ing Theory, pp. 37-1. JMLR Workshop and Conference Proceedings, 2012.
Sypetkowski, M., Rezanejad, M., Saberian, S., Kraus, O., Urbanik, J.. Taylor, J., Mabey, B., Victors, M., Yosinski, J., Sereshkeh, A. R., Haque, I., and Earnshaw, B. Rxrx1: A dataset for evaluating experimental batch correction methods. In Proceedings of the IEEE CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp- 4285 4294, 2023.
Tamkin, A., Taufeeque, M., and Goodman, N. D. Codebook features: Sparse and discrete interpretability for neural networks. arXiv preprint arXiv:2310. 17230, 2023.
Templeton, A. Scaling monosemanticity: Extracting in- terpretable features from claude 3 sonnet. Anthropic, 2024.
Yellin, Haeffele, D., and Vidal, Blood cell detec- tion and counting in holographic lens-free imaging by convolutional sparse dictionary learning and coding. In 2017 IEEE 14th International Symposium on Biomedical Imaging (ISBI 2017), pp. 650-653. IEEE, 2017.
Zhai, X., Kolesnikov, A., Houlsby, N., and Beyer, L. Scaling vision transformers. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 12104-12113, 2022.
12

---

Towards scientific discovery with dictionary learning
null space (1152)
row space (512)
random proj. (512)
relelvv difference
0.95
0.90
0.85
o1gha
Syer
bysr 0yer
0.80
0.75
15
30
layer
(a) reconstructions
(b) probing accuracies
Figure 7. Decoding of tokens from intermediate transformer layers. (A) Sample reconstructed images when decoding from different intermediate layers. (B) The relative linear probing accuracy when using the component from the null space, row space and a random 22-dimensional subspace as component compared to the full component. Both Figures use the MAE-G model.
A. Discussion: Linear concept directions in VIT MAEs
We have shown that DL is powerful approach for finding linear concept directions (features) that are strongly correlated with biological concepts such as cell-types and genetic perturbations. From an interpretability perspective, a question that remains, however, is whether these correlations solely appear due to first order effects of complex non-linear structures used by the model store abstract information, or whether linear directions are actually inherently meaningful to the model? While linear causal interventions offer strong evidence that the latter may indeed at least be partially true for large language models (see e.g... (Ferrando al. 2024) for overview), there exists relatively little evidence for VIT MAEs besides the high linear probing accuracies on e.g., natural and microscopy image classification tasks (Huang al. 2022 2024..
Alkin et al.
In this section, we provide an argument further supporting the hypothesis that MAES may rely on linear concept directions when processing data by analyzing at which point the model are the concepts are the most linearly separable.
Separation into row- and nullspace. We note that standard MAE architectures (Huang et al. 2022) use two different embedding dimensions for the encoder block and the decoder block. Both blocks are connected via an encoder-decoder projection matrix W Rd.xxd with, in our case, d, 1664 (ViT-G model from Zhai et 2022)) and dg 512. This projection matrix gives raise separation of the the tokens into the row-space and null space of W, x Joow + Faull where only the information stored in Frow passed to the decoder. ViTs and more generally transformer models have shown to align the basis across layers, allowing for decoding of tokens from intermediate layers (Alkin et al. 2024). We visualize this behavior in Figure 7a where we show the reconstructions when using the tokens from intermediate layers. Thus, we observe that the row-space component 2/om of tokens from early and intermediate layers ()() already store a reconstruction of the masked image that refined over the layers. Thus the question appears what is the role of the null space component
Tmail which won't be passed to the decoder and thus serves as "register" (in analogy to (Darcet al. 2023))
Component-wis linear probing We analyze in Figure 7b the different components, showing the relative linear probing accuracy of the probing accuracies using the null and row space components, compared to the entire token (dashed line 1) across different layers. As observed, the null space component consistently yields the same probing accuracy as the entire token, while the row space component yields significantly lower accuracy. For comparison, probing accuracy when using a random di-dimensional subspace (the same dimension as the row space), which consistently yields higher accuracy than that obtained from the row space. These findings suggest that biological concepts (i.c., genetic perturbations) are most linearly separable in the component used only for internal processing during the forward pass and not passed to the decoder, and therefore aligns with the hypothesis that the model represents abstract concepts as linear directions accessed by the layers while processing the data (Bricken et al. 2023).
also show the relative
13

---

Towards scientific discovery with dictionary learning
Task (1) cell type
Task (3) SiRNA perturbation
0.4-
Task (5) functional gene group
selectivity
06
selectivits
0.4
ICFL avg score
04-
82
0.2
ICL
Topk avg score
400
Topk
ordered labels
ordered labels
ordered labels
Figure& 8. The highest selectivity among all features for each label. ordered separately for each line starting with maximum
TopK
ICFL
cosine similarity
0.9
cosine similarity
cosine similarity
cosine similarity
0.94
0.6-
0.92
0.90
0.4
0.4
0.85
75
100 sparsity
125
15o
-4 learning rate log 10
0.88
50
----
75
100
125
150
-66
-4
sparsity
learning rate log 10
(a) w/ PCA whitening
(b) w/o PCA whitening
Figure 9. The cosine similarity between the original tokens and the reconstructed tokens for ICFL and TopK-SAE, (A) with PCA whitening and (B) without, as a function of the sparsity (first and third), i.e. # of non-zeros and 1o learning rate (second and fourth).
B. Interpretable features
In this section, we present additional visualizations of crops strongly correlated with selected feature directions. In the spirit of recent works for LLMs (Bricken et al. 2023), we only present a qualitative analysis that aims to highlight non-trivial, complex, and interpretable patterns captured by these features.
For completeness, Figure shows the same crops as Figure I but this time all 6 most correlated and anti-correlated crops. We further present in Figures 12 to 16 additional examples similar to Figure 10 for images strongly correlated with different features. In addition the heat-map and the entire crop, we also plot the patches that are most strongly correlated with the feature. We make two important observations: a) we can see clear interpretable patterns for which patches most strongly correlated with the cells, posing promising area for future research interpreting and validating concept directions found in large foundation models for microscopy image data; b) we see that the most correlated patches are robust to light artifacts, which can be seen best in the last column in Figure 12
C. Ablations
In this section we present ablations type of token, model size, sparsity and learning rate. If not further specified, always use features extracted from ICFL using PCA whitening.
Attention block It is common in the literature to use representations from the MLP output or the attention output (Bricken et al. 2023 Tamkin et al. 2023 Rajamanoharan et al. 2024a). We compare in Table the test balanced accuracy when taking representations from the residual stream and attention output. We observe that both result in similaraccuracies. We make the same observation in Figure 17a and 17b showing an ablation for the linear probes trained on the reconstruction using the same setting as described in Section Moreover, we compare in Figure 18 the selectivity scores in Figure g confirming further that the residual stream and the attention output show a similar behavior. The only exception is TopK for cell types, where the attention outputs result in significantly better selectivity scores, however, still substantially below the ones obtained by ICFL.
Residual stream
97.2%
87.8%
51.6% 52.5%
94.6%
32.1%
Attention output
96.8%
85.8%
94.6%
32.1%
Table 4. The balanced test accuracy (BTA) for representations taken from the residual stream (BTA row from Table and the attention output.
14

---

Towards scientific discovery with dictionary learning
Model size We further investigate the model size, shown in Figures 17a and 17b where compare the linear probes for the MAE-G (referred to Ph2 with 1.9B parameters) with the much smaller model MAE-L (referred Ph1 with 330M parameters). We observe that for simple tasks like classifying cell types, both models yield similar performances. However, we observe consistent improvements on complex classification tasks (3.5), both for the probes trained on the original representations, as well as the reconstructions from ICFL and TopK. This demonstrates that dictionary learning benefits from scaling the model size.
We further plot in Figure 19 the selectivity scores. For ICFL, we consistently observe improvements when increasing the model size, while for TopK SAE, we see significant drop. Interestingly, this drop does not occur for the probing accuracy the reconstructions in Figures and T7bl This suggests that, although capturing meaningful signals in the reconstructions, TopK SAE faces more difficulties in finding "interpretable" features with high selectivity scores from richer representations post-processed using PCA whitening.
Sparsity As a third ablation, we plot in Figure 9 the cosine similarity of the original tokens and the reconstructed token from the DL for both TSpK-SAR and ICFL. We observe that the reconstruction quality of ICFL is much higher than TopK SAE for the same sparsity constraints. This trend persists across all levels of sparsity, The unsupervised reconstruction quality measured by the cosine similarity (or the related fz-error) has been often used as a benchmark for
SABs (Rajamanoharan et al. 2024a Gao al. 2024).
Learning rate As last ablation, also plot in Figure 9 the cosine-similarity for different learning rates. Since PCA whitening leads more dense tokens, expect that decrease in the cosine semila-imillarities, which is also the case when comparing the solid lines (w/o PCA whitening) with the dashed lines (w PCA whitening). Except for TopK-SAE with PCA whitening the reconstruction quality slightly increases with the learning rate (likely due to too few training for small learning rates) and flattens after a learning rate of 5 10-5, which we choose for all experiments this paper. Moreover, we observe that TopK-SAE experiences high instabilities when combined with PCA whitening, which not the case for ICFL.
15

---

Towards scientific discovery with dictionary learning
(A) Composite Image
(B) Actin-Specific Channel
(C) Speccue-Specifii Channel
(D) Cell Segmentation Mask
(E) Token Alignment Heatmap
Image
Image
Image
Image
Image E
Figure 10. Sample visualizations of (A) composite images and their (B) nuclei- and (C) actin-staining channels which strongly correlate with a selected feature from a single functional gene group -oup-ar adherens junctions. Plotted by side are the (D) cell category-speciffe single-cell segmentation masks and (E) token-level hcatmaps of the inner products of the individual tokens with the selected feature direction for 5 out of 8 strongest correlated images per feature direction. Highlighted are correctly perturbed cells (yellow masks) and the cells which most likely remain unperturbed (magenta masks), which the only instances attempting establish cell-cell connections (cyan arrows) as they produce the gene to form functional adherens junctions.
16

---

Towards scientific discovery with dictionary learning
pos cor.
neg cor.
(a) Feature from Figure
pos COr.
neg COr.
(b) Feature 2 from Figure
pos COr.
neg cor.
(c) Feature 3 from Figure 1
Figure H1. Additional visualization of images most and least correlated with selected concept directions. For cach row in Figure 1 we also include the crops that are the most correlated with the feature direction in the opposite direction. More precisely, for each feature we show the most positively (first row) and negatively (second row) correlated crops. For each three features we observe clear concept shift along the feature direction (going from negatively correlated to positively correlated).
17

---

Towards scientific discovery with dictionary learning
Figure 12. Visualization of sample images aligned with feature direction with plausible biological interpretation. This feature appears to be focusing on the endoplasmic reticuli and nuclcoli channel (cyan area) surrounding the nuclcus. These are expanded relative the usual morphology of HUVEC cells.
Figure 13. Visualization of sample images aligned with feature direction with plausible biological interpretation. This feature appears to be firing for cells that are unusually large with spread out actin. Note that the feature focuses on the actin channel (red) surrounding the cell.
18

---

Towards scientific discovery with dictionary learning
Figure 14. Visualization of sample images aligned with feature direction with plausible biological interpretation. This feature
appears to be active for long spindly cells, with the features are most aligned for the long "stretched out" section of the cells.
Figure Visualization of sample images aligned with feature direction with plausible biological interpretation. This feature active for tightly clumped cells. The heatmaps are less clearly interpretable for these images, but appear to be active when neighboring nuclei are touching.
19

---

Towards scientific discovery with dictionary learning
Figure 16. Visualization of sample images aligned with feature direction with plausible biological interpretation. This feature shows a similar behavior to the feature in Figure 13
20

---

Towards scientific discovery with dictionary learning
TopK
ICFL
baseline
(1)
(2)
(3)
(4)
(5)
1.0
0.8
test bal. acc.
0.6
0.4
0.0
Ph2 res. w/-
Ph2 res. w/o
Ph2 attn. w/
Ph2 attn. w/o
Ph1 res, W/
Ph1 res. w/o
Ph1 attn. W/
Ph1 attn. w/o
Ph2 res. w/
Ph2 res. w/o
Ph2 attn. W/
Ph2 attn. w/o
Ph1 res. W/
Ph1 w/o
Ph1 attn. w/
Ph1 attn. w/o
ph2 w/o
Ph2 attn. w/
Ph2 attn. w/o
Ph res. w/
Ph1 res. w/o
Ph1 attn. w/
Phl attn. w/o
Ph2 res. w/
Ph2 res. w/o
Ph2 attn. W/
Ph2 attn. w/o
Ph1 res, W/
Ph1 res. w/o
Ph1 attn. w/
Phl attn. w/o
Phz res. w/
Ph2 attn. w/
Ph2 attn. w/o
Ph1 res. w/
Phl res. w/o
Ph1 attn. w/
Ph1 attn. w/o
Ph2 res.
Ph2
(a) Ablation for the BTA
(1)
(2)
(3)
(4)
(5)
rel. diff. to Ph2 w/
0.6
0.4
Ph2 res. w/o.
Ph2 attn. w/-
Ph2 attn. w/o
Ph1 res. w/-
Ph1 res. w/o-
Ph1 attn. w/-
Ph1 attn. w/o
Ph2 res. w/o.
Ph2 attn. w/
Ph2 attn. w/o-
Ph1 res. w/-
Ph1 res. w/o-
Ph1 attn. w/
Ph1 attn. w/o
Ph2 res. w/o-
Ph2 attn. w/
Ph2 attn. w/o
Ph1 res. w/
Ph1 res. w/o
Ph1 attn. w/-
Ph1 attn. w/o-
Ph2 res. w/o
Ph2 attn. w/-
Ph2 attn. w/o.
Ph1 res. w/-
Ph1 res. w/o
Phl attn. w/
Ph1 attn. w/O
Ph2 w/o
Ph2 attn. w/
Ph2 attn. w/o.
Ph1 res. w/-
Ph1 res. w/o.
Ph1 attn. w/-
Ph1 attn. w/o-
(b) Relative difference MAE-G
Figure 17. Comparison of ICFL with TopK SAEs performance. (A) The BTA of lincar probes trained on the original representation (solid lines) and reconstructions from ICFL features and TopK SAEs representations taken the residual stream and attention output of MAE-G (larger model) and MAE-L (smaller model), as well as with PCA whitening and without. (B) Same as (A) but depicting the relative difference linear probing accuracy compared to MAE-G residual stream using PCA.
21

---

Towards scientific discovery with dictionary learning
residual avg score
residual max score
attention avg score
attention max score
0.5
0.4
selectivity
0.6
Selectivity
selectivity
0.3
0.2
0.2
0.0
0.0
0.0
10
is
20
200
400
600
800
1000
10
30
ordered labels
ordered labels
ordered labels
(a) ICFL cell type
(b) ICFL. siRNA perturbation
ICFL functional gene group
0.35
0.3-
0.30
selectivity
selectivity
selectivity
0.25
0.15
0.2
0.10
0.05
0.0
0.00
10
20
200
400
600
800
1000
10
20
30
ordered labels
ordered labels
ordered labels
(d) TopK cell type
(e) TopK siRNA perturbation
(f) TopK functional gene group
Figure 18. The selectivity scores as in Figure B for ICFL (top row) and TopK (bottom row) when using representations from the residual stream (green) and the attention block (yellow).
22

---

Towards scientific discovery with dictionary learning
PH2 avg score
PH2 max score
PH1 avg score
PH1 max score
0.5
0.4
selectivity
0.6
selectivity
selectivity
0.3
0.4
0.2
0.1
0.0
0.0
0.0
10
i'
20
200
400
600
800
1000
10
30
ordered labels
ordered labels
ordered labels
(a) ICFL cell type
(b) ICFL siRNA pert.
(c) ICFL CRISPR group
0.6
0.5
0.4
0.15
Selectivity
selectivity
selectivity
0.4
0.3
0.3
0.10
0.05
0.2
1000
ordered labels
ordered labels
ordered labels
(d) TopK cell type
(c) TopK siRNA pert.
(f) TopK CRISPR group
Figure 19. The selectivity scores as Figure for for ICFL (top row) and TopK (bottom row) when using representations from the residual stream from MAE-G G (green) and MAE-L (yellow) using PCA whitening.
23


Reviews:
[{"review_id": "ElLd81nNJJ", "reviewer": "Reviewer_Qi4i", "summary": "This paper adapts techniques from the mechanistic interpretability literature to address the problem of discovering what concepts are learned by foundation models trained on microscopy data. Specifically, they develop a new dictionary learning method, which they call ICFL, based orthogonal matching pursuit to extract factors from a masked autoencoder trained on the Cell Painting dataset. A series of auxiliary probe tasks are used to evaluate the quality of the learned atoms. The features are compared with a manual feature extraction pipeline called CellProfiler.  Finally, the usefulness of the learned concepts is illustrated through an extended case study of the a genetic perturbation experiment.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and appreciate their positive feedback on the paper. We were especially pleased to read that the reviewer appreciated our extensive case study. We would like to respond below to a few comments made by the reviewer. \n\n> The only step in the method that I found unusual was how the feature directions are learned. Unless I am misunderstanding something, the objective (4) has a closed form solution for W and b. \n\nWe thank the reviewer for raising this point and would like to clarify: Given z, there is indeed a closed-form solution for W and b. However, there are two key reasons why W should not simply be updated to minimize Equation (4):\nFirst, given that we use more than 40M training tokens (samples), it would be infeasible to optimize (4) directly over all samples. Instead, we rely on mini-batches, making local updates such as gradient descent a natural choice.\n\nSecond, the objective in Equation (4) is convex in W and b only if z is fixed. However, z depends on W since it is computed as the solution to Algorithm 1. This raises the question: what is a practical way to update W? Directly solving Equation (4) may significantly alter W without accounting for its dependency on z. This issue is mitigated when we take only a gradient step with respect to the objective in (4) before updating z again. We refer the reader to [1] for a similar algorithm, where the decoder (or dictionary matrix) is also learned via gradient descent.\n\n> Using gradient descent with random resets struck me as an unecessarily complicated in an otherwise elegant approach.\n\nOne could think of the random resets as playing a role analogous to the projection step in projected gradient descent, in order to enforce the constraint that cosim(w_i, w_j) < 1-\\delta. However, unlike projected gradient, we are not deterministically projecting onto the boundary of the constraint set; instead, the random reset ensures that we are at some (random) point in the interior of the constraint set (with high probability). This trades off the compute costs and algorithmic complexity of a true projection step, for a cheap alternative that we find works well in practice. We will add this discussion to the updated paper.\n\n\n\n\n> These are not necessarily essential references, but the idea of identifying the major factors of variation in a control pool struck me as quite close to the ideas in: [reference to a paper by Zhang et al., 2024, Nat Biotech].\n\nThank you  for highlighting this interesting work and agree that discussing this line of research is a valuable addition to the paper.. We will extend our related work section accordingly.\n\n[1] Arora, Sanjeev, et al. \"Simple, efficient, and neural algorithms for sparse coding.\" Conference on learning theory. PMLR, 2015.", "claims_and_evidence": "I found the claims to be direct and well-supported. For example, it is easy to understand what dead features in SAE are, and the results of Table 1 seem to resolve any doubts that their approach suffers from this problem less.  Similarly, Figure 4 convinced me that the representations were of practical value in perturbation experiments.\n\nI was impressed by the depth of the case study. The identification of factors related to adherens junctions struck me as the sort of result that would be of genuine interest to computational biologists.", "methods_and_evaluation_criteria": "The idea of applying SAE to microscopy images from perturbation experiments is well-justified and the proposed probes are appropriate for evaluation. The method of subtracting variation present in control samples is well-motivated by the experimental design.\n\nThe only step in the method that I found unusual was how the feature directions are learned. Unless I am misunderstanding something, the objective (4) has a closed form solution for W and b. Using gradient descent with random resets struck me as an unecessarily complicated in an otherwise elegant approach.", "theoretical_claims": "There are no theoretical claims.", "experimental_designs_or_analyses": "The ablation experiments in the appendix are helpful for understanding the general properties of IFCL. Despite the comparison with Top-K SAE, some readers might find the lack of more detailed quantitative comparisons (on other datasets or against other baselines) potentially concerning. However, I found the novelty of the application and the qualitative finding from the case study outweighed any concern about precise quantitative benchmarking.", "relation_to_broader_scientific_literature": "This work brings together ideas from two cutting-edge areas -- genetic perturbation experiments and mechanistic interpretability. Better methodology for analyzing perturbation experiments is an area of active research.  Researchers have worked with both manual feature extraction and deep learning models for the associated microscopy images, but combining the interpretability/control of manual feature methods with the richness of self-supervised learning strikes me as potentially very impactful.\n\nIn the interpretability literature, the IFCL algorithm is closely related to Top-K SAE and OMP, and the authors appropriately represent their contribution.  From an interpretability perspective, much of the novelty is in the application domain. Many interpretability techniques have been understood on LLM embeddings, but few (none?) have successfully demonstrated utility on more advanced scientific imaging data.", "essential_references_not_discussed": "These are not necessarily essential references, but the idea of identifying the major factors of variation in a control pool struck me as quite close to the ideas in:\n\nhttps://doi.org/10.1038/s41587-024-02463-1\nhttps://doi.org/10.1093/biostatistics/kxv026", "other_strengths_and_weaknesses": "The paper is written in a direct and precise style.", "other_comments_or_suggestions": "I have no additional comments.", "questions_for_authors": "I have no additional questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DBKxWR9Qj6", "reviewer": "Reviewer_Sady", "summary": "This paper explores the application of dictionary learning (DL) to extract biologically meaningful concepts from large-scale masked autoencoders (MAEs) trained on microscopy images. The authors introduce Iterative Codebook Feature Learning (ICFL), a dictionary learning algorithm adapted from the Matching Pursuit (MP) algorithm, combined with PCA whitening on a control dataset. The approach aims to uncover sparse, interpretable features from cell imaging data that correspond to biological concepts such as cell types and genetic perturbations. The authors validate their method through classification tasks, interpretability analyses, and comparisons with Top-K sparse autoencoders (TopK SAEs) and handcrafted features from CellProfiler (CP). They show that ICFL improves selectivity, reduces “dead features,” and retains biologically relevant signals. The results suggest that dictionary learning can be effectively applied to bioimaging data, potentially enabling new insights in drug discovery and genetic perturbation analysis.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and for raising helpful questions. We would like to respond below to a few comments made by the reviewer. \n\n> Single way of assessing feature interpretation \n\nWhile linear probing is one important way to evaluate the quality of the features, we disagree with the reviewer’s comment that the paper does not use feature clustering analysis. In Figure 4 we show examples for how well clusters are separated along selected feature directions. Moreover, the selectivity scores (as displayed in Figure 2) essentially measures how well features can separate individual genetic perturbations from the average of all perturbations. May we ask the reviewer whether this answers the reviewers concern?\n\n\n> Lack of hyper-parameter sensitivity analysis… sparsity, number of iterations, PCA whitening strength\n\nWe agree that careful ablation studies are an important part of a paper. While the computational cost of these experiments prevents covering the full space of hyperparameters (the dictionary learning methods are trained over 40m tokes), we did include the hyperparameters that you suggested (and many more) in the paper:  \n- We ablate over the sparsity (Figures 3 and 8), model size of the MAE (Figure 17), type of representation (Figures 17 and 18) and learning rate (Figure 8).\n- We did not find the method to be sensitive to the number of iterations as we trained to approximate convergence.\n- There is no notion of “whitening strength” - the representations are whitened such that they have unit covariance and zero mean on the control cells \n\n\n> Lack of ablation experiments for sparse representations \n> ICFL relies on sparse representations to improve interpretability, but does not provide ablation experiments to demonstrate the necessity of sparsity for biological feature extraction.\n\nThe reviewer raises an interesting point. It is per se not clear why one should extract sparse representations instead of dense representations. First, to answer why not dense representations? We recall that the dimension of z (i.e. 8192) is much higher than the dimension of x (i.e. 1664). Without the sparsity constraint, there are infinitely many ways to implement the identity function (that is, Wz = x), and thus the problem is inherently ill-posed. As a result, there is also no reason why one should believe that the learned features should capture biologically meaningful signals. \n\nHowever, why sparse features? The motivation stems from the superposition hypothesis, which is described in Section 3. While this hypothesis clearly does not hold universally, the question remains - to what extent do it hold? With this paper, we provide at least some evidence that a weaker form of a superposition hypothesis may indeed hold. \n\n\n> but does not compare it with broader explanatory methods (e.g., topological data analysis, decomposable representation learning)\n\nWhile we agree that broader comparisons are generally useful; comparing across different classes of interpretability methods is very challenging because each class of methods has a different notion of interpretability. As a result, there are not clean quantitative comparisons across interpretability methods that can be made. That said, we believe that a well-designed user study that evaluated these methods for scientific discovery would be extremely valuable, but this goes well beyond the scope of this paper. \n\n\n> The experiments in this paper are mainly conducted on 100,000-level microscope data, and the lack of computational complexity analysis of million-level bio-image data affects the applicability of this method on large-scale datasets.\n\nWe believe there is a misunderstanding. We use a subset of the data for evaluation when computing the linear probes. However, to train the dictionaries W, we use 40M tokens (that is, 40M crops). Moreover, we train for 300k steps with a batch size of 8192 tokens. \n\n> 1. What is the computational complexity of ICFL on millions of microscope data? Has it been tested for efficiency on large-scale data?\n\nWe did not conduct an extensive run-time analysis given that the cost of compute is negligible compared to the cost of training and fine-tuning the MAEs\n. \n> 2. Can the extracted dictionary features be matched with known biomarkers?\n\nThe selectivity score analysis suggests that the extracted features align well with (some) genetic perturbations. May we ask the reviewer to be more precise what sort of experiment they had in mind?\n\n> 3. Does PCA whitening result in loss of certain biologically relevant information? Are there ablation experiments?\n\nThe contrary is true, PCA whitening helps to increase the biological signal in the extracted features. We would like to refer the reviewer to Figure 3. \n\n> 4. Have alternative methods of feature selection been considered...\n\nNot at this time, as such methods go beyond the scope of this paper, but we do regard such an analysis as an interesting and important future work.", "claims_and_evidence": "1. Lack of ablation experiments for sparse representations -- The paper emphasizes that sparse dictionary learning is essential for interpretability, but does not conduct a sparse vs. non-sparse comparison experiment. The lack of such an ablation experiment affects the credibility of the conclusion.\n2. Lack of hyper-parameter sensitivity analysis -- ICFL relies on several hyper-parameters (e.g., sparsity, number of iterations, PCA whitening strength), but the paper fails to conduct sensitivity analysis on these parameters, which affects the assessment of model robustness.\n3. Single way of assessing feature interpretation - the paper mainly uses Linear Probing to assess whether the features are biologically relevant, but does not use feature clustering analysis (whether biologically meaningful clusters can be formed spontaneously).\nThe use of these additional assessments can improve the persuasiveness of the paper.", "methods_and_evaluation_criteria": "The paper uses PCA whitening + sparse dictionary learning to extract features, but does not compare it with broader explanatory methods (e.g., topological data analysis, decomposable representation learning), which affects the full assessment of its methodological uniqueness.", "theoretical_claims": "ICFL relies on sparse representations to improve interpretability, but does not provide ablation experiments to demonstrate the necessity of sparsity for biological feature extraction.", "experimental_designs_or_analyses": "The experiments in this paper are mainly conducted on 100,000-level microscope data, and the lack of computational complexity analysis of million-level bio-image data affects the applicability of this method on large-scale datasets.", "relation_to_broader_scientific_literature": "None", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "None", "other_comments_or_suggestions": "It is recommended that the ICFL computational complexity analysis be supplemented with operational efficiency tests on large-scale data to enhance the utility of the method.", "questions_for_authors": "1. What is the computational complexity of ICFL on millions of microscope data? Has it been tested for efficiency on large-scale data?\n2. can the extracted dictionary features be matched with known biomarkers?\n3. does PCA whitening result in loss of certain biologically relevant information? Are there ablation experiments?\n4. have alternative methods of feature selection been considered (e.g., feature extraction based on topological data analysis)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oadTaNBtI3", "reviewer": "Reviewer_e1Nw", "summary": "This paper proposes a new algorithm for dictionary learning, namely Iterative Codebook Feature Learning (ICFL), which can be optionally augmented with the PCA whitening technique. This technique is then applied to interpret features learned in masked autoencoder models trained on microscopy images of cells. The authors present case studies on the extracted features.", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive review. We would like to respond below to a few comments made by the reviewer. \n\n> The experiment lacks comparisons between the proposed methods and other existing approaches to interpretability. It is not clear how much practitioners can gain from using ICFL compared with other methods.\n\nWhile a more extensive benchmark evaluation would indeed be valuable, we would like to stress that the emphasis of this paper lies in the use of dictionary learning / SAEs in the context of biological representation learning. In particular, we aim to explore how ideas developed in mechanistic interpretability for LLMs could potentially be used in the future for scientific discovery. A benchmark comparison therefore goes beyond the scope of this paper. \n\n>  show the generality on more foundation model families\n\nIn this paper, we prioritize an interdisciplinary approach and an in-depth analysis of features extracted from one specific model rather than multiple models. This is a common approach in the existing literature (see, e.g., [1,2]) for two reasons.\n- First, assessing the quality of extracted features is a highly domain-specific task, hence our usage of a SOTA model for cell microscopy data, rather than a generic vision foundation model. An attempt to apply the method to more foundation models, would have necessitated a less detailed study of the features that are learned. \n- Second, we note that the overwhelming majority of prior work [e.g. 1,2] on dictionary learning techniques have also only been applied to a single modality (typically text) with a single model. This approach is common for the reason outlined above (in any domain, careful analysis of the features requires extensive domain expertise), and access to large scale foundation models, the corresponding datasets, and the necessary infrastructure to run these experiments is often limited. \n\n> The presentation of the paper, the algorithm part in particular, looks sloppy to me.\n\nWe thank the reviewer for making us aware of the inconsistencies  in the notation. We switched between W and Wdec to better highlight the similarity between ICFL and TopK SAEs, however, there are indeed some minor inconsistencies due to smaller edits before submitting the paper, which we will resolve. Moreover, Wdec is initialized using the standard pytorch initialization for a linear transform (nn.Linear). We will mention this in the paper. \n\n> In line 163 (right), are \n\nIt should indeed be W_i^T, thank you for noticing this typo. \n\n> Are the MAE models open-sourced? What are the training configurations of the models? Those details should have been documented in the paper.\n\nUnfortunately, we are not able to open source the MAEs at the time of writing this response, but they are based on the standard implementation of the MAEs described in [3]. Moreover, the training details for the MAEs are as described in [4]. \n\n\n[1] Gao, Leo, et al. \"Scaling and evaluating sparse autoencoders.\" arXiv preprint arXiv:2406.04093 (2024).\n\n[2] Bricken, Trenton, et al. \"Towards monosemanticity: Decomposing language models with dictionary learning.\" Transformer Circuits Thread 2 (2023).\n\n[3] He, Kaiming, et al. \"Masked autoencoders are scalable vision learners.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.\n\n[4] Kraus, Oren, et al. \"Masked autoencoders for microscopy are scalable learners of cellular biology.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.", "claims_and_evidence": "The main claims are:\n- The proposed approach successfully retrieves biologically meaningful concepts, such as cell type and genetic perturbations.\n- The proposed approach can help better understand morphological changes in cells induced by genetic perturbations.\n\nThe first claim is supported via experimental results. But some comparisons are not convincing. For example, in Figure 2B, the results of ICFL are so sparse that I do not really know how to compare them against the baseline (CP) performances.\n\nThe second claim is based on a specific case study, which presents some evidence but does not adequately convince the readers on the generalizability of the approach.", "methods_and_evaluation_criteria": "The evaluation criteria seem sound to me.", "theoretical_claims": "No theorem in this paper.", "experimental_designs_or_analyses": "The experiment **lacks comparisons** between the proposed methods and other existing approaches to interpretability. It is not clear how much practitioners can gain from using ICFL compared with other methods.", "relation_to_broader_scientific_literature": "From the interpretability perspective, this paper introduces a new algorithm based on dictionary learning.", "essential_references_not_discussed": "I'm not aware of any missing related works.", "other_strengths_and_weaknesses": "## Strengths\n- Making scientific discovery with machine learning is an intriguing research topic.\n\n## Weaknesses\n- The presentation of the paper, the algorithm part in particular, looks sloppy to me.\n  - Abuse of notations. Section 4 uses $W$ and $W_{\\text{dec}}$ interchangeably without any additional explanations. I assume that $W_{\\text{dec}}$ follows the notations in SAE, but this is inappropriate since there is no $W_{\\text{enc}}$. In Equation (4), the use of $W$ and $W_{\\text{dec}}$ is not even consistent in a single equation.\n  - Missing details. The \"computing the features $z$ using batched-OMP\" paragraph assumes $W_{\\text{dec}}$ is given. However, no where in the main paper is the initialization of $W_{\\text{dec}}$ specified.\n  - Inconsistent formats. The OMP algorithm is described using bullet points. The batched-OMP is described in (very informal) texts and pseudo code. Making those descriptions consistent would help clarify the modifications in the new method.\n\n- The proposed Iterative Codebook Feature Learning (ICFL) seems a simple extension of existing methods to me.  The PCA whitening is also a known technique. Thus, **the technical contribution of the paper is limited**. While I value the simplicity in the design of ML algorithms, the authors need to present very solid comparisons against existing methods, show the generality on more foundation model families, and discuss insights on why the simple method is favorable.", "other_comments_or_suggestions": "Other minor presentation issues:\n\n- Please check the use of `\\citep` and `citet`. There are many mistakes in the paper, e.g., lines 133 (right), 134 (right), 151 (right), 156 (right), 249 (left), etc.\n\n- Line 125 (right): $\\\\{x\\\\}_{i=1}^N$ $\\to$ $\\\\{x\\_i\\\\}^N\\_{i=1} $.\n\n- Line 184 (left): MP $\\to$ OMP.\n\n- Line 205 (right): correspond $\\to$ correspond to.", "questions_for_authors": "- In line 163 (right), are $W_i$ and $x^{(t)}$ both column vectors? How is $|W_ix^{(t)}|$ defined?\n\n- Are the MAE models open-sourced? What are the training configurations of the models? Those details should have been documented in the paper.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The review are split on this paper. On one side, all reviewers agree that the paper propose and innovative application of interpretability for microscopy data, with some interesting results. But one reviewer points out that the methodology is not new and it is unclear that it would lead to good results on another task, limiting the scope of the paper for the broad audience of ICML.

I find that the paper is well written, with a clear exposition of both the methods and the application. The method is also simple and provides some insights on what the large model encodes in their intermediate layers, and the results on the microscopy are intriguing, even though it is not completely clear if they could really lead to scientific discovery (the results right now demonstrate that the model learns some known properties, but it is much harder to find novel ones from the learned representations).
On the other side, it is not completely clear what the novelty is in ICFL (it looks like classic DL with OMP sparse code and gradient update of the dictionary), and it is unclear whether it would lead to better results on other modalities. Moreover, no discussion is given about the runtime.

Overall, I lean mostly toward accepting the paper due to its novelty for the application, which shows a good example of scientific investigation in the results of Large models, with an original contribution. I would still recommend adding a discussion about runtime for both methods and fixing figure 2B to have more points for ICFL. 


**Extra remarks:**

- Using OMP for dictionary learning is far from new, and the random reset of unused atoms is a classical step in DL literature. OMP is typically one of the algorithm available in [scikit learn](https://scikit-learn.org/stable/modules/generated/sklearn.decomposition.DictionaryLearning.html#dictionarylearning) and the code implement a random reset of unused atoms.
- Side note on the first answer to **Qi4i**: the second point is not really valid. There is a large body of dictionary learning literature that uses alternate minimization, which alternates between $z$ and $W, b$, and there is a clear justification for this. But I agree that the first point is valid and requires learning with an iterative method, which can be one step of gradient descent. 
- l.170 - It is confusing that $k$ is both the number of features selected at each step and the predefined sparsity. I would advise changing it in the algorithm to select $l$ features?
- A reference about dictionary learning with cell imaging that could be interesting to include: Yellin, F., Haeffele, B. D. & Vidal, R. Blood cell detection and counting in holographic lens-free imaging by convolutional sparse dictionary learning and coding. in IEEE International Symposium on Biomedical Imaging (ISBI) (Melbourne, Australia, 2017).
- l.331.c2 - typo: "the by"
- Fig5-title : "actin"->action?
Decision: Accept (poster)