TvaraNet: A Lightweight Mamba Neural Network for Real-Time Medical Image Segmentation

1 Introduction

Medical image segmentation is a fundamental yet challenging task due to anatomical variability, low contrast in tissue images, and the presence of noise commonly found in medical imaging modalities. Encoder-decoder architectures such as UNet [Ronneberger et al., 2015] have established themselves as strong segmentation baselines by learning hierarchical feature representations. These architectures enable accurate delineation of anatomical structures and pathological regions.

Recent advances in backbone networks, including ResNetbased [Nisa and Ismail, 2022] and Transformer-based approaches [Chen et al., 2021], further enhance feature extraction and contextual reasoning capabilities. However, these gains often come at the expense of high computational complexity and large parameter counts, which hinder deployment in real-time resource-constrained clinical environments. The emergence of state-space models has shifted research focus toward lightweight architectures, leading to designs such as VMUNet and its variants [Ruan et al., 2024; Zhang et al., 2024b; Wu et al., 2025; Zhang et al., 2024a]. Although these models reduce computational overhead, they remain insufficiently compact for edge deployment. Ultralight architectures [Wu et al., 2024; Liao et al., 2024] improve efficiency but still exhibit critical limitations, namely imprecise boundary extraction, false positives, and spurious island predictions in the output masks.

1.1 Motivation and Contribution

Existing lightweight medical image segmentation models, such as Ultralight VMUNet [Wu et al., 2024], show poor shape fidelity and irregular feature activations around segmentation boundaries in predicted masks as shown in Fig. 1, which results in incomplete or distorted anatomical structures. In contrast, heavier segmentation networks mitigate this issue by leveraging deeper representations, but this comes with prohibitive computational and memory costs. These trade-offs motivate the following research question: Can an accurate, shape-preserving segmentation mask be achieved under strict parameter and computational complexity constraints?

To address this, we propose TvaraNet, an extremely lightweight architecture designed for accurate and boundaryconsistent medical image segmentation. Our contributions are summarized as follows:

• PASMA enriches feature interactions across channel groups in Parallel Vision Mamba (PVM) [Wu et al., 2024] by injecting globally aggregated channel-aware features prior to the Mamba update.
• Lightweight U-shaped segmentation networks propagate spurious features through skip connections, which lead to distorted shapes and irregular boundary activations. We address this using the DASMA block, which applies lightweight multi-scale spatial modulation to suppress irrelevant features and enhance spatially informative representations in skip pathways.
• Training-time feature suppression remains a challenging problem as spurious features are often reinforced during backpropagation. To address this, we introduce the AGGRL, which performs input-dependent, channelwise gradient inversion to selectively suppress redundant features while promoting more discriminative representation learning.
• Adaptive SVD loss enforces low-rank regularization on encoder features to suppress noise and redundancy, resulting in more compact and discriminative feature representations.
• Standard region-based losses often underemphasize boundaries, resulting in imprecise contour localization. To address this, BWCE loss is introduced to impose higher penalties for incorrectly predicted boundary pixels during training.

The remainder of this paper is organized as follows. Section 2 reviews related work, Section 3 details the proposed methodology, Section 4 discusses the experimental results, and Section 5 concludes the study.

2 RelatedWork

