Universal Learning of Distribution Flow using Ensemble Control Systems

Wei Zhang, Pengda Wang, Jr-Shin Li

Keywords: Ensemble control systems, flow matching, image restoration, moment kernelization

Abstract: Constructing interpretable representation models for learning flows of probability distributions has become a thriving research area in machine learning. This effort not only sparks new research avenues but also provides distinctive insights into established fields such as image processing. For example, the recently developed flow matching (FM) model has demonstrated its effectiveness in addressing image inverse problems. However, a unified and comprehensive framework for learning and generating probability flows from data in a general setting remains underexplored. In this work, inspired by an in-depth exploration of FM from a dynamical systems perspective, we develop an ensemble control system (ECS) model for probability flow learning. Our model is represented as an ensemble of heterogeneous control systems, with control inputs acting as time-dependent trainable parameters. The heterogeneous dynamics and time-dependent parameters significantly enhance the model capability of ECS, making it exceptionally powerful. To further capitalize on these strengths, we introduce a moment kernel transform, which generates a reduced kernel representation of ECS over a reproducing kernel Hilbert space, enabling efficient training. We demonstrate the significant advantages of the ECS model through various image restoration tasks and provide a detailed comparison with baseline FM-based image processing models.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 23690