A Lightweight U-like Network Utilizing Neural Memory Ordinary Differential Equations for Slimming the Decoder

1 Introduction

Deep learning has become increasingly prominent in the field of paramedicine, providing valuable assistance in disease detection and diagnosis. The UNet [Ronneberger et al., 2015] is a significant milestone in demonstrating the efficacy of encoder-decoder Convolutional Neural Networks (CNNs) with skip connections for medical image segmentation. Over time, UNet has become the foundational framework for most notable medical image segmentation method- Att-UNet(EED)

∗Corresponding author.

Att-UNet

20

30

40

Parameters (M)

(a) mIoU on PH2.

(b) mIoU on ISIC2018.

Figure 1: Visualization of the mIoU results on the PH2 and ISIC2018 datasets. X-axis corresponds to the number of parameters (lower the better). Y-axis represents the mIoU (higher the better). The circle's area is proportional to FLOPs (smaller the better).

ologies. To achieve higher precision, many studies have introduced complex modules or increased the number of parameters. For example, ResUNet [Zhang et al., 2018] combines the UNet architecture with residual connections inspired by ResNet [He et al., 2015]. Swin-UNet [Cao et al., 2022] leverages the Swin Transformer [Liu et al., 2021] architecture. The TransFuse [Zhang et al., 2021] model employs a dualpath configuration that seamlessly combines CNN and ViT [Dosovitskiy et al., 2020], enabling it to capture both local and global insights simultaneously.

Existing U-like networks have improved the performance of UNet, but most of them come with an increased computational cost. This can pose challenges when attempting to deploy them in practical application scenarios with limited computing resources. To address this issue, researchers have ventured into the research direction that goes beyond the pursuit of model performance, focusing on model lightweighting. In recent developments, UNeXt [Valanarasu and Patel, 2022] combined UNet and MLP [Tolstikhin et al., 2021], presenting a lightweight architecture that achieves remarkable performance while reducing both parameter count and computational requirements. MALUNet [Ruan et al., 2022] has successfully reduced model size by decreasing the number of model channels and incorporating multiple attention modules. The Efficient Group Enhanced UNet (EGEUNet) [Ruan et al., 2023] reduced the number of parameters and computational complexity through the integration of enhanced attention mechanisms and feature fusion modules.

While the aforementioned methods have demonstrated impressive abilities in reducing model parameters and computational complexity, they often face challenges when it comes to adaptability to existing frameworks and lack universality. To achieve lightweighting, most of these networks heavily reduce the number of channels in their network structure, resulting in a trade-off between network performance and the reduction of parameters and computations. However, these networks often pursue lightweighting and network performance as separate objectives, lacking a comprehensive solution. In order to bridge the gap between network performance and computational complexity, this paper aims to explore a universal method that has the potential to adapt to all U-like networks.

In this paper, we integrate the neural memory Ordinary Differential Equations (nmODEs) [Yi, 2023] decoders into U-like networks using three different discretization methods: explicit Euler's method, Heun's method, and linear multistep method. We then assess their effectiveness on multiple datasets. The application of the nmODEs decoders across different UNet variants showcases the versatility of our proposed solution. Moreover, experimental results substantiate that the nmODEs decoders effectively reduce parameters and computational complexity, without compromising the performance of the original model, and in some cases even slightly improving it. In summary, the proposed discretized nmODEs decoders can decrease the parameter count by approximately 20% ∼ 50% and the FLOPs by up to 74%.

2 RelatedWork

To delve into the integration of U-like networks with the theme of nmODEs, this section will first introduce several classic UNet networks and methods related to nmODEs.

2.1 UNet and Its Variants

