An MRF-UNet Product of Experts for Image Segmentation

Mikael Brudfors, Yaël Balbastre, John Ashburner, Geraint Rees, Parashkev Nachev, Sebastien Ourselin, M. Jorge Cardoso

Keywords: CNN, U-Net, MRF, products of experts, image segmentation

TL;DR: We describe a novel approach for combining a UNet with a low parameter, first-order MRF prior.

Abstract: 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.

Registration: I acknowledge that publication of this at MIDL and in the proceedings requires at least one of the authors to register and present the work during the conference.

Authorship: I confirm that I am the author of this work and that it has not been submitted to another publication before.

Paper Type: methodological development

Primary Subject Area: Segmentation

Secondary Subject Area: Transfer Learning and Domain Adaptation

Source Code Url: https://github.com/balbasty/nitorch/blob/3c66ccc5793ac76c51c48a9c063b8638df324e48/nitorch/nn/modules/segmentation.py#L394

Data Set Url: https://my.vanderbilt.edu/masi/workshops/, https://mrbrains18.isi.uu.nl/

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2104.05495/code)