LET-Net: locally enhanced transformer network for medical image segmentation

With the great development of deep learning, especially convolutional neural networks (CNNs), various CNN-based methods, such as U-Net [7], have significantly improved the performance of medical image segmentation. These approaches possess the popular U-shaped encoder–decoder structure. To further assist precise segmentation, a battery of innovative improvements based on encoder–decoder architecture has emerged [26‐30]. One direction is to design a new module for enhancing the encoder or decoder ability. For instance, Dai et al. [26] designed Ms RED network, which, respectively, employs a multi-scale residual encoding fusion module (MsR-EFM) and a multi-scale residual decoding module (MsR-DFM) in the encoder and decoder stages to improve skin lesion segmentation. In the work [27], a selective receptive filed module (SRFM) was designed to obtain suitable sizes of receptive fields, thereby boosting breast mass segmentation. Another direction is optimizing skip connection to facilitate the recovery of spatial information. UNeXt [28] proposed an encoder–decoder structure involving convolutional stages and tokenized MLP stages, achieving better segmentation performance while also improving the inference speed. However, these methods directly fuse unaligned features from different levels, which may hamper accuracy, especially for small objects. In this paper, we propose a powerful feature-aligned local enhancement module, which ensures that feature maps at adjacent levels can be well aligned and then explore substantial local cues to optimally enhance the discriminative details.

Feature alignment has drawn much attention and is now an active research topic in computer vision. Numerous researchers have devoted considerable effort to addressing this challenge [6, 31‐37]. For instance, SegNet [6] utilized max-pooling indices computed in the encoder to perform an upsampling operation in the corresponding decoder stage. Mazzini et al. [32] proposed a guided upsampling module (GUM) that generates learnable guided offsets to enhance the upsampling operation. IndexNet [33] built a novel index-guided encoder–decoder structure in which pooling and upsampling operators are guided by self-learned indices. AlignSeg [34] learned 2D transformation offsets by a simple learnable interpolation strategy to alleviate feature misalignment. Huang et al. [35] designed an FaPN framework consisting of feature alignment and feature selection modules, achieving substantial and consistent performance improvements on dense prediction tasks. SFNet [31] presented a flow alignment module that effectively broadcasts high-level semantic features to high-resolution detail features by its semantic flow. Our method shares a similar aspect with the work [31], in which efficient spatial alignment is achieved by learning offsets. However, unlike these methods, we further enhance discriminative representations by subtraction under the premise of aligning low-resolution and high-resolution features, which facilitates excavating imperceptible local cues related to small objects.

Attention-based algorithms have been developed to assist in segmentation. In general, attention mechanisms can be categorized into channel attention, spatial attention, and self-attention according to different focus perspectives. Inspired by the success of SENet [38], various networks [39‐41] have incorporated the squeeze-and-excitation (SE) module to recalibrate features by modeling channel relationships, thereby improving segmentation performance. K. Wang et al. [42] proposed a dual attention network (DANet), which combines position spatial attention and channel attention modules to capture rich contexts.

Additionally, Transformer networks based on self-attention have been popular in medical image segmentation [43‐48]. For instance, TransUnet [43] inserted Transformer layers between CNN-based encoder and decoder stages to model global contexts, achieving excellent performances in multi-organ and cardiac segmentation. Wu et al. [48] proposed FAT-Net with a dual encoder that is, respectively, based on CNNs and Transformers for skin lesion segmentation. However, the loss of local contexts may still hinder the prediction accuracy of Transformer-based methods. In this paper, we propose a feature-aligned local enhancement module and progressive local-induced decoder, which, respectively, emphasize local information and adaptively recover spatial information to improve predictions.

Although CNN-based methods have achieved great success in medical image segmentation, they have general limitations in modeling global contexts. In contrast, pyramid vision transformer (PVT) [19] inherits the advantages of both Transformer and CNN while proving to be more effective for segmentation. Thus, we choose PVT as the backbone to obtain global receptive fields and learn effective multi-scale features.