UNet [Ronneberger et al., 2015] was the cornerstone of all UNet variants for semantic image segmentation. It has a Ulike structure with a contracting path for context extraction and an expanding path for precise localization. Utilizing skip connections, UNet effectively preserved fine feature details. Attention UNet (Att-UNet) [Oktay et al., 2018] incorporated an attention mechanism into the UNet, significantly improving its efficacy in extracting image features. UNeXt [Valanarasu and Patel, 2022], being the pioneering lightweight medical image segmentation network that integrates convolutional and multilayer perceptrons, had shown remarkable success in reducing network parameters and computational workload. MALUNet [Ruan et al., 2022] demonstrated outstanding performance in skin cancer segmentation by incorporating attention mechanisms and deep separable convolution into UNet through carefully designed modules. Building upon MALUNet, EGE-UNet [Ruan et al., 2023] further refined the attention mechanism and introduces new feature fusion modules, ensuring network performance while significantly reducing network complexity. EGE-UNet surpassed many large-scale networks in terms of both performance and efficiency.

2.2 Neural Ordinary Differential Equations

Ordinary Differential Equation (ODE) systems, recognized as a distinctive class of dynamical systems, have long been subject to extensive exploration and empirical investigation in the realms of mathematics and physics. Neural ODEs (NODEs) [Chen et al., 2018] provided mathematical principles that elucidate ResNet, transforming it from an enigmatic black-box network into a comprehensible framework. They presented a novel approach to conceptualize neural networks as representations of ODEs. NODEs established a foundation for the unification of neural networks and ODEs. One notable advantage of NODEs is their ability to eliminate the storage of intermediate quantities during the forward propagation, resulting in a substantial reduction in parameters and computational overhead. Nevertheless, there are certain limitations in mapping data through NODEs. For instance, they face challenges in representing mappings like g(1) = -1, g(-1) = 1. This limitation arises from the fact that NODEs models utilizing data inputs as initial values can only learn features within the same topological space as the input data [Dupont et al., 2019]. Furthermore, when modeling problems with differential equations, creating a dynamical system, it has been demonstrated that attractors in dynamical systems are believed to be linked to memory capacity [Poucet and Save, 2005] [Wills et al., 2005]. However, conventional NODEs lack the capability to effectively leverage the memory capacity provided by attractors.

The nmODEs [Yi, 2023] are specialized variants of NODEs designed to overcome the limitations inherent in traditional NODEs and harness the full memory capabilities offered by dynamical systems. It enhanced the neural network's nonlinear expression capability by employing implicit mapping and utilizing nonlinear activation functions. Simulating the dynamical system governing neocortical neuronal memory, nmODEs introduced a distinct feature-a clear dynamical characterization achieved through the segregation of learning neurons and memory neurons. Unlike previous NODEs approaches, nmODEs treated input data as external parameters rather than utilizing them as initial values for ODEs. By separating the neuron's function into learning and memory components, learning exclusively occurs in the learning part, while the memory part maps the input to its global attractor, establishing a mapping from input space to memory space. The nmODEs have found successful applications in various segmentation tasks. For instance, the nmPLS-Net [Dong et al., 2023] leveraged the robust nonlinear representation and memory capabilities of nmODEs to construct an edge segmentation-based decoding network. This approach has enabled accurate lung lobe segmentation. Additionally, the integration of nmODEs into UNet, employing a straightforward discretization method, has demonstrated promising outcomes in tasks such as diabetic kidney segmentation [Wang et al., 2024] and liver segmentation [He et al., 2023]. [Hu et al., 2023] enhanced the robustness of medical image segmentation using nmODEs and yielded favorable outcomes. These applications showcase the versatility and effectiveness of nmODEs across different medical image segmentation challenges.

3 Methods

In this section, we leverage the discretization methods of nmODEs to construct a lightweight U-like network architecture. nmODEs have the nature to be adaptive to the decoders of U-like networks. Therefore, in Section 3.1, we will introduce the U-like networks with the discretized nmODEs decoders. These U-like networks utilize skip connections to extract low-level features and inject them into internal states for feature aggregation, resulting in a series of low parameter nmODEs decoders. In Section 3.2, we will introduce three discretized ODE solvers to approximately calculate the output of the nmODEs decoders.

