Compositional Representation Learning for Brain Tumour Segmentation

TL;DR

This paper introduces vMFNet, a mixed supervision framework that learns robust brain tumour segmentation models using mostly weak labels and minimal pixel-level annotations, leveraging compositional representations and unsupervised learning.

Contribution

The paper proposes a novel approach combining weak supervision, unsupervised learning, and compositional vMF distributions for effective brain tumour segmentation with limited annotations.

Findings

High segmentation accuracy with mostly weak labels

Emergent anatomical structure learning from pathology supervision

Effective use of minimal fully-annotated data

Abstract

For brain tumour segmentation, deep learning models can achieve human expert-level performance given a large amount of data and pixel-level annotations. However, the expensive exercise of obtaining pixel-level annotations for large amounts of data is not always feasible, and performance is often heavily reduced in a low-annotated data regime. To tackle this challenge, we adapt a mixed supervision framework, vMFNet, to learn robust compositional representations using unsupervised learning and weak supervision alongside non-exhaustive pixel-level pathology labels. In particular, we use the BraTS dataset to simulate a collection of 2-point expert pathology annotations indicating the top and bottom slice of the tumour (or tumour sub-regions: peritumoural edema, GD-enhancing tumour, and the necrotic / non-enhancing tumour) in each MRI volume, from which weak image-level labels that indicate the presence or absence of the tumour (or the tumour sub-regions) in the image are constructed. Then, vMFNet models the encoded image features with von-Mises-Fisher (vMF) distributions, via learnable and compositional vMF kernels which capture information about structures in the images. We show that good tumour segmentation performance can be achieved with a large amount of weakly labelled data but only a small amount of fully-annotated data. Interestingly, emergent learning of anatomical structures occurs in the compositional representation even given only supervision relating to pathology (tumour).