As shown in Fig. 2, the PVT-based encoder has four stages with a similar architecture. Each stage contains a patch embedding layer and multiple Transformer layers. Benefiting from its progressive shrinking pyramid and spatial-reduction attention strategy, the PVT-based encoder can produce multi-scale feature maps with fewer memory costs. Specifically, given an input image \(X \in {\mathbb {R}}{^{H \times W \times 3}}\), it produces features \(\left\{ {\textbf{E}}_{i}\left| {1 \le i \le 4} \right. \right\}\), in which \({\textbf{E}}_{i} \in {\mathbb {R}}^{{H /{{2^{i + 1}}}} \times W/2^{i + 1} \times {C_i}}\). Therefore, we obtain high-resolution detail features and low-resolution semantic features, which are beneficial for segmentation.

The powerful global receptive field of PVT-based encoder makes it challenging for our model to adequately capture critical local details. Although low-level features can provide some local context, directly transmitting them to the decoder via a simple skip connection is problematic, as this may introduce a large amount of irrelevant background information. As a solution, leveraging high-level features is an effective way, but one significant issue, i.e., feature alignment, should be fully considered in this procedure to prevent local contexts from being overshadowed by global contexts. To this end, we propose a feature-aligned local enhancement (FLE) module, in which informative detailed features are effectively captured under the premise of feature alignment, producing discriminative representation. The internal structure of FLE is illustrated in Fig. 3, and it consists of two steps: feature-aligned discriminative learning and local enhancement.

Feature-aligned discriminative learning Due to the information gap between semantics and resolution, feature representation is still suboptimal when directly upsampling high-level feature maps to guide low-level features. To obtain strong feature representations, more attention and effort should be given to position offset between low-level and high-level features. Inspired by previous work [31], we propose a feature-aligned discriminative learning (FDL) block that aligns adjacent-level features and further excavates discriminative features, leading to high sensitivity to small objects. Within FDL, two \(1 \times 1\) convolution layers are first employed to compress adjacent-level features (i.e., \({{\textbf{E}}_i}\) and \({{\textbf{E}}_{i - 1}}\)) into the same channel depth. Then, a semantic flow field is calculated by a \(3 \times 3\) convolution operation, as described in Eq. 1where \({f_{s \times s}}\left( \cdot \right)\) indicates \(s \times s\) convolution layer followed by batch normalization and a ReLU activation function, while \(\copyright\) and \(\textrm{U}\left( \cdot \right)\), respectively, represent concatenation and upsampling operation. Next, according to learned semantic flow \({\varDelta _{i - 1}}\), we obtain a feature-aligned high-resolution feature \({{\tilde{\textbf{E}}}_i}\) with semantic cues, Mathematicallywhere \(Warp( \cdot )\) indicates the mapping function, \({{\textbf{E}}_i}\) is a \({C_i}\) dimensional feature map defined on the spatial grid \({\varOmega _{i}}\) of the specific size \(\left( {{H/{{2^{i + 1}}}},{W /{{2^{i + 1}}}}} \right)\). Schematically as shown in Fig. 4, the warp procedure consists of two steps. In the first step, each point \({p_{i - 1}}\) on the spatial grid \({\varOmega _{i - 1}}\) is mapped to \({p_i}\) on low-resolution feature, which is formulated by Eq. 3It is worth mentioning that due to the resolution gap between the flow field and features (see Fig. 4), Eq. 3 contains a halved operation to reduce the resolution. In the second step, we adopt the differentiable bilinear sampling mechanism [49] to approximate the final feature \({\tilde{\textbf{E}}}_i\) by linearly interpolating the scores of four neighboring points (top-right, top-left, bottom-right, and bottom-left) of \({p_i}\).

