DXM-TransFuse U-Net for Nerve Identification

• The paper introduces a dual-encoder, transformer-fusion U-Net that integrates co-registered RGB and birefringence data for enhanced intraoperative nerve identification.
• It employs cross-modal multi-head attention at the bottleneck to fuse modality-specific features, achieving a 2% Dice improvement over multimodal baselines.
• The architecture demonstrates real-time potential with sub-30ms inference while offering promising avenues for clinical adaptation despite dataset limitations.

The DXM-TransFuse U-Net is a dual-encoder, transformer-fusion variant of the classic U-Net architecture, purpose-built for automated nerve tissue identification from multimodal optical imaging. Specifically, it leverages co-registered RGB and Mueller-matrix polarimetric birefringence images, integrating them via cross-modal multi-head attention at the encoder–decoder bottleneck. This architecture is designed to extract and fuse complementary modality-specific features, targeting enhanced nerve segmentation and detection performance in intraoperative settings (Xie et al., 2022).

1. Architectural Design and Feature Fusion

The network structure consists of parallel encoder branches for each input modality—RGB and Jet-mapped birefringence (derived via Mueller matrix decomposition)—feeding into a shared decoder. Each encoder processes its modality through four down-sampling stages, where each stage comprises two $3{\times}3$ convolutional layers with ReLU activations and a $2{\times}2$ max-pooling operation (stride two). Feature map depths double at each down-sampling (e.g., $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$), and skip connections preserve intermediate features for later fusion with the decoder.

At the bottleneck, a Dual Cross-Modal Transformer Fusion block processes the deepest encoder features. Each modality's feature map $F^a, F^b \in \mathbb{R}^{d \times h \times w}$ is linearly projected to sequences of flattened tokens:

• $Q^a, K^a, V^a$ from $F^a$
• $Q^b, K^b, V^b$ from $F^b$, with $N = h \cdot w$ token length.

Cross-attention operations enable each modality to attend to the other: $$\mathrm{Attention}(Q^a, K^b, V^b) = \mathrm{softmax}\left(\frac{Q^a (K^b)^\top}{\sqrt{d_k}}\right) V^b,$$ and analogously for the reverse direction. Residual connections, layer normalization, and dropout are applied, followed by a position-wise feed-forward network using GeLU/ReLU activation. The two fused representations are then reconstituted to spatial grids, concatenated, and projected via a $2{\times}2$0 convolution to produce a unified bottleneck embedding.

The decoder then performs four stages of bicubic upsampling (factor two), each followed by $2{\times}2$1 convolution + ReLU, concatenation with the corresponding encoder skip feature map, and two further $2{\times}2$2 convolutions. The output is a one-channel probability mask, post-processed through a sigmoid activation.

2. Input Modalities, Data Preprocessing, and Augmentation

Input consists of pixel-wise aligned $2{\times}2$3 RGB images (three channels) and Jet-mapped birefringence images (three channels). The latter are computed from co-registered polarimetric data using Mueller matrix decomposition. Preprocessing includes channel-wise standardization (subtracting mean, dividing by training-set standard deviation).

Data augmentation for training comprises vertical and horizontal flips, random brightness/contrast shifts, {$2{\times}2$4}-multiple rotations, and Gaussian noise; geometric transforms are mirrored onto ground-truth segmentation masks. Five-fold cross-validation is implemented, distributing $2{\times}2$538 images per fold for validation and 150 for training.

3. Training Protocol and Loss Design

The loss function is a multi-task sum:

• Weighted binary cross-entropy (BCE):

$2{\times}2$6

where $2{\times}2$7 compensates for class imbalance.

• Edge-loss (Sobel-based):

$2{\times}2$8

where $2{\times}2$9 is a Sobel edge operator at orientation $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$0.

• The total loss is $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$1.

Stochastic gradient descent (SGD) is used as optimizer, with initial learning rate $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$2 and multi-step decay (learning rate reduced by $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$3 at $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$4 and $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$5 through 250 epochs). Batch size is 8 per modality, trained on four NVIDIA Tesla V100 SXM2 GPUs within PyTorch.

4. Dataset Construction and Evaluation Protocol

The dataset comprises 188 co-registered RGB/Jet-mapped birefringence image pairs from a porcine model, focusing on various neck nerve segments. Ground-truth nerve masks are annotated through marker-based watershed segmentation in OpenCV, refined by expert surgeons.

Evaluation employs:

• Segmentation metrics:
  • Dice coefficient: $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$6
  • $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$7 score: $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$8
• Detection metrics (image-level, nerve present/absent at Dice $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$9 0.5): accuracy, sensitivity (recall), specificity, precision, balanced accuracy.

5. Quantitative Results and Comparative Evaluation

Results are summarized in the table below:

Model	Dice (%)	$F^a, F^b \in \mathbb{R}^{d \times h \times w}$0 (%)
U-Net (Jet only)	$F^a, F^b \in \mathbb{R}^{d \times h \times w}$1	$F^a, F^b \in \mathbb{R}^{d \times h \times w}$2
Attention U-Net (Jet only)	$F^a, F^b \in \mathbb{R}^{d \times h \times w}$3	$F^a, F^b \in \mathbb{R}^{d \times h \times w}$4
Cross-Attention U-Net	$F^a, F^b \in \mathbb{R}^{d \times h \times w}$5	$F^a, F^b \in \mathbb{R}^{d \times h \times w}$6
Dual U-Net (late concat)	$F^a, F^b \in \mathbb{R}^{d \times h \times w}$7	$F^a, F^b \in \mathbb{R}^{d \times h \times w}$8
Co-Learn U-Net (Kumar et al.)	$F^a, F^b \in \mathbb{R}^{d \times h \times w}$9	$Q^a, K^a, V^a$0
DXM-TransFuse U-Net	$Q^a, K^a, V^a$1	$Q^a, K^a, V^a$2

The DXM-TransFuse U-Net achieves a $Q^a, K^a, V^a$3 improvement in Dice and $Q^a, K^a, V^a$4 in $Q^a, K^a, V^a$5 over the best multimodal baselines. Inference cost is $Q^a, K^a, V^a$653M parameters at approximately 29.3ms per image pair (NVIDIA V100). Ablation experiments indicate that cross-modal Transformer fusion consistently outperforms pure concatenation or co-learning modules, and that single-modality attention yields modest improvements, with maximum gains realized through multimodal Transformer-based fusion.

6. Clinical Relevance and Current Limitations

DXM-TransFuse U-Net demonstrates efficacy for real-time, noninvasive intraoperative nerve identification, obviating the need for disruptive electrical nerve stimulation. The fused RGB and birefringence representation improves specificity, discriminating nerves from similarly birefringent tissues such as tendon and collagen, and operates at sub-30ms frame rates suitable for surgical guidance. Limitations include the relatively small dataset (188 frames), restricting generalizability; a single transformer block at the bottleneck (future variants may use multi-scale cross-modal fusion); and augmentation/head-count hyperparameters optimized on limited data.

7. Future Directions

Planned extensions include increasing the dataset size for higher statistical generalization, developing multi-scale transformer fusion insertions within the U-Net, and re-tuning augmentation and architectural hyperparameters for broader anatomical diversity and clinical settings. Enhanced data acquisition and advanced cross-modal attention mechanisms are anticipated to further boost segmentation accuracy and clinical applicability.

For full architectural and experimental detail, refer to "DXM-TransFuse U-net: Dual Cross-Modal Transformer Fusion U-net for Automated Nerve Identification" (Xie et al., 2022).