Towards scientific discovery with dictionary learning: Extracting biological concepts from microscopy foundation models

We thank the reviewer for their constructive review. We would like to respond below to a few comments made by the reviewer.

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.

While 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.

show the generality on more foundation model families

In 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.

• 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.
• 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.

The presentation of the paper, the algorithm part in particular, looks sloppy to me.

We 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.

In line 163 (right), are

It should indeed be W_i^T, thank you for noticing this typo.

Are the MAE models open-sourced? What are the training configurations of the models? Those details should have been documented in the paper.

Unfortunately, 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].

[1] Gao, Leo, et al. "Scaling and evaluating sparse autoencoders." arXiv preprint arXiv:2406.04093 (2024).

[2] Bricken, Trenton, et al. "Towards monosemanticity: Decomposing language models with dictionary learning." Transformer Circuits Thread 2 (2023).

[3] He, Kaiming, et al. "Masked autoencoders are scalable vision learners." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[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.

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.

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.

We 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): First, 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.

Second, 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.

Using gradient descent with random resets struck me as an unecessarily complicated in an otherwise elegant approach.

One 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.

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].

Thank 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.

[1] Arora, Sanjeev, et al. "Simple, efficient, and neural algorithms for sparse coding." Conference on learning theory. PMLR, 2015.

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.

Single way of assessing feature interpretation

While 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?

Lack of hyper-parameter sensitivity analysis… sparsity, number of iterations, PCA whitening strength

We 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:

• 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).
• We did not find the method to be sensitive to the number of iterations as we trained to approximate convergence.
• There is no notion of “whitening strength” - the representations are whitened such that they have unit covariance and zero mean on the control cells

Lack of ablation experiments for sparse representations ICFL relies on sparse representations to improve interpretability, but does not provide ablation experiments to demonstrate the necessity of sparsity for biological feature extraction.

The 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.

However, 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.

but does not compare it with broader explanatory methods (e.g., topological data analysis, decomposable representation learning)

While 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.

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.

We 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.

1. What is the computational complexity of ICFL on millions of microscope data? Has it been tested for efficiency on large-scale data?

We 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 .

2. Can the extracted dictionary features be matched with known biomarkers?

The 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?

3. Does PCA whitening result in loss of certain biologically relevant information? Are there ablation experiments?

The 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.

4. Have alternative methods of feature selection been considered...

Not 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.