Neural network parametrized level sets for image segmentation

1. Problem Statement

Image segmentation and classification aim to partition an image into regions with distinct properties. The Chan-Vese model is a state-of-the-art variational method for this task, minimizing an energy functional to find piecewise constant approximations of an image. Traditionally, this is implemented using unparametrized level-set methods, where the level-set function is represented pixel-wise (a vector of values). While powerful, these methods can be computationally expensive and lack structural constraints.

The paper addresses the following challenges:

• Efficiency: Can neural networks serve as efficient parametrizations for level-set functions?
• Transparency: Why do specific neural network architectures (e.g., depth, activation functions) work well for representing geometric shapes (polygons/segments)?
• Formalism: How can the mathematical foundations of the Chan-Vese model be rigorously linked to neural network theory to justify using them for segmentation?

2. Methodology

The authors propose a Parametrized Chan-Vese (PCV) framework where level-set functions are represented by Deep Neural Networks (DNNs) rather than pixel grids.

A. Neural Network Architectures

The paper defines specific network structures to approximate level sets:

• Heaviside Networks: Use the Heaviside step function ($\sigma$) as the activation.
  • One-layer (Shallow) Networks: Represent affine linear combinations of Heaviside functions. These are shown to approximate polygonal regions.
  • Two-layer Networks: Defined as a composition of two layers. The authors prove that two-layer networks are sufficient and most efficient for approximating polyhedral segments and multiphase classification functions.
• Sigmoid Networks: For numerical implementation, the discontinuous Heaviside function is replaced by a smooth Sigmoid function ($\sigma_\varepsilon$). This allows for gradient-based optimization.

B. Mathematical Formulation

1. Parametrization: Instead of optimizing a pixel-wise function $\ell$, the method optimizes the parameters (weights and biases) of a neural network $s(\vec{x})$ such that the zero-level set of $s$ defines the segmentation boundary.
2. Multiphase Representation: For $m$ classes, the image is partitioned into $2^m$ regions using $m$ level-set functions. The piecewise constant approximation $m(\vec{x})$ is constructed as a linear combination of products of these level-set functions.
3. The Functional: The authors define a Parametrized Chan-Vese Functional ($CV_{f, \varepsilon}^{LS}$). This functional includes:
   • Data Fidelity: The $L^2$ distance between the image and the piecewise constant approximation.
   • Regularization: Length of the boundaries (using the Dirac delta approximation) and area terms.
4. Optimization Algorithm:
   • Step 1 (Initialization): A neural network is pre-trained (Algorithm 1) on synthetic data (e.g., circles) to generate a good initial structure for the level-set parameters.
   • Step 2 (Refinement): The network parameters are optimized using Mini-batch Stochastic Gradient Descent (SGD) (Algorithm 2) to minimize the PCV functional directly on the target image.

3. Key Contributions

1. Theoretical Link Between Level Sets and Neural Networks

The paper provides a rigorous mathematical proof that two-layer neural networks are the optimal architecture for approximating the characteristic functions of convex polygons and, by extension, arbitrary regions with finite perimeter.

• Lemma 2.4 & Lemma 3.6: Demonstrates that a customized two-layer Heaviside network can exactly represent the characteristic function of a convex polygon defined by $n_1$ lines.
• Theorem 5.2: Proves that any multiphase level-set function (with regions of finite perimeter) can be approximated in $L^2(\Omega)$ by a two-layer neural network.

2. Efficiency of Two-Layer Networks

The authors show that while one-layer networks can define the boundaries (lines), a two-layer network is required to combine these boundaries into specific regions (polygons) and multiphase classifications. This establishes that depth 2 is sufficient for image segmentation tasks, avoiding the need for deeper, more complex architectures.

3. Approximation Chain and Convergence

The paper establishes a convergence chain:
$$m_\varepsilon \approx t_\varepsilon \approx t \approx m$$
Where:

• $m$: The ideal piecewise constant segmentation.
• $t$: The two-layer Heaviside network approximation.
• $t_\varepsilon$: The two-layer Sigmoid network approximation.
• $m_\varepsilon$: The smoothed multiphase function.
  The authors prove that as the smoothing parameter $\varepsilon \to 0$, the sigmoid network converges to the Heaviside network, and the smoothed functional $\Gamma$-converges to the original Chan-Vese functional under regularity assumptions.

4. Novel Algorithm

The paper introduces a two-stage algorithm (Algorithms 1 & 2) that combines supervised pre-training (to learn a good initialization) with unsupervised variational optimization (to refine the segmentation based on the specific image energy).

4. Numerical Results

The authors tested the method on synthetic images (circles) and compared different initialization strategies and hyperparameters:

• Initialization:
  • Pre-trained Initialization: Using a network trained on similar shapes resulted in cleaner boundaries and faster convergence to the global minimum.
  • Random Initialization: While capable of finding the solution, it produced noisier, more irregular boundaries and was more susceptible to local minima or stochastic noise from mini-batch optimization.
• Smoothing Parameter ($\varepsilon$):
  • Smaller $\varepsilon$ (e.g., 0.5) provided faster convergence and clearer results in early iterations compared to larger $\varepsilon$ (e.g., 1.0), though results converged to similar quality after sufficient iterations.
• Regularization ($\mu$):
  • The parameter controlling boundary length significantly affected the smoothness of the result. $\mu=0.1$ yielded the clearest boundaries in the tested scenarios.
• Optimizer:
  • Stochastic optimizers (AdamW) were found to be faster and less sensitive to initialization than deterministic methods (like L-BFGS), though they introduced slight variability in the final solution.

5. Significance

• Bridging Variational Methods and Deep Learning: The paper successfully unifies the classical variational approach (Chan-Vese) with modern deep learning. It moves beyond "black box" segmentation by using neural networks as explicit, mathematically grounded parametrizations of geometric shapes.
• Computational Efficiency: By representing level sets with a small number of network parameters (weights/biases) rather than a full pixel grid, the method reduces the dimensionality of the optimization problem, potentially leading to faster computation and lower memory usage.
• Interpretability: The work explains why neural networks work for segmentation: they are naturally suited to represent polyhedral regions (the building blocks of complex shapes) via two-layer architectures.
• Practical Application: The proposed algorithm offers a robust framework for image segmentation that leverages the speed of stochastic optimization and the structural priors of neural networks, outperforming random initialization in terms of boundary quality and convergence stability.

In summary, this paper provides a foundational theoretical justification for using two-layer neural networks as parametrized level sets, proving their efficiency in approximating geometric segments and offering a practical, high-performance algorithm for image segmentation and classification.