Flow Stochastic Segmentation Networks

2025-07-24

AI Summary

To address the challenge of efficiently modeling high-dimensional inter-pixel covariance in generative image segmentation, this paper proposes Flow-SSN—a novel generative segmentation model supporting both discrete-time autoregressive and continuous-time flow paradigms. Unlike conventional low-rank covariance parameterizations, Flow-SSN directly models arbitrary high-rank pixel-wise covariance via differentiable flow transformations, eliminating the need for pre-specified rank constraints or explicit storage of distribution parameters—thereby enhancing representational capacity and sampling efficiency. The model integrates probabilistic density transformation, differentiable sampling, and domain-specific medical image priors. Evaluated on multiple medical segmentation benchmarks, Flow-SSN achieves state-of-the-art performance, significantly outperforming existing diffusion-based methods in both accuracy and inference speed.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Deep Generative Models & Autoencoders
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs

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.

Problem

• Overcoming low-rank limitations in generative segmentation models
• Enhancing efficiency in sampling from flow-based segmentation models
• Achieving state-of-the-art results in medical imaging benchmarks

Innovation

• Generative segmentation model with flow variants
• Estimates high-rank covariances without assumptions
• Efficient sampling with expressive prior learning