RepMode: Learning to Re-parameterize Diverse Experts for Subcellular Structure Prediction

TL;DR

RepMode is a novel deep learning framework that dynamically re-parameters a mixture of experts to accurately predict 3D subcellular structures from transmitted-light images, addressing multi-scale and partial labeling challenges.

Contribution

It introduces a task-aware re-parameterization method with a mixture-of-diverse-experts architecture for improved subcellular structure prediction.

Findings

Achieves state-of-the-art performance in SSP

Effectively handles multi-scale variability

Addresses partial labeling issues

Abstract

In biological research, fluorescence staining is a key technique to reveal the locations and morphology of subcellular structures. However, it is slow, expensive, and harmful to cells. In this paper, we model it as a deep learning task termed subcellular structure prediction (SSP), aiming to predict the 3D fluorescent images of multiple subcellular structures from a 3D transmitted-light image. Unfortunately, due to the limitations of current biotechnology, each image is partially labeled in SSP. Besides, naturally, subcellular structures vary considerably in size, which causes the multi-scale issue of SSP. To overcome these challenges, we propose Re-parameterizing Mixture-of-Diverse-Experts (RepMode), a network that dynamically organizes its parameters with task-aware priors to handle specified single-label prediction tasks. In RepMode, the Mixture-of-Diverse-Experts (MoDE) block is designed to learn the generalized parameters for all tasks, and gating re-parameterization (GatRep) is performed to generate the specialized parameters for each task, by which RepMode can maintain a compact practical topology exactly like a plain network, and meanwhile achieves a powerful theoretical topology. Comprehensive experiments show that RepMode can achieve state-of-the-art overall performance in SSP.