After that, to enhance the discriminative local context representation, we further utilize subtraction, absolute value, and residual learning procedures. Conclusively, the final optimized feature \({{\hat{\textbf{E}}}_{i - 1}}\) can be expressed as follows:Local enhancement In the PVT-based encoder, attention is established between each patch, allowing information to be blended from all other patches, even if their correlation is not high. Meanwhile, since small targets only occupy a portion of the entire image, the global interaction in transformer architecture cannot fully meet the requirements of small target segmentation where more detailed local contexts are needed. Considering that the convolution operation with a fixed receptive field can blend the features of each patch’s neighboring patches, we construct a local enhancement (LE) block to increase the weights associated with adjacent patches to the center patch using convolution, thereby emphasizing the local features of each patch.

As shown in Fig. 3, LE has a convolution-based structure and consists of four stages. Each stage includes a \(3 \times 3\) convolutional layer followed by batch normalization and a ReLu activation layer (denoted as \({f_{3 \times 3}}( \cdot )\)). Additionally, dense connections are added to encourage feature reuse and strengthen local feature propagation. As a result, the feature map obtained by LE contains rich local contexts. Let \({x_0}\) denote the initial input, and the outputs of \({i^{th}}\) stage within LE can be formulated as follows:where \([\ ]\) represents the concatenation operation. In summary, LE utilizes the local receptive field of the convolution operation and dense connections to achieve local enhancement.

Efficient recovery of spatial information is critical in medical image segmentation, especially for small objects. Inspired by previous works [50, 51], we propose a progressive local-induced decoder to adaptively restore feature resolution and detailed information. As shown in Fig. 2, the decoder consists of three cascaded local reconstruction and refinement (LRR) modules. The internal structure of LRR is illustrated in Fig. 5, where two steps are performed: local-induced reconstruction (LR) and split-attention-based refinement (SAR).

Local-induced reconstruction LR aims to transfer the spatial detail information from low-level features into high-level features, thereby facilitating accurate spatial recovery of high-level features. As shown in Fig. 5, LR first produces a reconstruction kernel \(\kappa \in {\mathbb {R}}{^{{k^2} \times {H_{i - 1}} \times {W_{i - 1}}}}\) based on low-level feature \({{\textbf{F}}_{i - 1}}\) and high-level feature \({{\textbf{D}}_{i}}\), in which \(k\) indicates the neighborhood size for reconstructing local features. The procedure of generating the reconstruction kernel \(\kappa\) can be expressed as follows:where \({f_{s \times s}}\left( \cdot \right)\) represents an \(s \times s\) convolution layer followed by batch normalization and a ReLU activation function. \(\textrm{U}\left( \cdot \right)\), \(\copyright\), and \(Soft\left( \cdot \right)\), respectively, indicate upsampling, concatenation, and Softmax activation operations. Meanwhile, another 3 \(\times\) 3 convolution and upsampling operation are applied on \({{\textbf{D}}_i}\) to obtain \({{\tilde{\textbf{D}}}_i}\) with the same resolution size as \({{\textbf{F}}_{i - 1}}\). Mathematically, \({{\tilde{\textbf{D}}}_i} = \textrm{U}\left( {{f_{3 \times 3}}\left( {{{\textbf{D}}_i}} \right) } \right)\). Note that, \({{\textbf{D}}_{4 }}={{\textbf{E}}_4}\) here. Next, we optimize pixel \({{\tilde{\textbf{D}}}_i}\left[ {u,v} \right]\) under the guidance of reconstruction kernel \({\kappa _{\left[ {u,v} \right] }} \in {\mathbb {R}}{^{k \times k}}\), producing refined local feature \({{\hat{\textbf{D}}}_i}\left[ {u,v} \right]\). This can be written as Eq. 7, where \(r = \left\lfloor k/2 \right\rfloor\)Subsequently, \({{\hat{\textbf{D}}}_i}\) and \({{\textbf{F}}_{i - 1}}\) are concatenated together and then passed through two convolutional layers to produce an optimized feature. Conclusively, LR overcomes the limitations of traditional upsampling operations in precisely recovering pixel-wise prediction, since it takes full advantage of low-level features to adaptively predict reconstruction kernel and then effectively combines semantic contexts with spatial information toward accurate spatial recovery. This can strengthen the recognition of small objects.

