ProtoMTG: Prototypical Multi-Task Learning for the Generation of Multiple Stained Immunohistochemical Images

Multiplex immunohistochemistry (mIHC) has potential to asses the complex tumor microenvironment by simultaneously detecting multiple markers within a single tissue section, however, the acquisition of mIHC images in clinical labs are both time-consuming and costly. Hence, applying machine learning-based virtual staining techniques for rapid generation of different mIHC bio-markers has become a considerable alternative. The existing bio-image based virtual staining models generate the distribution of different markers independently, which has limited explainable ability and overlooks the fact that the exploration of potential interrelationships among these markers can help determine the localization of each individual marker. To address the above issues, we propose an explainable prototypical multi-task learning framework (i.e., ProtoMTG) to simultaneously generate multiple mIHC markers from the DAPI staining of nuclei. Specifically, ProtoMTG involves a multi-task prototype layer that can capture the relationship among different virtual staining tasks by learning the shared and task-specific prototypes. Then, in the proto-attention layer, both task-specific and shared prototypes will be re-weighted and combined to instruct the generation of different mIHC markers. In ProtoMTG, we also design the novel prototypical activation loss and diversity loss to learn better prototype representation for the virtual staining task. To evaluate the performance of our method, we develop three benchmark mIHC datasets from different organs (i.e., colon, liver and stomach) of humans. The experimental results indicate that our method can not only outperform the existing image generation models, but also have good explainable ability for the virtual staining of mIHC markers.