Parallel Multi-Dimensional LSTM, With Application to Fast Biomedical Volumetric Image Segmentation

TL;DR

This paper introduces PyraMiD-LSTM, a parallelizable multi-dimensional LSTM architecture that improves volumetric image segmentation by efficiently capturing global context, outperforming previous methods on brain MRI datasets.

Contribution

The paper presents a novel pyramidal re-arrangement of MD-LSTM computations, enabling efficient GPU parallelization for 3D image segmentation tasks.

Findings

Achieved state-of-the-art results on MRBrainS13 dataset.

Demonstrated efficient parallelization on GPUs for 3D data.

Outperformed previous MD-LSTM variants in segmentation accuracy.

Abstract

Convolutional Neural Networks (CNNs) can be shifted across 2D images or 3D videos to segment them. They have a fixed input size and typically perceive only small local contexts of the pixels to be classified as foreground or background. In contrast, Multi-Dimensional Recurrent NNs (MD-RNNs) can perceive the entire spatio-temporal context of each pixel in a few sweeps through all pixels, especially when the RNN is a Long Short-Term Memory (LSTM). Despite these theoretical advantages, however, unlike CNNs, previous MD-LSTM variants were hard to parallelize on GPUs. Here we re-arrange the traditional cuboid order of computations in MD-LSTM in pyramidal fashion. The resulting PyraMiD-LSTM is easy to parallelize, especially for 3D data such as stacks of brain slice images. PyraMiD-LSTM achieved best known pixel-wise brain image segmentation results on MRBrainS13 (and competitive results on EM-ISBI12).