LambdaUNet for 2.5D DWI Segmentation

• The paper introduces LambdaUNet, a novel 2.5D segmentation model that replaces standard convolutions with Lambda⁺ layers to disentangle intra-slice and inter-slice feature aggregation.
• LambdaUNet integrates efficient 1×1 convolutions, local 2D windows, and sparse inter-slice operations to address the challenges of anisotropic DWI data in acute ischemic stroke assessment.
• Experimental results on clinical DWI datasets show that LambdaUNet outperforms conventional 2D and 3D models, achieving up to an 86.51% DSC and improved lesion delineation.

LambdaUNet is a neural network architecture explicitly designed for segmentation of highly-discontinuous 2.5D medical images, particularly diffusion-weighted MR images (DWI) used in acute ischemic stroke assessment. LambdaUNet extends the standard U-Net framework by replacing convolutional layers with Lambda⁺ layers, a novel module that separately and efficiently models intra-slice and inter-slice context relevant for volumetric data exhibiting substantial inter-slice discontinuities. LambdaUNet is demonstrated to provide significant segmentation performance improvements over both conventional 2D and 3D models on clinical DWI datasets (Ou et al., 2021).

1. Motivation: Segmentation Challenges in 2.5D DWI

DWI studies are typically acquired as stacks of 2D slices characterized by high in-plane resolution (e.g., $1 \times 1$ mm²) but much greater slice thickness (e.g., 6 mm) and often with inter-slice gaps. This anisotropy results in spatially dense anatomical information within slices but sparse, often discontinuous structures across slices. Standard 2D convolutional networks (CNNs) such as U-Net, Attention U-Net, and TransUNet process slices independently and fail to leverage inter-slice information. Conversely, full 3D CNNs aggregate context isotropically, which can introduce noise and result in overfitting by incorporating irrelevant or misleading inter-slice data. DWI therefore benefits from an intermediate or “2.5D” approach: (a) dense and translation-equivariant feature extraction within slices, and (b) explicit, sparse modeling of inter-slice context.

LambdaUNet addresses these requirements with Lambda⁺ layers that disentangle and regulate intra-slice and inter-slice information propagation, ensuring robust segmentation in the presence of abrupt slice-to-slice variations (Ou et al., 2021).

2. Lambda⁺ Layer: Mathematical Formulation and Mechanisms

Given input volume $I \in \mathbb{R}^{T \times H \times W \times C}$ ($T$ slices, $H \times W$ per slice, $C$ channels), each pixel $n \in \mathcal{P}$ ($|\mathcal{P}|=T \cdot H \cdot W$) is assigned a feature vector $x_n$. The Lambda⁺ layer maps $X \in \mathbb{R}^{|\mathcal{P}| \times C}$ to $Y \in \mathbb{R}^{|\mathcal{P}| \times V}$ using three principal steps:

1. Linear Projections:

Compute queries, multi-depth keys, and values: $I \in \mathbb{R}^{T \times H \times W \times C}$0, $I \in \mathbb{R}^{T \times H \times W \times C}$1, $I \in \mathbb{R}^{T \times H \times W \times C}$2, where $I \in \mathbb{R}^{T \times H \times W \times C}$3, $I \in \mathbb{R}^{T \times H \times W \times C}$4, $I \in \mathbb{R}^{T \times H \times W \times C}$5. Keys are normalized spatially with softmax.

2. Lambda Construction:

For each pixel $I \in \mathbb{R}^{T \times H \times W \times C}$6, - Global within-slice lambda ($I \in \mathbb{R}^{T \times H \times W \times C}$7): Aggregates features from all pixels in the same slice. - Local 2D lambda ($I \in \mathbb{R}^{T \times H \times W \times C}$8): Aggregates features from an $I \in \mathbb{R}^{T \times H \times W \times C}$9 local window within the slice, with learned relative-position weights. - Inter-slice (sparse 3D) lambda ($T$0): Aggregates features from the same $T$1 position in neighboring slices (excluding the current slice), with learned inter-slice weights.

Each lambda is a linear map in $T$2.

3. Query-Lambda Product:

The output $T$3. This produces pixel-level 2.5D features for downstream decoding.

Lambda⁺ layers are implemented efficiently using 1x1 convolutions, softmax, 2D and 3D convolutions for local and inter-slice aggregation, and batched tensor operations (Ou et al., 2021).

3. LambdaUNet Architecture

LambdaUNet follows the canonical encoder–decoder U-Net topology with skip connections, with key modifications:

