Flow Stochastic Segmentation Networks

TL;DR

Flow Stochastic Segmentation Networks (Flow-SSN) introduce a flexible generative segmentation approach capable of modeling complex pixel covariances efficiently, outperforming previous methods on medical imaging benchmarks.

Contribution

The paper presents Flow-SSNs, a novel generative segmentation model family that overcomes limitations of low-rank parameterizations and improves sampling efficiency.

Findings

Achieved state-of-the-art results on medical imaging benchmarks.

Can estimate high-rank pixel covariances without assumptions.

More efficient sampling than diffusion-based models.

Abstract

We introduce the Flow Stochastic Segmentation Network (Flow-SSN), a generative segmentation model family featuring discrete-time autoregressive and modern continuous-time flow variants. We prove fundamental limitations of the low-rank parameterisation of previous methods and show that Flow-SSNs can estimate arbitrarily high-rank pixel-wise covariances without assuming the rank or storing the distributional parameters. Flow-SSNs are also more efficient to sample from than standard diffusion-based segmentation models, thanks to most of the model capacity being allocated to learning the base distribution of the flow, constituting an expressive prior. We apply Flow-SSNs to challenging medical imaging benchmarks and achieve state-of-the-art results. Code available: https://github.com/biomedia-mira/flow-ssn.