Medical image segmentation is a critical field that integrates artificial intelligence into the medical domain [Yao et al., 2024; Lee et al., 2024; Shen et al., 2023; Song et al., 2024; Kim and Yoo, 2025]. This aids clinical studies in various applications, such as dermatology and neuroanatomy. Advancements in this area began with the development of the encoderdecoder based architecture called UNet [Ronneberger et al.,

2015], which inspired numerous subsequent variants [Huang et al., 2020; Yan et al., 2022; Zhou et al., 2018; Peng et al., 2025]. This evolution was followed by the insertion of recently developed backbones, which provided stronger feature extraction capabilities, such as Resnet [Nisa and Ismail, 2022], and ResNext [Xie et al., 2020], which primarily rely on Convolutional Neural Network (CNN). Subsequently, research shifted toward transformers, which had the power of self-attention and cross-attention while maintaining the architecture style of UNet. This led to the development of transformer-based backbones such as TransUNet [Chen et al., 2021], Segformer [Xie et al., 2021], Segmenter [Strudel et al., 2021], and Swin-UNet [Cao et al., 2022]. Later, the emergence of state-space models shifted research toward lightweight architectures, leading to the development of models such as VM-UNet [Ruan et al., 2024], VMUNetv2 [Zhang et al., 2024b], HMTUNet [Zhang et al., 2024a], Swin-UMamba [Liu et al., 2024], HVMUNet [Wu et al., 2025], UNet-mamba [Wang et al., 2024], which performed better than the earlier models while reducing the parameter count. Subsequently, research efforts focused on developing extremely lightweight architectures with minimal parameter counts and computational requirements, such as Ultralight VMUNet [Wu et al., 2024] and LightM-UNet [Liao et al., 2024]. Due to the reduction in parameters, these models often suffer from inconsistent boundary predictions in the segmentation masks of medical images. To address this issue, we propose TvaraNet in the following section.

3 Methodology

In this section, we present the architecture of TvaraNet, an ultra-lightweight segmentation network designed to preserve boundary fidelity under extreme efficiency constraints. TvaraNet follows an encoder-decoder architecture while introducing novel enhancements in the encoder, decoder, skip connections, and loss design to overcome common limitations of lightweight models, including inaccurate boundary localization and redundant feature propagation. We briefly describe each component of TvaraNet in the following subsections. A visualization of the TvaraNet model is shown in Fig. 2.

3.1 Parallel SkipMamba Averaging

PVM processes channels independently, limiting its capability to capture semantics spread across channels. PASMA introduces globally aggregated features into the channel chunks of PVM. This is achieved by computing the channel-wise average and injecting it into the mamba layer. Let Xin ∈ RB×C×H×W be the input to PASMA, which is passed through the following operations,

In Eq. (1), the input is processed by a 1 × 1 convolution followed by layer normalization (LN). The normalized feature map is then split along the channel dimension into four equal groups of size C/4 using the split operation Sp, producing Y C/4 i , where i ∈ {1, 2, 3, 4}.

In Eq. (2), the average of the channel-split outputs is computed as Y C/4 avg .

The averaged channel representation is then added to each channel chunk and passed through the Mamba layer, enabling inter-channel features across the split channel groups. To preserve the original feature representation, a skip connection with a scaling parameter θ is incorporated, as shown in Eq. (3). The resulting output is denoted as YM C/4 i .

The outputs obtained from Eq. (3) are concatenated along the channel dimension, as shown in Eq. (4), to obtain Yout.

The concatenated output is then passed through the layer normalization LN to obtain the normalized output, and then it is projected to the required output dimension using the multilayer perceptron projection (Proj) layer in Eq. (5).

3.2 Dilated Asymmetric Spatial Mixer Attention

DASMA enhances the skip connection by introducing scaleenhanced features. This is done by introducing three convolution branches with different dilation rates. Let Zin be the input which belongs to RB×C×H×W .

The input Zin is processed by the first convolution branch, which consists of three layers. Specifically, Zin is processed by asymmetric depthwise convolutions with kernel sizes 1×3 and 3 × 1, each employing a unidimensional padding of 1, as defined in Eqs. (6) and (7). This is followed by a pointwise convolution (PW) with a 1 × 1 kernel, which projects the incoming features to a single-channel representation, as shown in Eq. (8).

Similarly, depthwise asymmetric convolutions with dilation rates of (1, 3) and (3, 1) are applied to the input Zin, as defined in Eqs. (9) and (10). The resulting dilated feature maps are then passed through a 1 × 1 pointwise convolution to produce a single-channel output, as shown in Eq. (11).

In the third branch of dilated asymmetric convolution, we pass the input Zin through the dilated asymmetric depthwise convolution of kernels 1 × 3 and 3 × 1 with a unidimensional padding of 5 and unidimensional dilation of 5 in each case in Eq. (12) and (13). The dilated output is subsequently passed through a 1 × 1 convolution kernel to obtain the channel-modulated output, which is further modulated to a single-channel output, Z 3 out , as expressed in Eq. (14).