3.1 The U-like Networks with nmODEs Decoders

ሻ𝑡ሺݕ

(a) NODE block.

To begin with, let's revisit the disparity in inputs between conventional NODE and nmODE. Their structures are delineated in Fig. 3(a) and Fig. 3(b), respectively. In the typical NODEs configuration, the initial value y(0) is derived from the data itself, and the output represents the numerical solution of the ODE. In contrast, within the nmODEs architecture, the initial value y(0) is set to a random value, e.g., 0. Data served as an sequential external input x(t). nmODEs take two inputs, aligning seamlessly with the decoders of U-like networks. Information from the skip connections serves as the sequential external inputs to the nmODEs decoders, while information from the upward path serves as y(t), enabling the full utilization of nmODEs. This is exactly the reason why we chose nmODEs. The differential equation for nmODEs is formulated as follows:

For t ≥ 0, where y(t) ∈ Rn represents the network's state, x(t) ∈ Rm stands for external input, and θt means the parameters of the skip connections. We use the nmODEs to formulate the decoders of the U-like networks, resulting in an modified U-like network as illustrated in Fig. 2. This network applies parameterized computation in the skip connections, which is denoted as g(x(t), θt). The duty of the skip connections is to transform the low-level features to the fixed-size high-level feature maps. The U-like network with nmODEs decoders performs parameter-less feature aggregation. Given an L-layer U-like network using nmODEs decoders, define the output of the network encoder, which is the skip-connected input of the decoder, as xl , and denote the upward path input of the decoder as yl , 1 ≤ l ≤ L.Initializing x(0) = xL and y(0) = yL = 0, the target of the nmODEs decoders is to obtain the numerical solution at the point (τ, x(τ) = x1 ) along the trajectory of the Eq. (1). In the rest of this section, we will introduce three discretized ODE solvers, e.g. explicit Euler's method, Heun's method, and linear multistep method.

For simplicity in subsequent discussions, let's denote Eq. (1) as F(t, y) , the function F remains continuous and adheres to certain Lipschitz conditions, ensuring the solution's existence and uniqueness. Given an initial value y0, the trajectory traced by Eq. (1) starting from y0 is denoted as yt, encompassing all instances where t ≥ 0. A vector y∗ is defined as an equilibrium point of the NODEs if it satisfies the equation F(t∗ , y∗ ) = 0. An equilibrium point y∗ earns the designation of a global attractor when, for any given y0, the corresponding trajectory yt converges towards y∗ as t →∞. With the existence of global attractor within nmODEs, the establishment of a favorable nonlinear mapping from x to y∗ becomes feasible.

ሻ𝑡ሺݕ 3.2 DiscretizationMethods

𝜃ሻ ,ሻ𝑡ሺݕ ,ሻ𝑡ሺݔ𝑓ሺ = ሻ𝑡ሺݕሶ In the modified lightweight U-like networks, the upward path is parameterless. Only a small number of parameters are responsible for information integration of skip connections, matching the feature sizes at the upward path. We discretize the nmODEs decoders in three different methods and replace the original decoders of U-like networks. The structure inside the decoders varies with the discretization method, with specific reference to the mathematical derivation. All three discretization methods address initial value problems (IVPs) for nmODEs, sharing common steps.

Theorem 1. Explicit Euler's Method [Euler, 1845]. Given the derivative ˙y(t) = F(t, y(t)), choose a value δ for the size of every step along t-axis and set tn+1 = tn + δ. The yn+1 from yn and tn is

where yn is an approximate solution at time tn, i.e., yn ≈ y(tn). yn+1 is an explicit function of yi for i ≤ n.

The explicit Euler's method, defined in Theorem 1, is the most straightforward in concept and also the most convenient discretization approach to apply. Let t = tn, we can derive from the Eq. (1) and (2) that