Split-attention-based refinement To enhance feature representation, we implement an SAR block in which grouped sub-features are further split and fed into two parallel branches to capture channel dependencies and pixel-level pairwise relationships through two types of attention mechanisms. As shown in Fig. 5, SAR is composed of two basic components: a spatial attention block and a channel attention block. Given an input feature map M, SAR first divides it along the channel dimension to produce \(M = \left\{ {{M_1},{M_2}, \cdot \cdot \cdot ,{M_G}}\right\}\). For each \({M_i}\), valuable responses are specified by attention mechanisms. Specifically, \({M_i}\) is split into two features, denoted as \({M}_i^1\) and \({M}_i^2\), which are separately fed into the channel attention block and spatial attention block to reconstruct features. This allows our model to focus on “what” and “where” are valuable through these two blocks.

In channel attention block, global average pooling (denoted as \(GAP(\cdot )\)) is performed to produce channel-wise statistics, which can be formulated asThen, channel-wise dependencies are captured according to the guidance of a compact feature, which is generated by a Sigmoid function (i.e., \(Sig( \cdot )\)). Mathematicallyin which parameters \({W_1}\) and \({b_1}\) are used for scaling and shifting S.

In spatial attention block, spatial-wise statistics are calculated using Group Norm (GN) [52] on \(M_i^2\). The pixel-wise representation is then strengthened by another compact feature calculated by two parameters \({W_2}\) and \({b_2}\) and a Sigmoid function. This process can be expressed asNext, \(\tilde{M}_i^1\) and \(\tilde{M}_i^2\) are optimized by an additional consistency embedding path and then concatenated. This procedure is represented asAfter aggregating all sub-features, a channel shuffle [53] is performed to facilitate cross-group information exchange along the channel dimension.

As stated in the previous study [54], training models with only pixel-wise loss may limit segmentation performance, especially resulting in prediction errors for small objects. This is due to class imbalance between foreground and background, such that task-relevant information is overwhelmed by irrelevant noise. Therefore, to facilitate preserving task-relevant information, we explore novel supervision at the feature level to further assist accurate segmentation. Let X and Y denote the input medical image and its corresponding ground truth, respectively. Z represents the deep feature extracted from input X.

Mutual information (MI) Mutual information is a fundamental quantity that measures the amount of information shared between two random variables. Mathematically, the statistical dependency of Y and Z can be quantified by MI, which is expressed aswhere \(p\left( {Y,Z} \right)\) is the probability distribution between Z and Y, while \(p\left( Z \right)\) and \(p\left( Y \right)\) are their marginals.

Mutual Information Loss Our primary objective is to maximize the amount of task-relevant information about Y in the latent feature Z while reducing irrelevant information. This is achieved by two mutual information terms [55, 56]. Formally Owing to the notorious difficulty of the conditional MI computations, these terms are estimated by existing MI estimators [56, 57]. In detail, the first term is accomplished through the use of Pixel Position Aware (PPA) loss [57] (\({L_{\textrm{PPA}}}\)). Since PPA loss assigns different weights to different positions, it can better explore task-relevant structure information and give more attention to important details. The second term is estimated by Variational Self-Distillation (VSD) [56] (\({L_{\textrm{VSD}}}\)) that uses KL-divergence to compress Z and remove irrelevant noises, thereby addressing the effect of imbalances in the number of foreground and background pixels caused by small targets. Thus, our total loss can be expressed as

To investigate the effectiveness of our proposed method, we validate LET-Net in two applications: polyp segmentation from coloscopy images and breast lesion segmentation from ultrasound images.

In this section, we conduct a series of ablation studies to verify the effectiveness of each critical component in our proposed LET-Net, including FLE, LRR, and mutual information loss.