The outputs of the three convolution branches are concatenated along the channel dimension to form the three-channel feature map Zcat, as defined in Eq. (15).

The concatenated output Zcat is then passed through the batch normalization layer to obtain the batch-normalized output Hmix BN in Eq. (16).

The batch-normalized output Hmix BN is then passed through a 1×1 convolution to increase the number of output channels to 4 times that of Hmix BN , resulting inHmix 1 , as shown inEq. (17).

The output of the 1×1 convolution layer Hmix 1 is then passed through a 3 × 3 depthwise convolution layer to obtain Hmix 2 , as shown in Eq. (18).

The output Hmix 2 is then passed through the Gelu activation layer to obtain Hmix Gelu in Eq. (19).

The output Hmix Gelu is then passed through a 1 × 1 convolution for channel modulation, resulting in Hmix 3 , as shown in Eq. (20). This is done to ensure that the extracted attention mask has dimensions similar to those of Zin.

The output Hmix 3 is then passed through a sigmoid layer to obtain the mask M in Eq. (21) which ranges from 0 to 1.

The attention mask M obtained is then multiplied elementwise with input Zin (Hadamard multiplication) in Eq. (22) and then added with input Zin. The output Zout contains scale-enhanced features while suppressing irrelevant activations.

3.3 Adaptive Gated Gradient Reversal Layer

AGGRL contains a channel-wise modulated gradient reversal layer, which is driven by a gating layer in the forward propagation.

In the forward propagation, this layer acts as an identity function with no effect on the input as in Eq. (23).

The gating parameter provides channel-wise modulation in backpropagation, which is calculated in the forward propagation. This is done by applying global average pooling on the input x along the spatial dimensions H and W in Eq. (24).

Further, the global averaged output is then passed through a gating layer, which consists of a multi-layer perceptron (MLP) layer followed by a ReLU activation (σrelu) layer, followed by another MLP layer. This gives us the raw gating output graw in Eq. (25).

The raw gating values graw are then passed through the sigmoid layer (σsigmoid) to obtain the gating mask (g) in Eq. (26). The g contains values between 0 to 1 and acts as a mask for the gradients in the backpropagation during training.

During forward propagation, this layer behaves as an identity mapping, yielding the gradient relationship between the input and output shown in Eq. (27). During backpropagation, the incoming gradients are multiplied by the negative gating mask, as defined in Eq. (28), enabling channel-wise gradient reversal and selective suppression of redundant feature responses.

3.4 Adaptive Singular Value Decomposition Loss

The adaptive SVD loss leverages the singular values obtained from the bottleneck feature representation. Let F be the feature output of the encoder that belongs to RB×C×H×W .

Each sample feature F[b], which belongs to RC×H×W in a batch of size B, is reshaped into RC×V , where V = H ×W, as shown in Eq. (29).

In Eq. (30), S contains the singular values arranged in descending order.

The energy ratio Rb(k) is computed using the cumulative squared singular values ∑ki=1 S2 i and the sum of all squared singular values ∑min(C,V ) j=1 S2 j as per the Eq. (31).

The smallest index kt that satisfies Eq. (32) is then determined from Rb. The threshold parameter τ in Eq. (32) is learnable and constrained to the range [0.5, 0.99] to ensure training stability.

The loss for each sample in Eq. (33) is computed as the cumulative sum of the squared singular values from index kt, determined using Eq. (32), to min(C, V ) according to the dimensions of S. This formulation penalizes smaller singular values, thereby suppressing redundant feature components and encouraging more compact feature representations. The final adaptive SVD loss is obtained by averaging the samplewise losses Lb over the batch size B, as shown in Eq. (34).

3.5 Boundary-Weighted Cross-Entropy Loss

The binary ground-truth mask G is defined by G ∈ {0, 1}H×W . To emphasize learning near object boundaries, we utilize the signed distance transform [Kervadec et al., 2019]. For each pixel (i, j), the distance transform is defined as