Eq. (3) is formulated for continuous values, and adapting it to discretized neural networks necessitates a conversion. In the context of a U-like network structure, where the value yl in the lowest layer serves as the initial value, the initial value problem is solved during the operation of the upward path. This process results in a discrete form, with the number of solution steps being inversely proportional to the number of network layers:

In the context of applying U-like networks, the total number of layers is L. l is the number of layers of the current operation, 1 ≤ l ≤ L. yl in Eq. (4) represents the input of the l-th layer of the upward path, serving as the initial value for solving yl-1 . Here, xl denotes the information transmitted through the skip connection, and δ is 1/L. θl means the parameters for skip connection in l-th layer. Based on Eq. (4) we show the internal realization of the decoder as Fig. 4 (a).

Theorem 2. Heun's method [Heun, 1900]. Given the derivative ˙y(t) = F(t, y(t)), choose a value δ for the size of every step along t-axis and set tn+1 = tn +δ. First calculate the intermediate value yn+1 and then the final approximation yn+1 at the next integration point as follows:

where yn is an approximate solution at time tn, i.e., yn ≈ y(tn). yn+1 is the predicted median, yn+1 is the final result corrected for yn+1.

The Heun's method, defined in Theorem 2, replaces the derivative at the point yn in the explicit Euler's method with the average of derivatives at yn and yn+1. Therefore, from a mathematical perspective, this method enhances accuracy compared to the explicit Euler's method. Let t = tn, we can derive from the Eq. (1) and (5) that

Since we design reusable modules for ˙y(tn), there is no more substitution of expansion in Eq. (6). Similarly, we rewrite the continuous numerical form into a discretized network form to obtain the following equation:

The decoder structure based on this construction is illustrated in Fig. 4 (b). It is worth noting that the Heun's method requires two layers of skip connection information as input xl and xl-1 : the current layer and the next layer. Which means 2 ≤ l ≤ L. Therefore, the decoder of the layer 1 applies an EED.

Theorem 3. Linear Multistep Method [Bashforth and Adams, 1883]. Given the derivative ˙y(t) = F(t, y(t)), choose a value δ for the size of every step along t-axis and set tn+1 = tn + δ as follows:

where yn is an approximate solution at time tn, i.e., yn ≈ y(tn). yn+1 is the final result associated with both yn and yn-1.

The linear multistep method, defined in Theorem 3, involves a linear combination of the derivatives of multiple selected points, the most commonly used number of selected points is two, and the calculation of the unknown points is performed by the information of the known two points, this linear multistep method is also known as the Eulerian twostep method. Let t = tn, we can derive from the Eq. (1) and Eq. (8) that

Rewrite it in discretized form:

The internal structure of the decoder designed based on Eq. (10) is shown in Fig. 4 (c). It is worth noting that the linear multistep method requires two layers of upward path information as input yl and yl+1 , the current layer and the previous layer. Which means 1 ≤ l ≤ L - 1, so the decoder of the layer L applies a EED.

Function Selection & Initial Value Problem

In the network framework we devised, the upward path is parameter-free, meaning the f-function has no parameters to alter the shape of input yl . Consequently, the responsibility of reshaping xl to match yl falls on the g-function. Therefore, the g-function includes essential components like the up-sampling function to modify height and width, and the convolution operation to adjust the number of channels. Moreover, we have the flexibility to incorporate parameterfree operations, such as batch normalization and activation functions, within both the g and f functions. In this paper, the internal implementations of the g-function include, in sequence, the convolution, upsampling, and activation functions, while the f-function exclusively comprises batch normalization. For the initialization of y(0), we set it to a zero matrix, maintaining consistency with the initial input width and height. However, there is some flexibility in determining the channels of y(0). Typically, the number of channels is chosen to be consistent with the original input, but considerations are also given to the number of groups to facilitate grouped convolution operations. As demonstrated in the experimental section, augmenting the number of channels enhances performance to some extent, albeit accompanied by an escalation in the number of model parameters and FLOPs.

