Sigma Flows for Image and Data Labeling and Learning Structured Prediction

TL;DR

This paper presents the sigma flow model for structured data labeling on Riemannian manifolds, integrating geometric PDEs with machine learning to improve prediction expressivity and connect to deep learning architectures.

Contribution

The paper introduces the sigma flow model that combines geometric PDEs with learnable Riemannian metrics for structured prediction on manifolds, bridging differential geometry and machine learning.

Findings

Proof of concept experiments show the model's expressivity.

Prediction performance demonstrates the effectiveness of the approach.

Structural similarities to transformer architectures are identified.

Abstract

This paper introduces the sigma flow model for the prediction of structured labelings of data observed on Riemannian manifolds, including Euclidean image domains as special case. The approach combines the Laplace-Beltrami framework for image denoising and enhancement, introduced by Sochen, Kimmel and Malladi about 25 years ago, and the assignment flow approach introduced and studied by the authors. The sigma flow arises as Riemannian gradient flow of generalized harmonic energies and thus is governed by a nonlinear geometric PDE which determines a harmonic map from a closed Riemannian domain manifold to a statistical manifold, equipped with the Fisher-Rao metric from information geometry. A specific ingredient of the sigma flow is the mutual dependency of the Riemannian metric of the domain manifold on the evolving state. This makes the approach amenable to machine learning in a specific way, by realizing this dependency through a mapping with compact time-variant parametrization that can be learned from data. Proof of concept experiments demonstrate the expressivity of the sigma flow model and prediction performance. Structural similarities to transformer network architectures and networks generated by the geometric integration of sigma flows are pointed out, which highlights the connection to deep learning and, conversely, may stimulate the use of geometric design principles for structured prediction in other areas of scientific machine learning.