where dij,mn = √(i -m)2 + (j -n)2 denotes the Euclidean distance between pixels (i, j) and (m, n). The distance map DT(i, j) obtained from Eq. (35) is normalized,

In Eq. (36), ϵ is added in the denominator for stability purposes, and its value is set to 10-6 . A boundary-aware weight map is defined as

which assigns higher weights to pixels closer to the object boundary. The boundary-weighted binary cross-entropy loss LBWCE is then formulated as

Eq. (38) gives pixel-wise binary cross-entropy loss LBCE. In Eq. (39) B denotes the batch size, M and N denote the height and width of the image. Here, W (k) b denotes the boundary-aware weight map for the k-th sample, while ˆ P (k) and G(k) represent the predicted probability map and groundtruth mask, respectively. The weighting and loss computations are applied pixel-wise over spatial locations (i, j).

3.6 Final Objective Function

The final objective function is given by L. In this, we combine the weighted addition of losses from binary crossentropy, dice loss [Milletari et al., 2016], proposed adaptive SVD loss, and BWCE loss.

The dice loss is given in Eq. (40). While the trainable weight λθ for adaptive SVD loss is clamped from 0 to 0.01. The final objective function for our proposed TvaraNet architecture is given by Eq. (41). BCE loss regulates global region-level probabilistic calibration. BWCE loss introduces spatially adaptive weighting to emphasize boundary pixels. Their combination enables joint optimization of region consistency and boundary precision. The theoretical convergence of the final objective is given in the supplementary material.

4 Experimental Results

4.1 Experimental Setup

The AdamW optimizer with a learning rate of 0.0001, a batch size of 32, and a weight decay of 0.00001 is used in all experiments. We used the CosineAnnealingLR approach to schedule the learning rate. The model was trained on an NVIDIA A100 GPU for a total of 400 epochs. During both training and inference, all input images were resized to 256×256×3. During training, we applied augmentations including horizontal flip, vertical flip, and random rotation. Experimental results are reported using the metrics mIoU, Dice score (DSC), boundary intersection over union (B-IoU), and Hausdorff distance (HD95) [Taha and Hanbury, 2015; Crum et al., 2006; Ribera et al., 2018].

4.2 Dataset Description

ISIC 2017 [Codella et al., 2018]: The ISIC 2017 dataset is a skin lesion segmentation benchmark comprising 1,500 training images and 650 validation images, each accompanied by an expert-annotated binary segmentation mask.

ISIC 2018 [Codella et al., 2019]: The ISIC 2018 dataset comprises 2,694 dermoscopic images with corresponding ground-truth lesion masks. The dataset is divided into 1,886 training images and 808 validation images.

Spinal segmentation [Chu et al., 2015]: The spinal segmentation dataset comprises 897 spine images and corresponding annotated masks of the lower spine region. The dataset is divided into 819 training images and 78 validation images.

PH2 [Mendonc¸a et al., 2015]: The PH2 dataset is a dermoscopic image dataset consisting of 200 high-resolution images, each accompanied by a pixel-level lesion segmentation mask. The dataset is used for cross-dataset evaluation of the proposed model.

4.3 Results and Discussion

Quantitative Analysis

In Table 1, TvaraNet is compared with two categories of models. The first category comprises heavy models (parameters ≥ 1 million), including UNet [Ronneberger et al., 2015], UNetV2 [Peng et al., 2025], VM-UNet [Ruan et al., 2024], and HVMUNet [Wu et al., 2025]. The second category comprises recent lightweight baseline models (parameters ≤ 1 million), including LightM-UNet [Liao et al., 2024] and Ultralight VMUNet (U. VMUNet) [Wu et al., 2024]. On ISIC17, TvaraNet outperforms both lightweight and heavy models, as presented in Table 1. Specifically, it achieves superior performance in mIoU (77.76 vs. 76.24), Dice score (87.49 vs. 86.52), and B-IoU (57.56 vs. 55.87). A similar trend is observed on ISIC-18, where TvaraNet outperforms both lightweight and heavy models in terms of mIoU, Dice score, and B-IoU (51.85 vs. 49.96).