4 Experiments

4.1 Datasets

PH2. The PH2 1 dataset is a compilation of 200 dermoscopic images focused on dermatology, specializing in melanocytic lesions. It enables research in skin lesion segmentation and classification, with a particular emphasis on nevus and melanoma.

ISIC2017. The ISIC2017 2 dataset is a comprehensive compilation of 2150 dermoscopic images of skin lesions. It serves as a valuable resource for advancements in dermatology and computer-aided diagnosis. The dataset covers a wide range of skin conditions, including both benign and malignant lesions. It is primarily designed for tasks related to melanoma detection and skin cancer classification.

1 https://www.fc.up.pt/addi/ph2%20database.html

2 https://challenge.isic-archive.com/data/#2017

ISIC2018. The ISIC2018 3 dataset is a valuable resource in dermatological image analysis, focusing on melanoma detection and skin cancer classification. It consists of approximately 2700 diverse dermoscopic images.

4.2 Implementation Details

In our experiments, we replaced the decoders of UNet, AttUNet, MALUNet, EGE-UNet and UNeXt with nmODEs decoders using various discretization methods. All experiments were executed on a single RTX 4090 GPU using PyTorch 4 . Each dataset was randomly split into training and testing sets with a 7:3 ratio, and images were normalized and resized to 256 × 256 for consistency. Data augmentation techniques, including horizontal flipping, vertical flipping, and random rotation, were applied. The optimizer AdamW [Loshchilov and Hutter, 2017] was used, along with the CosineAnnealingLR [Loshchilov and Hutter, ] scheduler, setting the maximum iteration count to 50 and the minimum learning rate to 1e-5. Training spanned 300 epochs with a batch size of 8. Mean Intersection over Union (mIoU) and Dice similarity score (DSC) were employed as evaluation metrics. Experiments were repeated five times, and results were reported as mean and standard deviation for each dataset.

4.3 Comparative Results

We chose five prevalent networks from the existing U-like networks, each with documented experimental results on relevant datasets. Subsequently, we replaced their decoders with nmODEs decoders employing different discretizations. The comparison with the original networks and other major models is illustrated in Table 1. The experimental results clearly demonstrates the undeniable effectiveness of our proposed nmODEs decoders in terms of parameter and FLOPs reduction. For instance, the Att-UNet, which initially boasts the highest number of parameters, sees a reduction of approximately 46% in parameters and 74% in FLOPs after applying EED. Similarly, the classical UNet experiences a reduction of about 30% in parameters and 40% in FLOPs with the application of HD. Even in the case of lightweight networks, a consistent reduction of around 50% in parameters and 20%-30% in FLOPs is observed after integrating the nmODEs decoders.

In terms of performance, the networks with the nmODEs decoders generally outperform the original networks, while those without performance improvement largely maintain the level of the original networks. On the PH2 dataset, both mIoU and DSC of Att-UNet with EED applied improve by about 5 percentage points over the original network, reaching the state-of-the-art. Similarly, mIoU and DSC of UNet with LMD applied improve by 2.7 and 1.6 percentage points over the original network, also reaching the state-of-theart, as demonstrated in the ablation experiment section. On the ISIC2017 dataset, the networks employing nmODEs decoders largely retain the performance of the original networks. On the ISIC2018 dataset, EGE-UNet with LMD applied outperforms the original network and reaches the stateof-the-art level.

3 https://challenge.isic-archive.com/data/#2018

4 https://pytorch.org/get-started/locally/

All three of our proposed decoders effectively reduce the number of parameters and FLOPs without compromising the original network's performance. However, among these three decoders, HD, with one additional g-function compared to EED and LMD, incurs more parameters and FLOPs. Surprisingly, its performance is not as strong as that of the other two decoders, raising questions that warrant further investigation.

