SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Training vs Fine-tuning. FixedChannels and channel-adaptive models can be trained from scratch by initializing weights randomly, or fine-tuned by initializing with ImageNet pre-trained weights. For fine-tuning, the weights of the first layer are replicated to cover more than three channels as needed. FixedChannels-fine-tuned models improve performance about  $22\%$  over randomly initialized models from 0.500 to 0.612 in the performance score. Note some tasks that under-perform after fine-tuning: generalization tasks C_Task3 and C_Task4 in the Cell Painting data (Fig. 5A). This reveals how FixedChannels models overfit to variation specific to one type of image, reducing a model's effectiveness on OOD problems. While fine-tuned FixedChannels models are a higher bar to improve, all channel-adaptive strategies in this regime are able to improve the performance of generalization tasks H_Task3, C_Task2, and C_Task3 (peaks outside the cyan polygon in Fig. 5C). This highlights the benefits of using channel-varying data instead of fine-tuning for fixed subsets.
Most channel-adaptive strategies outperform the FixedChannels baseline when fine-tuning, indicating that generalist models can have better performance. On average, fine-tuning channel-adaptive models improves performance by  $23\%$  relative to training from scratch (Tab. 1). The large difference between training and fine-tuning models indicates that CHAMMI still benefits from pre-training on larger datasets, even if they are not cellular images. CHAMMI is a small dataset with well balanced, and well annotated images useful for fine-tuning and evaluation rather than for pre-training at large scale. Our future work includes the extension and maintenance of CHAMMI with a large-scale non-annotated dataset that is useful for pre-training.
Data augmentation and loss functions. Common data augmentations for RGB images are not necessarily effective for microscopy because of differences in color space. We implemented a set of basic augmentations that have been found to be successful for fluorescence microscopy, including random cropping and flips. In addition, we evaluated the performance of thin-plate-spline (TPS) transformations [54] as a potential strategy to simulate technical artifacts and improve model robustness to noise. We observed that more complex data augmentations can increase performance in all architectures up to  $2\%$  relative to basic augmentations (Tab. 1), highlighting the importance of image manipulation for OOD generalization [55-59].
We also investigated self-supervised learning (SSL) [60-65] for solving the tasks in the CHAMMI benchmark. The results in Fig. 5 were obtained with a supervised loss using class annotations in the CHAMMI dataset. Here, we evaluate SimCLR [60, 66] using the basic set of augmentations plus TPS transformations [54]. We observe that SimCLR alone severely underperforms (SSL FixedChannels performance score is 0.360 vs 0.623 when using supervision). However, in combination with supervised learning, SimCLR improves performance by about  $2\%$  (FixedChannels supervision+SSL performance score is 0.633, Tab. 1). These results reinforce our observation that CHAMMI is a great dataset for fine-tuning and benchmarking, but not necessarily for large-scale pre-training with SSL.
Figure 5: Model comparison on the CHAMMI benchmark. Radar plots have nine axes representing the benchmark tasks. Numbers in color boxes indicate the performance score obtained by the model. A) Comparison of three baseline models. The tasks in the outer circle are colored by the dataset they correspond to, with gray meaning validation tasks, and colors meaning generalization tasks. B,C) Comparison of five, channel-adaptive strategies trained or fine-tuned with the CHAMMI dataset. Number in the circle center indicates in how many tasks the model is better than the baseline.
Table 1: Computational cost and performance score for each method on CHAMMI. The cost is the number of trainable parameters and the average number of forward passes during inference. Performance scores include relative improvement percentages with respect to the previous column. From left to right: Trained, Fine-tuned, Fine-tuned + TPS, Fine-tuned + TPS + SSL.
Model	Computational Cost ↓	Performance Score ↑
Nparams	Cinference	Trained	Fine-tuned +	+TPS +	+SSL +
ChannelReplication	27.82M	4.0X	0.344	-	-	-
FixedChannels	83.47M	1.0X	0.500	0.614 22%	0.623 1%	0.633 2%
Depthwise	27.84M	1.0X	0.517	0.616 19%	0.622 1%	0.608 2%
SliceParam	27.86M	1.0X	0.457	0.592 29%	0.600 1%	0.600 0%
TargetParam	27.84M	1.0X	0.496	0.622 25%	0.628 1%	0.618 2%
TemplateMixing [51]	28.02M	1.0X	0.466	0.615 31%	0.616 0%	0.621 1%
HyperNet [53]	28.29M	1.0X	0.537	0.616 14%	0.630 2%	0.632 0%
Computational cost. Training a channel-adaptive model not only has the potential to improve performance in downstream tasks, but can also be computationally more efficient, reducing the carbon footprint of microscopy image analysis at scale. We estimated the computational cost of models both for training and testing times, which are reported in Tab. 1. The training cost is estimated with the number of parameters of the underlying architecture, assuming that the specific training strategy adds a constant cost to the overall training complexity. The inference cost is estimated with the average number of forward passes required by a model to produce features for a single-cell image in the CHAMMI test sets. Note that ChannelReplication has the lowest training cost, but the highest inference cost, because it needs to process channels individually.
The training cost of channel-adaptive models is very similar to each other, because all are based on the same ConvNext architecture, which is trained once for all channel configurations. The differences are primarily in the arrangement of the first convolutional layer, which may require more or less filters depending on the strategy. Inference cost is also the largely same as all models process each single-cell image once. Both TemplateMixing [51, 52] and HyperNet [53] have a computational overhead due their layer weight generation requirements. However, this increase is negligible when amortized across a batch of images (e.g., less than  $0.2\%$  for a batch size of 64) [52], and needs to only be computed once when processing multiple batches of images with the same set of channels. Thus, the main difference between models is on performance accuracy, which is reported in Tab. 1 as the overall score (performance on the six generalization tasks). Most models maintain the same level of performance or improve upon the FixedChannels baselines.
4.3 Generalization capabilities
One of the main goals of extending models to be channel-adaptive is that training with more data can generalize better to new domains or out-of-distribution data. There are two major generalization problems in microscopy: the first is biological and the second is technical. The biological generalization problem deals with experimental data that has novel biological conditions not seen during training. CHAMMI includes six biological generalization tasks with increasing difficulty, which introduce novel organelles, cell lines, and compounds in the prediction problems (Fig. 3). Generalization to new biological conditions (or domains) remains an open challenge of active research [59, 67, 68]. The technical generalization problem deals with model performance on new datasets with different number of channels, which is a new challenge introduced with CHAMMI.
Biological generalization. We evaluate two training strategies that can be used to improve OOD generalization performance on the biological tasks. First, we consider Stochastic Weight Averaging Densely (SWAD) [69], which achieves generalization by seeking the flat minima of loss landscapes through dense weight averaging. Second, we consider Mutual Information Regularization with Oracle (MIRO) [70], which exploits the generalizability of a large-scale pre-trained model.
First, we found that SWAD improves the performance of FixedChannels on all biological generalization tasks, while MIRO improves performance on half of them (Fig. 6A). The main weakness of MIRO is that it relies on a reference ImageNet pre-trained model as an oracle to learn image representations, which we have found to be suboptimal (Fig. 3A). Next, we compare these results with TargetParam trained in the standard way, and observe that it is able to perform similarly to FixedChannels trained with SWAD or MIRO overall while improving in some tasks (appendix). This highlights the potential of channel-adaptive models to improve generalization performance out-of-the-box. Finally, we trained the TargetParam model with SWAD and MIRO to evaluate the possibility of observing a synergistic effect of training with more data and optimizing the model with OOD techniques. TargetParam slightly improves performance in three tasks out of six using SWAD or MIRO, which supports the idea that channel-adaptive models can improve generalization ability while also benefiting from advances in OOD generalization.
Technical generalization. Here, we evaluate transfer learning performance of models trained with two datasets while testing on a holdout, third dataset that has a different number of channels not seen during training. The main hypothesis is that training a model with varying-channel images can perform better in a new image dataset with a different number of channels when no training data is available. From the five channel-adaptive models evaluated in this work, Depthwise can be extended to a new set of channels without additional training by simply replicating / removing filters. Therefore, we evaluate the Depthwise approach in this transfer learning setting, and use the same approach with ImageNet pre-trained weights as a baseline. Training on varying-channel images indeed improves performance up to  $11\%$ ,  $12\%$ , and  $5\%$  relatively, when the WTC, HPA, and CP datasets are held-out, respectively (Fig. 6B). Note that the baseline model has only seen three-channel RGB images. For comparison, Fig. 6B reports the performance of Depthwise and FixedChannels models trained with images from the evaluated dataset as an estimation of the upper-bound performance for transfer learning. The difference reveals ample room for improvement in methods and training strategies. All the results in this experiment are the average F1-score of the generalization tasks of each dataset (Sec. 3).
Technical variation biases in the microscopy images. Microscopy images, like other biological data, exhibit artifacts due to technical variation (e.g., equipment, microscope, technician, time of day). These artifacts can introduce positive biases that enhance performance in IID data and confounding factors that challenge accurate OOD generalization. The benchmark tasks were designed to include OOD data for improved estimation of model performance. However, technical variation may still play a role in confounding the results, which is also the topic of active concurrent research work.
A relatively small-scale dataset compared to the actual needs of biological research. Given the dataset balance, biological relevance, and known technical variations, our benchmark serves as a limited case study compared to the broader requirements of contemporary biological research. For instance, our tasks involve fewer than 10 classes, while real-world problems often encompass thousands of classes. We hope CHAMMI facilitates the rapid development and evaluation of models
Biological generalization: novel conditions / domains
B.
Technical generalization: leave one dataset out / new channels
Transfer Learning
Full training (upper bound)
Fixed
Channels
Figure 6: Generalization evaluation. A) Performance on the six biological generalization tasks of CHAMMI for two models: FixedChannels (left) and TargetParam (right). Each model is further trained with domain generalization methods SWAD [69] and MIRO [70]. B) Performance of training on two datasets and leaving one dataset out to evaluate technical generalization (avg. F1-score).
that can scale to larger varied-channel datasets that could be curated in the future, in which case the benchmark will still be relevant given its well-standardized set of tasks and metrics.
We presented CHAMMI, a new dataset for investigating a novel task: creating channel-adaptive models for flexible and scalable microscopy image analysis. We curated single-cell images from three publicly available sources, standardized their format, and designed nine biologically relevant downstream tasks to quantify the progress made by channel-adaptive models. We conducted extensive experiments that include baselines for channel-adaptive, non-adaptive, and domain generalization methods sampled from prior work on our task. We find that all channel-adaptive methods are able to improve or match the performance of non-adaptive specialists models at a smaller computational cost. Generalization tasks with OOD data benefit more from channel-adaptive models, and the HyperNet model improved performance substantially with additional data augmentations and the incorporation of self-supervision.
There are many open challenges to realize the full potential of channel-adaptive models in microscopy and beyond. First, exploring new architectures specifically designed for varying-channel images may further improve performance. Second, training and optimization strategies as well as data augmentations that target varying-channel images present many opportunities for innovation. Third, biological and technical generalization may benefit from novel domain adaptation techniques. Finally, creating unannotated, large-scale datasets (e.g., with many more microscopy modalities and diverse image resolutions) for pre-training channel-adaptive models can also complement our well-balanced and well-annotated benchmark. We leave these important research questions for future work, and expect the proposed CHAMMI benchmark to support many of these exciting developments.
Broader Impacts. The acquisition and analysis of biological data have the potential for both positive and negative societal impacts when deployed in applications such as drug development and basic research. Bad actors may use this research to develop biotechnology that harms humans. We integrate openly accessible microscopy images acquired with different experimental conditions and technical formats. By solving the channel-adaptive model problem, this effort has the potential to bring together various microscopy imaging communities. Channel-adaptive models can also impact various other fields, including satellite, spectral, thermal, and ultrasound imaging, among others.
This study was supported, in part, by the Broad Institute Schmidt Fellowship program and by National Science Foundation NSF-DBI awards 2134695 and 2134696. Figure created with BioRender.com.