PSR-Nets: Deep neural networks with prior shift regularization for PET/CT based automatic, accurate, and calibrated whole-body lymphoma segmentation

Lymphoma is a type of lymphatic tissue originated cancer. Automatic and accurate lymphoma segmentation is critical for its diagnosis and prognosis yet challenging due to the severely class-imbalanced problem. Generally, deep neural networks trained with class-observation-frequency based re-weighting loss functions are used to address this problem. However, the majority class can be under-weighted by them, due to the existence of data overlap. Besides, they are more mis-calibrated. To resolve these, we propose a neural network with prior-shift regularization (PSR-Net), which comprises a UNet-like backbone with re-weighting loss functions, and a prior-shift regularization (PSR) module including a prior-shift layer (PSL), a regularizer generation layer (RGL), and an expected prediction confidence updating layer (EPCUL). We first propose a trainable expected prediction confidence (EPC) for each class. Periodically, PSL shifts a prior training dataset to a more informative dataset based on EPCs; RGL presents a generalized informative-voxel-aware (GIVA) loss with EPCs and calculates it on the informative dataset for model finetuning in back-propagation; and EPCUL updates EPCs to refresh PSL and RRL in next forward-propagation. PSR-Net is trained in a two- stage manner. The backbone is first trained with re-weighting loss functions, then we reload the best saved model for the backbone and continue to train it with the weighted sum of the re-weighting loss functions, the GIVA regularizer and the L2 loss function of EPCs for regularization fine-tuning. Extensive experiments are performed based on PET/CT volumes with advanced stage lymphomas. Our PSR-Net achieves 95.12% sensitivity and 87.18% Dice coefficient, demonstrating the effectiveness of PSR-Net, when compared to the baselines and the state-of-the-arts.