An MRF-UNet Product of Experts for Image Segmentation
This addresses overfitting in medical image segmentation for neuroimaging applications, offering an incremental improvement by integrating classic MRF priors into CNNs.
The paper tackles the problem of CNNs overfitting on out-of-distribution data in semantic segmentation by fusing a UNet with an MRF using a product of distributions and iterative mean-field approximation. The MRF-UNet improves generalization to out-of-distribution samples and reduces parameters while preserving high accuracy on 3D neuroimaging data.
While convolutional neural networks (CNNs) trained by back-propagation have seen unprecedented success at semantic segmentation tasks, they are known to struggle on out-of-distribution data. Markov random fields (MRFs) on the other hand, encode simpler distributions over labels that, although less flexible than UNets, are less prone to over-fitting. In this paper, we propose to fuse both strategies by computing the product of distributions of a UNet and an MRF. As this product is intractable, we solve for an approximate distribution using an iterative mean-field approach. The resulting MRF-UNet is trained jointly by back-propagation. Compared to other works using conditional random fields (CRFs), the MRF has no dependency on the imaging data, which should allow for less over-fitting. We show on 3D neuroimaging data that this novel network improves generalisation to out-of-distribution samples. Furthermore, it allows the overall number of parameters to be reduced while preserving high accuracy. These results suggest that a classic MRF smoothness prior can allow for less over-fitting when principally integrated into a CNN model. Our implementation is available at https://github.com/balbasty/nitorch.