• Encoder: Four downsampling levels, each composed of a Lambda⁺ layer (increasing latent dim: $T$4), R=3 local window, and inter-slice context over $T$5 slices, followed by $T$6 max-pooling.
• Bottleneck: Single Lambda⁺ layer at $T$7 dimensions.
• Decoder: Four upsampling stages, each using $T$8 transposed convolution, concatenation with corresponding (reshaped) encoder features, $T$9 convolution, ReLU, and batch normalization.
• Final Output: $H \times W$0 convolution followed by sigmoid activation yields per-voxel lesion probabilities.

Slices are batched for Lambda operations; the decoder remains slice-wise (2D). Slices are recombined only at inference (Ou et al., 2021).

4. Experimental Results and Comparative Performance

LambdaUNet was evaluated on a clinical dataset of 99 acute ischemic stroke patients, stratified by lesion size and vascular territory, with manual expert annotations on DWI ($H \times W$1 s/mm²) and eADC. Training/validation/test splits comprised 67/32 cases with threefold cross-validation; input channels are DWI and eADC.

Preprocessing included per-slice intensity normalization and random sampling of 8-slice segments ($H \times W$2). Training used binary cross-entropy loss, RMSProp ($H \times W$3, linear decay), batch size 12, aggregated over four Quadro RTX 6000 GPUs. Convergence was achieved in $H \times W$44 hours ($H \times W$540 epochs).

Quantitative results on held-out test folds are summarized below:

Method	Dim	DSC (%)	Recall / Precision (%)	F₁ (%)
U-Net	2D	82.15	80.28 / 86.29	81.61
Attention U-Net	2D	81.83	77.45 / 86.74	80.82
TransUNet	2D	83.45	83.24 / 87.15	84.48
3D U-Net	3D	78.20	83.54 / 78.39	78.21
LambdaUNet-2D	2D only	84.03	82.27 / 87.10	84.19
LambdaUNet-3D	full 3D	84.76	79.92 / 89.86	84.09
LambdaUNet (Ours)	2.5D	86.51	81.76 / 89.39	84.84

LambdaUNet outperformed all baselines by up to 8 DSC points. Qualitative analysis revealed that LambdaUNet captured irregular lesion boundaries and maintained cross-slice consistency; 2D methods missed inter-slice cues, and 3D methods oversegmented by overfitting noisy inter-slice context (Ou et al., 2021).

5. Implementation Details and Computational Profile

• Framework: PyTorch with PyTorch Lightning.
• Core Operations: 1×1 convolution for projections; softmax for key normalization; 2D and 3D convolutions for lambda aggregation; batched einsum and reshaping for efficiency.
• Parameter Count: Comparable to standard U-Net of similar width ($H \times W$6–$H \times W$7M parameters).
• FLOPs & Runtime: Approximately 4 hours for training (4 RTX 6000 GPUs); inference per 8-slice volume takes $H \times W$8 seconds on a single GPU.
• Official Code: Released at https://github.com/YanglanOu/LambdaUNet (Ou et al., 2021).

6. Limitations and Areas for Extension

LambdaUNet was specifically tuned for DWI data, whose anisotropic slice structure defines the 2.5D regime. Its generalizability to isotropic modalities (T1, T2) or other organ systems remains untested. While Lambda⁺ layers provide a computationally more efficient mechanism than 3D convolutions, they still incur extra overhead relative to pure 2D networks; further reduction of memory and computation, possibly via lightweight lambdas or hybrid attention, is a potential direction. The default choice for inter-slice lambda aggregation ($H \times W$9) provided optimal accuracy on this dataset, but adaptive or dynamic $C$0 may improve robustness to heterogeneity in other clinical protocols. Integrating vision-transformer encoders (as in TransUNet) with LambdaUNet’s 2.5D lambdas is proposed as a direction to enhance segmentation of small lesions (Ou et al., 2021).

7. Broader Context and Research Impact

LambdaUNet provides a principled framework that bridges the performance gap between 2D and 3D CNN-based segmentation for non-isotropic, slice-based volumetric data. Its explicit architectural disentanglement of intra-slice and inter-slice features constitutes a general approach for other forms of sparse 3D or pseudo-3D image data prevalent in biomedical imaging. LambdaUNet’s release and performance on acute stroke lesion data positions it as a foundation for further development in 2.5D segmentation, exploration of adaptive context modeling, and integration with transformer-based architectures (Ou et al., 2021).