UCell: rethinking generalizability and scaling of bio-medical vision models

TL;DR

UCell introduces a small, parameter-efficient model for biomedical image segmentation that achieves performance comparable to larger models without extensive pretraining, emphasizing scalability and generalizability.

Contribution

The paper presents UCell, a tiny recursive model that matches larger models' performance in biomedical segmentation, trained from scratch without large-scale pretraining.

Findings

UCell matches larger models in segmentation accuracy across multiple benchmarks.

UCell can be trained solely on microscopy data without pretraining on natural images.

UCell demonstrates strong adaptability through one-shot and few-shot fine-tuning.

Abstract

The modern deep learning field is a scale-centric one. Larger models have been shown to consistently perform better than smaller models of similar architecture. In many sub-domains of biomedical research, however, the model scaling is bottlenecked by the amount of available training data, and the high cost associated with generating and validating additional high quality data. Despite the practical hurdle, the majority of the ongoing research still focuses on building bigger foundation models, whereas the alternative of improving the ability of small models has been under-explored. Here we experiment with building models with 10-30M parameters, tiny by modern standards, to perform the single-cell segmentation task. An important design choice is the incorporation of a recursive structure into the model's forward computation graph, leading to a more parameter-efficient architecture. We found that for the single-cell segmentation, on multiple benchmarks, our small model, UCell, matches the performance of models 10-20 times its size, and with a similar generalizability to unseen out-of-domain data. More importantly, we found that ucell can be trained from scratch using only a set of microscopy imaging data, without relying on massive pretraining on natural images, and therefore decouples the model building from any external commercial interests. Finally, we examined and confirmed the adaptability of ucell by performing a wide range of one-shot and few-shot fine tuning experiments on a diverse set of small datasets. Implementation is available at https://github.com/jiyuuchc/ucell