HD-RDS-UNet: Leveraging Spatial-Temporal Correlation between the Decoder Feature Maps for Lymphoma Segmentation

Lymphoma is cancer originated in the lymphatic system. Clinically, automatic and accurate lymphoma segmentation is critical yet challenging. Recently, UNet-like architectures are widely used for medical image segmentation. The pure UNet-like architectures can model the spatial correlation between the feature maps very well, whereas they discard the critical temporal correlation. Some prior works combine UNet with recurrent neural networks (RNNs) to utilize the spatial and temporal correlation simultaneously. However, it is inconvenient to incorporate some advanced techniques proposed for UNet to RNNs, which hampers their further improvements. In this paper, we propose a recurrent dense siamese decoder architecture, which simulates RNNs and can densely utilize the spatial temporal correlation between the decoder feature maps following a 'UNet' approach. We combine it with a modified hyper dense encoder. Therefore, the proposed model is a UNet with a hyper dense encoder and a recurrent dense siamese decoder (HD-RDS-UNet). To stabilize the training process, we propose a weighted Dice loss with stable gradient and self-adaptive parameters. We perform patient-independent five-fold cross-validation on 3D volumes collected from whole-body PET/CT scans of patients with lymphomas. The experimental results show that the volume-wise average Dice score and sensitivity are 85.58% and 94.63%, respectively. The patient-wise average Dice score and sensitivity are 85.85% and 95.01%, respectively. The different configurations of HD-RDS-UNet consistently show superiority in the performance comparison. Besides, a trained HD-RDS-UNet can be easily pruned, resulting in significantly reduced inference time and memory usage, while keeping very good segmentation performance.