Despite being the lightest model at only 0.060 GFLOPs, TvaraNet consistently outperforms both lightweight and heavier segmentation architectures. This indicates that its gains stem from effective architectural design rather than model capacity. The consistent improvements across region and boundary-based metrics indicate robust representation learning and improved contour preservation. TvaraNet attains superior performance among lightweight models on the spinal segmentation dataset across all evaluation metrics. Nonetheless, it does not surpass heavier models in regionbased measures such as mIoU and Dice score, presumably due to the presence of multiple anatomical structures within a single spinal image, which may require greater model capacity. Despite its compact design of only 0.037 million parameters, TvaraNet surpasses heavier models in boundarysensitive evaluation. It attains a superior B-IoU (61.32 vs. 61.08), indicating better boundary delineation. The comparison of latency and inference speed between TvaraNet and other lightweight segmentation models is reported in Table 3. Additional comparison results with other lightweight models are provided in the supplementary material.

Cross-Dataset Evaluation

The model trained on the ISIC-18 dataset is evaluated on the PH2 dataset as part of cross-dataset benchmarking, as presented in Table 2. We observe that TvaraNet consistently outperforms not only lightweight models but also heavier models by a considerable margin. It achieves superior performance in mIoU (85.74 vs. 84.87), Dice score (92.32 vs. 91.81), and B-IoU (53.13 vs. 49.59). These results validate the model's generalization capability beyond the dataset on which it was trained.

Qualitative Analysis

Fig. 3, Fig. 4, and Fig. 5 present qualitative comparisons of the segmentation masks predicted by TvaraNet and the corresponding baseline models on the ISIC-17, ISIC-18, and Spinal segmentation datasets. On both ISIC-17 and ISIC-18, TvaraNet performs better than the baseline models, LightMUNet and Ultralight VMUNet. This is evidenced by more accurately delineated boundaries and fewer spurious activations in the attention maps near object boundaries, indicating the suppression of irregular feature activations. Similarly, in Fig. 5, TvaraNet demonstrates improved shape prediction along with reduced false-positive predictions in the presence of multiple anatomical structures. This highlights the model's effectiveness in handling multiple objects of interest within a single image.

Challenges in Prediction

TvaraNet faces challenges in predicting object masks when the lesion occupies a large portion of the image frame. This limitation primarily arises from the limited number of such samples in the training dataset. Another challenge is the generation of false-positive spurious regions in the predictions. This issue occurs because certain non-target regions exhibit visual characteristics similar to those of the medical object of interest, leading to incorrect predictions.

Ablation Study

In Table 4, the ablation study of TvaraNet trained and evaluated on ISIC-18 is presented. The experiment begins with the stage-wise incorporation of each component. Starting with only the PASMA module, the baseline performance improves mIoU (78.24 to 78.42) and DSC (87.79 to 87.91) with the addition of the DASMA module, which enhances the skip connections. However, performance gains in boundary-aware metrics are limited at this stage, suggesting that DASMA primarily enhances region-level representations. Furthermore, incorporating AGGRL improves performance by suppressing irrelevant channel-wise features during backpropagation. The addition of the adaptive SVD loss further improves mIoU (78.77 to 78.83), DSC (88.12 to 88.16), and B-IoU (51.56 to 52.63) by reducing the influence of non-dominant and irrelevant features in the bottleneck representation. Finally, the introduction of the BWCE loss achieves the best overall segmentation performance, yielding the highest mIoU (79.31) and Dice score (88.46) while reducing HD95 to 16.30. These results demonstrate that the proposed components contribute complementary benefits toward improving segmentation accuracy and boundary localization.

5 Conclusion

We present TvaraNet, an ultra-lightweight medical image segmentation network with 0.037M parameters and 0.060 GFLOPs that achieves boundary-consistent predictions. By integrating PASMA, DASMA, AGGRL, adaptive SVD loss, and BWCE loss, TvaraNet effectively mitigates boundary distortion, feature redundancy, and shape inconsistency. Extensive experiments and ablation studies demonstrate its effectiveness across multiple datasets while maintaining suitability for resource-constrained deployment. Future work will focus on reducing false-positive predictions and improving model interpretability.