4.4 Representative Results

Figures 5 illustrates the qualitative segmentation results obtained by replacing the decoders of MALUNet, EGE-UNet, Att-UNet and UNet with the nmODEs decoders. These comparisons are made against the original network using partially representative images. In each figure, the last two columns highlight that the mask predicted by the network with the nmODEs decoders is closer to the ground truth than the mask predicted by the original network.

4.5 Ablation Study

In our ablation experiments using UNet as the base network on the PH2 dataset, our proposed decoders demonstrated a significant reduction in parameters. This reduction can be attributed not only to the efficient computational mechanism of NODEs but also to the fact that the number of channels in the initial y(0) is only 3, maintaining consistency with the upward path. The ablation experiments involved starting with the original UNet and progressively reducing the number of channels in its upward path to 3. Subsequently, we simplified the nmODEs decoders by retaining only the part of the skip connection summation, denoted as yl-1 = yl + g(xl , θl ), ultimately utilizing all three decoders.

Analyzing the experimental results presented in Table 2 for the PH2 dataset, it is observed that the segmentation performance experiences a slight degradation when reducing the channels in the upward path of the original UNet to 3. This slight impact could be attributed to the unchanged downward path of the network, crucial for extracting feature informa- MALUNet MALUNet

(EED)

(a) Visualization on the

ISIC2018

EGE-UNet EGE-UNet

(LMD)

(b) Visualization on the

ISIC2018

Att-UNet Att-UNet

(EED)

(c) Visualization on the

PH2

(d) Visualization on the

ISIC2017

Table 2: Ablation experiments on network upward path structure.

tion across different levels. In UNet, the upward path primarily integrates information and produces output, with the final output channels consistently minimal. We hypothesize that the number of channels in the upward path of UNet might not significantly affect information integration capabilities, with the primary influence residing in the algorithm inside the decoder. Upon applying the simplified nmODEs decoder, a slight improvement in the network's performance is observed. However, substantial enhancement occurs with the complete nmODEs decoders using three different discretizations, confirming the significant influence of the decoder algorithm on information integration capabilities.

The number of channels in the initial y(0) plays a pivotal role in determining the overall network complexity, exerting a significant impact on its information integration capabilities. Regarded as a crucial hyperparameter, its selection is of paramount importance. Table 3 provides a thorough comparison of nmODEs decoders employing different numbers of channels in the initial y(0). The table illustrates that a three-channel y(0), aligning with the input information, proves more effective than a single-channel y(0). Similarly, increasing the number of channels in y(0) to 8 leads to an improvement in the network's performance. Opting for an even higher number of channels, such as 16 or 32, would lead to an escalation in the number of parameters and network complexity, deviating from our original intention. Therefore, such configurations are not utilized.

5 Conclusion

In this paper, we introduce three specialized nmODEs decoders tailored for U-like networks. By discretizing nmODEs using methods like explicit Euler's, Heun's, and linear multistep methods, we exploit the benefits of the upward path's low-channel parameterlessness and the memory efficiency of neural ordinary differential equations. Our decoders drastically reduce parameters and computational demands in Ulike networks. To thoroughly assess their impact, we undertake a comprehensive series of segmentation experiments across widely-used datasets. We conduct an extensive comparative analysis between prevalent U-like networks and their nmODEs-enhanced counterparts. The outcomes of our study shed light on the streamlined complexity exhibited by networks that integrate nmODEs decoders, showcasing their remarkable capability to either maintain or enhance performance. Furthermore, our findings indicate that in specific tasks and network configurations, there exists a promising potential to achieve state-of-the-art performance levels. We hope that our work will provide fresh insights for the development of lightweight medical image segmentation models.

Acknowledgements

This work was supported by the National Major Science and Technology Projects of China under Grant 2018AAA0100201, the National Natural Science Foundation of China under Grant 62206189, and the China Postdoctoral Science Foundation under Grant 2023M732427.