In this work, we propose a general framework for designing neural network architectures inspired by dynamic diferential 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 diferential 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 diferential equation. By adding regularization to UNet, we can derive the PottsMGNet also through our proposed framework.
(L. Li) https://raymondhonfu.github.io/ (R. H. Chan) https://www.norceresearch.no/personer/xue-cheng-tai (X. Tai); https://www.math.hkbu.edu.hk/~haoliu/ (H. Liu);
CEUR
ceur-ws.org image segmentation across noise levels via a unified framework grounded in multigrid and control perspectives. Notably, their analysis revealed that many encoder-decoder networks implicitly implement operator-splitting algorithms for control problems. Further advancing this interplay, Liu et al. [20] proposed a double-well network that learns region force terms in the Potts model through neural representations. We also want to mention that the work [21] has provided another way to interpret neural networks through multigrid techniques.
In this work, we aim to present a general framework for designing neural networks based on dynamic diferential equations using operator-splitting technique. The basic idea is discretizing a continuous control problem using the dynamic system as a constraint which is then approximated by some operatorsplitting schemes. We unroll this splitting scheme as a neural network with some control variables parameterized as learnable modules. We also illustrate this approach through examples of the UNet and the PottsMGNet.
A central insight of our work is that deep neural networks (DNN) can be naturally derived as discretizations of continuous control problems governed by partial diferential equations. This perspective provides a unifying framework to design and interpret network architectures through dynamical systems and optimal control.
Consider a continuous control problem where the state evolution is described by a PDE of the form:
= (, )
for (x, ) ∈ Ω × [0, ],
(0) = 0 in Ω,
(
matches the target state. The optimization procedure is the same as network training.
This is a general framework that applies to all networks. From another perspective, it also implies that for any given network, we can find a corresponding control problem so that this network is a numerical scheme (with proper discretization and unrolling) solving this problem. This framework ofers several advantages: • Theoretical Analysis: Tools from PDE theory (e.g., stability, convergence, and regularity analysis) can be applied to study the behavior of DNNs. • Architecture Design: New network architectures can be derived by choosing appropriate PDE models and discretization schemes. It enables to equip the designed neural networks with diferent kinds of physical properties which are often desired, but have not been able to achieve with existing networks. • Scalability, Interoperability and Explainability: The continuous viewpoint provides insights into the role of depth, initialization, and parameterization in DNNs.
In the following sections, we give details of this idea and demonstrate how to use operator-splitting methods and unrolling to discretize the control problems. Applications of this framework on UNet [1] and PottsMGNet [18] will be discussed.
Assume that we have decomposed the diferential operator in the following form: =1 = ∑ (, ), (0) = 0 where each (, ) is possibly nonlinear and acts as a distinct operator. Operator-splitting methods approximate the full solution by evolving through each operator sequentially or in parallel over a time step . For example, the sequential splitting [22, 23] evolves the solution by applying the sub-solvers sequentially: where , ()̂ = ( + ) denotes the solution operator of +1 = , ∘ ⋯ ∘ 1,
( ), = (, ), ( ) = . ̂ +1 = 1 ∑ =1 , ( ).
, splitting scheme [24] solves each sub-problem in parallel and combines them together ()̂ can either be solved explicitly or approximated numerically. The parallel We may also apply the sequential and parallel schemes together to form hybrid splitting schemes as explained in Appendix A, see also [18, Appendix D] for some more details.
After splitting, unrolling methods are used to construct neural networks. Unrolling maps the splitting scheme to a neural network architecture, where each operator becomes a network layer. Specifically, we have the following construction: 1. Layer structure: Each time step → +1 is unrolled into sub-layers, one per operator . If is known (e.g., a Laplacian), we implement as a fixed layer. If contains unknown learnable parameters, we represent as a trainable neural network block.
(, ) with parameters ().
The operator becomes a learnable layer 3. Unrolled network: By replacing with , model Θ with learnable parameters Θ = { ( )}
=1,=1 . , ,
.
, , in the splitting scheme, we get a neural network
This unrolled network can be trained by minimizing the discrepancy between predictions from a given initial condition 0 and the corresponding ground-truth solutions (e.g., MSE): ̂ =1 ℒ (Θ) =∑ ‖ Θ( 0) − ‖2 .
where { 0,
=̂ 1 is a set of training samples. Let us emphasis that we can use other loss functions
}
instead of MSE, such as the cross-netropy loss. Diferent deep neural networks in the literature are just
diferent approximations of some specially chosen PDEs or ODEs. In the following section, we will
show that the well-known UNet [1] is just an unrolling of a numerical approximation of a special PDE,
see (
The well-known deep neural network UNet [1] is designed to segment images into foreground and background represented by a binary image. By applying the framework in Section 2-3, we show that it is just a numerical approximation of a PDE.
Given an input image defined on the image domain Ω, we consider the following initial value problem { ( x, 0) = ( ( x)),x ∈ Ω, ( x,) = ( x, ) ∗ ( x, ) + () − ln ( x,)
, (x, ) ∈ Ω × (0, ],
1−( x,)
(
Equation (
Σ = { ∶ ( x, ) ≥ 0 for (x, ) ∈ Ω × (0, ]},
ℐΣ is the indicator function of Σ and ℐ Σ denotes the subdiferential of ℐΣ. By solving (
To solve (
In traditional multigrid methods, a popular framework is the ”fine grid → coarse grid → fine grid” strategy [27, 28]. Such a form of V-cycle multigrid method can be interpreted as space decomposition and subspace correction [29, 30].
Motivated by the fact that images can be viewed as piecewise constant functions in practice, in this
paper, we consider piecewise constant approximation of diferent resolutions. Let > 0 be some integer.
For the finest resolution, we consider discretizing into (2 )2 small squares, each of which is of size
(2−+1 ) × (2−+1 ). They are called pixels in image processing and rectangular elements in finite elements
methods. In real applications, the finest mesh is the grid the input image is given. Denote ℎ1 = 2−+1
and this set of pixles by 1. Starting with 1, we can define a sequence of coarse pixels { }+=11 .
Specifically, let ℎ = 2−1 ℎ1 = 2− . The set consists of coarse pixels with size ℎ × ℎ that form a
partition of . For each , we can define a set of piecewise constant basis functions. For the -th pixel
in , denoted by , we define (x)such that it equals to 1 if x ∈ and equals to 0 otherwise. We
denote the space spanned by this set of functions by . We have that
+1 ⊂ ⊂ ⋯ ⊂ 1.
(
We will decompose { ( x, ), ()} into a sum of variables with diferent scales over the multigrids.
Then, we use a hybrid splitting method to solve (
( x, ) = ( x, ) + ( ̃x, ), () = () + (̃).
These variables will be further decomposed next. Above, , are sums of control variables in the left branch of the V-cycle, and , ̃ ̃are sums of the control variables in the right branch, see an illustration in Figure 1. We also decompose the nonlinear operator as follows: − ln
Here, () contains nonlinear operations in the left branch and (̃) contains nonlinear operations in the right branch. In particular, we put − ln in (̃) only, i.e., () = ℐ Σ() and (̃) = 1− UNet, in which the operation − ln ℐ Σ() − ln . Later, we will show that our operator splitting method recovers a simplified corresponds to the sigmoid layer at the end of UNet. 2. We further decompose the operators into components at diferent grids as: ( x, ) = ∑ (x, ), () = ∑ (), () =
∑ (),
=1
(
is illustrated in Figure 1. Each equals to over if the support sets do not overlap. We also decompose and into a sum of components.
(x, ) = ∑ (x, ), = ∑ (), = ∑
(). =1 =1 =1 we have tered at (0, 0)is a shifted version of Note that ’s have diferent supports, making specifying the support for each function complicated. In order to simplify the settings, we shift for each , let be a shifting kernel so that = so that they are centered at (0, 0). Specifically, ∗ ̂ where ̂ supported on a square cen . According to the shift-invariant property of convolution, (x, ) ∗ ( x, ) = ∑ (x, ) ∗ ( x, ) = ∑(
(x, ) ∗ ̂ (x, )) ∗ ( x, ) × Thus learning a kernel of with 22(+1−) coeficients is converted to learning kernels ̂ (x, ) with coeficients arond the origin. In implementation, corresponds to the number of channels in a CNN and { }=1 corresponds to CNN convolution kernels of size × . In our experiments, for simplicity, we omit the shifting operator and replace ( (x, ) ∗ ( x, )) by ( x, )) and it gives good results. This is also what is done in the original UNet. The same decomposition is done for ̃, ̃and ̃. 4. For each grid level and each channel , we further decompose
̂ (x, ) = ∑ , (x, ), and , (x, ) has support around the origin with × position is to increase the number of parameters for the training variables. In our algorithm, each channel will compute an intermediate result. In the computation, , is used to convolve with the intermediate solution in the -th channel of grid level − 1 . The same decomposition is coeficients. The purpose of this decomBefore the decomposition, the control variables are ( x, ), () . After these decompositions, we see that the control variables are decomposed as in the following and the control variables are now
, (x, ), (), ̃, (x, ), ̃(),
∗, (x, ), ∗(),: ( x, ) = ∑ ∑
∑ , (x, ), ( ̃x, ) = ∑ ∑
∑ ̃, (x, ) + ∑ ∗(x, ), −1 =1 =1 =1 =1 =1 and the operators (), (̃) are decomposed as: ( x, ) = ∑ ∑ (x, ), (̃x, ) = ∑
∑ ̃(x, ) + ̃∗(x, ), () =
∑
∑ =1 =1 (), (̃) = with () = ̃() = ℐ Σ(), ∗() = − ln ̃() + ∗()
.
The original PDE (
{ ( x, 0) = ( ), x ∈ Ω.
= ∗ + ∗̃ + +
+̃ () + (̃), ( x, ) ∈ Ω × [0, ],
To solve (
. Denote the computed solution at time = by . The resulting algorithm that updates to +1 is summarized in Algorithm 1, where and represent the downsampling and upsampling operator respectively. For simplicity, variable dependencies on x are omitted.
where ReLU() = max{, 0̄} is the rectified linear unit.
Problem (
= max{, 0̄} = ReLU()̄, − ∗ = ∗̂ ∗ + −̂ ln
Data: The initial condition solution .
Compute ∈ by solving end for Compute = 1 ∑
=1 .
end for
4.3. On the solution to (
− ∑ ̂ ∗ ∗ − +̂ ℐ Σ() ∋ 0, where is some constant, ∗ = 1 ∑
=1 ∗ is known, ̂ is a convolution operator, is the number of
input channel, ̂ is a bias function. The solution to (
=̄ ∗ + (∑=1 ̂ ∗ ∗ + ̂ ),
sub-step, it is, in fact, a projection. Its closed-form solution is given as
In (
1 where Sig() = 1+ − is the sigmoid function. By repeating (26) so that +1 converges to some function ∗, we set = ∗. In particular, since 0 = ,̄ the updating formula (26) always gives 1 = 0.5. If we only consider a two-step fixed point iteration, we get = Sig (− 0.5 − ̄
) = Sig ( −̄ 0.5 ) .
We first show that a building block of Algorithm 1 is equivalent to a layer of a simplified UNet. Each layer of UNet is a convolution layer activated by ReLU:
{ − ̄ = − ln .
1−
=̄ ∗ + (∑=1 ̂ ∗ ∗ + ̂ ), − ̄
= − ln +1 = Sig (−
+1 1 − +1 − ̄
) ,
{
=̄
= ReLU( )̄ ,
∑
=1 ∗ ∗ + ,
=1
=1
where is a convolutional kernel and is the bias. In Algorithm 1, the building block is (
in Algorithm 1.
We have =̄ ̄ , and = . Essentially, Algorithm 1 and UNet are the same. Thus, we have shown that a
simplified UNet structure (with only 1 convolution layer at each data resolution) is equivalent to one
iteration of Algorithm 1. UNet architecture consists of four components: encoder, decoder, bottleneck,
and skip-connections, each of which has a corresponding component in the structure of Algorithm 1:
1. Encoder: Encoder in UNet corresponds to the left branch of the V-cycle in Algorithm 1. The
number of data resolution levels corresponds to the number of grid levels .
2. Decoder: Decoder in UNet corresponds to the right branch of the V-cycle in Algorithm 1.
3. Bottleneck: Bottleneck in UNet corresponds to the computations at the coarsest grid level (grid
4. Skip-layer connection: Skip-layer connections in UNet correspond to the relaxation steps (
Algorithm 2: A hybrid splitting method to solve the control problem (
Compute ∈ by solving end for +1 − 1̄ − ∗( ) ∗ 1̄ − ∗( ) − ∗( +1 ) ∋ 0.
above given split of operators.
The second substep is then defined as
In the previous section, we presented how to use the framework introduced in Section 2-3 to show that
the operator splitting scheme of (
We consider the control problem (
, ∗() = − 1 ln
variance 2. To solve this new system, we consider a multi-step operator splitting algorithm (Algorithm
2) using the hybrid splitting method in Appendix A for approximating the solution of (
Compared with the UNet structure, the PottsMGNet starts with a control problem with a diferent nonlinear term and ,̃which results in a diferent activation function at each layer of the resulted neural network. The position of the relaxation step is also diferent, which leads to a diferent skip-connection structure. Additionally, the UNet is a one-step algorithm while the PottsMGNet is a multi-step algorithm. Similar to UNet, we can also extend Algorithm 2 to a multi-channel case by further split the control variables.
In this work, we present a framework to design novel neural networks using operator splitting of some control problems. Networks designed in this way are usually robust to disturbances in the input, because concerned control problems usually include some regularization terms. We also presented how to derive UNet and PottsMGNet from this framework. In the future, we would apply this framework to other types of networks like graph neural networks.
During the preparation of this work, the authors used DeepSeek in order to: Grammar and spelling check. After using these tools, the authors reviewed and edited the content as needed and take full responsibility for the publication’s content.
A. Hybrid Splitting Schemes for Initial Value Problems and Neural
We discuss a hybrid splitting scheme proposed in [18] as a powerful technique for solving initial value problems, which can be efectively leveraged to design and understand the architecture of deep neural networks. A hybrid splitting scheme combines the principles of parallel and sequential splitting schemes to decompose complex problems into more manageable subproblems.
Consider the initial value problem: ∑ =1 , (x, ; ) + ∑ (x, ; ) + (x, )) = 0 on Ω × [0, ], (0) = 0 .
We can first split it into sequential steps, where each step consists of parallel substeps. At the ( + 1) -th time step, denote the intermediate result computed in the -th branch of the -th sequential , and define +(−1)/ = 1 ∑ =1 +(−1)/ . The function +/ is computed by +
∑ (∑
∑ =1 =1 =1 step by solving:
+/ − +(−1)/ =1 =1 = − ∑ , (x, ; +(−1)/ ) − (x, +1 ; +/ ) − (x, ).
Theorem D.1], meaning that the error ‖ +1 − ( +1 )‖∞ is of the order ( ) . from 1 to and from 1 to to compute +1 .
, are treated explicitly, while the operators are treated implicitly. Functions can be treated explicitly as shown here. Starting with ,0 = , this algorithm iterates through When
, and are Lipschitz continuous, this hybrid splitting scheme is first-order accurate [ 18, [23] R. Glowinski, T.-W. Pan, X.-C. Tai, Some facts about operator-splitting and alternating direction methods, in: Splitting Methods in Communication, Imaging, Science, and Engineering, Springer, 2016, pp. 19–94. [24] T. Lu, P. Neittaanmaki, X.-C. Tai, A parallel splitting-up method for partial diferential equations and its applications to navier-stokes equations, ESAIM: Mathematical Modelling and Numerical Analysis 26 (1992) 673–708. [25] X. Tai, L. Li, E. Bae, The potts model with diferent piecewise constant representations and fast algorithms: a survey, Handbook of Mathematical Models and Algorithms in Computer Vision and Imaging: Mathematical Imaging and Vision (2021) 1–41. [26] X.-C. Tai, H. Liu, R. H. Chan, L. Li, A mathematical explanation of unet, Mathematical Foundations of Computing (2024) 0–0. [27] T. F. Chan, T. P. Mathew, Domain decomposition algorithms, Acta numerica 3 (1994) 61–143. [28] J. Xu, Iterative methods by space decomposition and subspace correction, SIAM review 34 (1992) 581–613. [29] X.-C. Tai, Rate of convergence for some constraint decomposition methods for nonlinear variational inequalities, Numerische Mathematik 93 (2003) 755–786. [30] X.-C. Tai, J. Xu, Global and uniform convergence of subspace correction methods for some convex optimization problems, Mathematics of Computation 71 (2002) 105–124.