Source-linked AI summary
Parallel Multi-Dimensional LSTM, With Application to Fast Biomedical Volumetric Image Segmentation
Marijn F. Stollenga, Wonmin Byeon, Marcus Liwicki, Juergen Schmidhuber
TL;DR
Biomedical volumetric segmentation needs broad pixel context, while prior MD-LSTM variants were difficult to parallelize on GPUs. The paper introduces PyraMiD-LSTM, which reorganizes computations in a pyramidal topology for efficient parallel processing. It achieves state-of-the-art segmentation results in challenging benchmarks, including the best known MRBrainS13 result and competitive EM-ISBI12 performance.
Problem
Prior MD-LSTM methods could capture full spatio-temporal pixel context but were difficult to parallelize on modern GPUs, especially for volumetric biomedical data.
Method
PyraMiD-LSTM rearranges MD-LSTM computation into pyramidal contexts using six directional C-LSTMs whose outputs are combined.
Results
PyraMiD-LSTM achieves state-of-the-art segmentation results in challenging benchmarks, with the best known MRBrainS13 result and competitive EM-ISBI12 performance.
Takeaways & Limitations
Highly parallel PyraMiD-LSTM makes GPU-parallel multidimensional recurrent processing available for challenging biomedical volumetric segmentation tasks.
Abstract
from arXiv · showhide
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).
1 Introduction
LSTM-based recurrent models can aggregate broad image context, but prior MD-LSTM approaches were difficult to parallelize for volumetric biomedical segmentation. PyraMiD-LSTM addresses this computational challenge and achieves strong results on two biomedical segmentation tasks.
- LSTM networks retain information over long time intervals through gated error-carousal units, supporting sequence-processing tasks.Their prior successes include handwriting recognition, large-vocabulary speech recognition, and machine translation.
- MD-LSTM connects hidden units across image grids, allowing each pixel’s representation to recursively gather information from the entire image.This makes two-dimensional MD-LSTM applicable to assigning foreground or background classes during image segmentation.
- 3D-LSTM captures each volumetric pixel’s full context through eight sweeps performed by eight LSTMs along directed volume diagonals.Biomedical volumes include CT, MR, and EM data, which prior methods often process slice by slice.
- Prior MD-LSTM variants were difficult to parallelize because recurrent computations are sequential, particularly for volumetric data.The paper introduces PyraMiD-LSTM as a topology and update strategy designed to improve parallelization, reduce computations, and scale on GPUs.
- PyraMiD-LSTM achieves competitive results on EM-ISBI12 and the best known results on MRBrainS13.The method is evaluated on two challenging biological volumetric-image segmentation tasks.
2 Method
The method reorganizes MD-LSTM computation into pyramidal contexts and plane-wise C-LSTM updates, reducing directional recurrence and improving GPU parallelization for volumetric data.
- Pyramidal Connection Topology: 45°-rotated connections make whole grid rows independently computable, while extra inputs fill the resulting context gaps.The added connections allow more than two context elements to be processed.
- Pyramidal Connection Topology: PyraMiD-LSTM replaces cuboid MD-LSTM contexts with pyramids, reducing the number of 3D directional LSTMs from 8 to 6.For dimension d, MD-LSTM requires 2^d LSTMs, whereas PyraMiD-LSTM requires 2 × d.
- PyraMiD-LSTM: The architecture can stack PyraMiD-LSTM layers before a fully connected output stage using tanh and, finally, softmax class probabilities.Inputs are sampled from random locations with random rotations and flips before entering the directional C-LSTMs.
- PyraMiD-LSTM: In 3D, six C-LSTMs scan planes along the positive and negative directions of the three axes and jointly create each pixel’s full context.Each C-LSTM is a complete LSTM RNN that processes the entire volume in one direction.
- PyraMiD-LSTM: RNN-tailored convolutions propagate information sideways within each volume axis, unlike CNN convolutions that pass information upward between layers.The outputs from all directional C-LSTMs are added to form one PyraMiD-LSTM layer.
3 Experiments
The experiments evaluate PyraMiD-LSTM on EM and MR brain volumetric segmentation using sub-volumes, augmentation, RMSprop with momentum, and organizer-defined metrics. The method outperforms other methods in EM rand error without post-processing and leads the MR brain final ranking.
- Datasets: The experiments use EM and MR brain volumetric image segmentation datasets, with EM membrane labels and MR scans divided into training and test volumes.The EM dataset has one training and one testing stack; the MR dataset has five training and fifteen test volumes.
- Data processing: GPU memory limits require training and testing on randomly extracted sub-volumes, with dataset-specific rotation and flipping augmentation during training.Test-time sub-volumes are stitched using a Gaussian kernel, while rotations and flipping are disabled.
- Training: RMSprop with momentum uses normalized gradients based on a running gradient-variance average, which also helps address vanishing gradients.The training objective is squared loss, and normalized updates are applied to the weights.
- Architecture: The model uses three PyraMiD-LSTM layers with 16, 32, and 64 hidden units, followed by fully connected layers and a class-sized output layer.The same architecture and hyperparameters are used for both datasets.
- EM results: Without post-processing, PyraMiD-LSTM networks outperform other methods in EM rand error and remain competitive in warping and pixel errors.The authors note that post-processing could further improve performance.
- MR brain results: PyraMiD-LSTM leads the MR brain final ranking with a new state-of-the-art result and outperforms other methods for CFS in all metrics.The ranking sums team ranks across nine organizer-computed measures, with ties broken using standard deviation.
4 Conclusion
PyraMiD-LSTM addresses the GPU-parallelization limitations of previous MD-LSTM implementations while retaining their ability to use entire spatio-temporal context. The approach has achieved state-of-the-art segmentation results on challenging benchmarks.
- MD-LSTM recursively incorporates each pixel’s entire spatio-temporal context, unlike CNNs’ typically smaller local receptive fields.
- Figure 5 shows three MR scan inputs, their pre-processed versions, and the resulting segmentation, while omitting one artifact-ridden input.
- Previous MD-LSTM implementations could not exploit the parallelism of modern GPU hardware.
- PyraMiD-LSTM changes this by providing a highly parallel MD-LSTM architecture for segmentation.
- PyraMiD-LSTM has achieved state-of-the-art segmentation results in challenging benchmarks.