New Ways to Design Deep Neural Networks

In this work, we propose a general framework for designing neural network architectures inspired by dynamic differential equations, utilizing the operator-splitting technique. The central idea is to treat neural network design as a discretizations of a continuous-time optimal control problem, where the underlying dynamics are governed by differential equations serving as constraints which is then unrolled as our network. These dynamics are discretized through operator-splitting schemes, which allow complex evolution equations to be decomposed into simpler substeps. Each step in the splitting scheme is then unrolled and interpreted as a layer in a neural network, with certain control variables modeled as learnable parameters. This formulation provides a principled way to incorporate prior knowledge about dynamics and structure into the network design. Using our theory, we give a rigorous mathematical explanation of the well-known UNet and show that it is a discretizations of a simple differential equation. By adding regularization to UNet, we can derive the PottsMGNet also through our proposed framework.