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

Konstantin Donhauser, Gemma Elyse Moran, Aditya Ravuri, Kian Kenyon-Dean, Kristina Ulicna, Cian Eastwood, Jason Hartford

Keywords: scientific discovery, mechanistic interpretability, microscopy image data, computer vision

TL;DR: We take the first step towards discovering scientifically valuable concepts from large foundation models using tools from mechanistic interpretability.

Abstract:

Dictionary learning (DL) has emerged as a powerful interpretability tool for large language models. By extracting known concepts (e.g., Golden-Gate Bridge) from human-interpretable data (e.g., text), sparse DL can elucidate a model's inner workings. In this work, we ask if DL can also be used to discover *unknown* concepts from less human-interpretable scientific data (e.g., cell images), ultimately enabling modern approaches to scientific discovery. As a first step, we use DL algorithms to study microscopy foundation models trained on multi-cell image data, where little prior knowledge exists regarding which high-level concepts should arise. We show that sparse dictionaries indeed extract biologically-meaningful concepts such as cell type and genetic perturbation type. We also propose a new DL algorithm, Iterative Codebook
Feature Learning~(ICFL) and combine it with a pre-processing step which uses PCA whitening from a control dataset. In our experiments, we demonstrate that both ICFL and PCA improve the selectivity or ``monosemanticity'' of extracted
features compared to TopK sparse autoencoders.

Track: Main track

Submitted Paper: No

Published Paper: No

Submission Number: 55