Brudfors, M;
Balbastre, Y;
Ashburner, J;
Rees, G;
Nachev, P;
Ourselin, S;
Cardoso, MJ;
(2021)
An MRF-UNet Product of Experts for Image Segmentation.
In:
Proceedings of Machine Learning Research.
(pp. pp. 48-59).
Proceedings of Machine Learning Research (PMLR): Lübeck, Germany.
Preview |
Text
2104.05495v1.pdf - Accepted Version Download (1MB) | Preview |
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.
Archive Staff Only
 |
View Item |
