Neural Schrödinger Bridge with Sinkhorn Losses: Application to Data-driven Minimum Effort Control of Colloidal Self-assembly
This work addresses a specific problem in controlling colloidal self-assembly for materials science, offering a data-driven solution where physics-based models are difficult to obtain, but it is incremental as it builds on existing neural network and Schrödinger bridge methods.
The paper tackles the minimum effort control of colloidal self-assembly by formulating it as a generalized Schrödinger bridge problem with nonaffine coefficients, proposing a neural Schrödinger bridge framework that learns these coefficients from data and uses Sinkhorn losses for distributional constraints, achieving effectiveness in a numerical case study.
We show that the minimum effort control of colloidal self-assembly can be naturally formulated in the order-parameter space as a generalized Schrödinger bridge problem -- a class of fixed-horizon stochastic optimal control problems that originated in the works of Erwin Schrödinger in the early 1930s. In recent years, this class of problems has seen a resurgence of research activities in the control and machine learning communities. Different from the existing literature on the theory and computation for such problems, the controlled drift and diffusion coefficients for colloidal self-assembly are typically nonaffine in control, and are difficult to obtain from physics-based modeling. We deduce the conditions of optimality for such generalized problems, and show that the resulting system of equations is structurally very different from the existing results in a way that standard computational approaches no longer apply. Thus motivated, we propose a data-driven learning and control framework, named `neural Schrödinger bridge', to solve such generalized Schrödinger bridge problems by innovating on recent advances in neural networks. We illustrate the effectiveness of the proposed framework using a numerical case study of colloidal self-assembly. We learn the controlled drift and diffusion coefficients as two neural networks using molecular dynamics simulation data, and then use these two to train a third network with Sinkhorn losses designed for distributional endpoint constraints, specific for this class of control problems.