Vision Transformer, Mixture-of-Head attention, Adaptive Feature Extraction, Window-level Routing

1 Introduction↩︎

Perineural invasion (PNI), the insidious infiltration of cancer cells along nerve sheaths, is a critical route of metastasis in cholangiocarcinoma. Its presence significantly escalates the risk of recurrence, correlates with poor survival, and dictates surgical planning  [1]–[3]. Accurate preoperative identification of PNI is therefore paramount for personalized treatment.

Despite its clinical urgency, non-invasive PNI prediction remains challenging due to subtle MRI features [4], [5]. Furthermore, most PNI studies are hampered by small cohorts, often numbering in the low hundreds [2], [6], which frequently leads to class imbalance.

Methodologically, existing approaches often rely on radiomic features [7], which may fail to capture complex 3D spatial patterns. While end-to-end deep learning offers potential, standard architectures struggle. CNNs [8], [9] are limited in modeling long-range dependencies, while Vision Transformers (ViT) [10], [11] and their 3D adaptations [12] incur prohibitive computational costs in 3D MRI.

Hierarchical approaches like Swin Transformer [13] improve efficiency by computing attention within local windows. Building on this, architectures incorporating parallel multi-scale window processing, such as MViT [14] or Focal Transformers [15], have emerged to explicitly capture features at different scales simultaneously. However, the uniform processing of standard Multi-Head Self Attention (MSA) in these methods limits adaptive feature selection and causes attention head redundancy.[16]–[18].

The Mixture-of-Head attention (MoH) mechanism [16], inspired by Mixture-of-Experts (MoE) principles [19], [20], addresses this redundancy by treating attention heads as experts. MoH employs a router to dynamically select a subset of specialized heads for each input. This enhances specialization and efficiency. However, MoH typically operates at the token level, which remains computationally intensive for 3D data.

We propose the Multi-window Mixture-of-Head Attention Transformer (MMA-Former), an adaptive 3D architecture that synergizes multi-scale processing with a novel adaptation of MoH. We introduce the Coarse-Fine Transformer (CFT) structure to capture parallel multi-scale information. Crucially, we integrate Window-Specific MoH (WS-MoH). Instead of routing every token, WS-MoH dynamically routes the entire window to specialized heads. This enables spatially adaptive feature extraction where different regions leverage different attention heads efficiently.

Our main contributions are these:

• MMA-Former, a novel 3D end-to-end architecture utilizing a CFT structure for parallel multi-scale feature extraction for PNI prediction.

• WS-MoH, a novel approach applying MoH routing at the window level, enabling adaptive feature extraction based on the spatial context of the window.

• Xpace Block, an encoder utilizing hierarchical feature fusion between different stages.

2 Methodology↩︎

2.1 Datasets and Preprocessing↩︎

We utilized a retrospective dataset of anonymized T1-weighted, contrast-enhanced MR images (hepatobiliary phase) acquired from multiple MRI scanners at Samsung Medical Center over a decade. Images were provided in NIfTI format with ground truth annotations for the liver and tumor. After quality control, the final analysis cohort comprised 168 cholangiocarcinoma patients. The presence of PNI was confirmed by post-surgical histopathological examination, comprising 67 PNI-positive cases and 101 PNI-negative cases.

We employed a localization strategy, extracting cropped volumes of size \(96\times96\times48\) centered on the tumor. The necessity of this localization is validated in Table 2 (A).

2.2 MMA-Former Architecture↩︎

The MMA-Former, illustrated in Fig. 2 (a), employs a hierarchical design. Input 3D volumes undergo patch partitioning and linear embedding. The architecture consists of an encoding phase followed by \(K\) stages. Each stage comprises patch merging layer and an MMA block for feature extraction. An Xpace block, shown in Fig. 2 (c), fuses hierarchical features across stages by integrating downsampled representations from the previous stage with features from the current stage.

2.3 The MMA Block: CFT with WS-MoH↩︎

2.3.1 CFT Structure↩︎

The Coarse-Fine Transformer (CFT) structure captures multi-scale information by processing features through parallel pathways with different window sizes.

Given an input feature map \(z^{l-1}\) at sub-block \(l\), we apply layer normalization (LN). The normalized features \(z'\) are processed in parallel through two pathways. The Fine path, denoted as \(P_f\), utilizes a small window size \(W_F\) (e.g., \(3\times3\times3\)). The Coarse path, denoted as \(P_c\), employs a larger window size \(W_C\) (e.g., \(6\times6\times6\)). Both pathways use WS-MoH as the attention mechanism. \[\label{eq:cft95parallel} A_F = \text{WS-MoH}_{W_F}(z'), \quad A_C = \text{WS-MoH}_{W_C}(z')\tag{1}\] The outputs are fused via concatenation and linear projection (Fusion), then combined with the input via residual connections, followed by an MLP: \[\label{eq:cft95aggregate} A_{fused} = \text{Fusion}(A_F, A_C) + z^{l-1}\tag{2}\] \[\label{eq:cft95mlp} z^{l} = \text{MLP}(\text{LN}(A_{fused})) + A_{fused}\tag{3}\] Subsequent blocks (Block \(l+1\)) employ shifted window partitioning (SWS-MoH) [13] for cross-window communication (Fig. 2 (b)).

2.3.2 WS-MoH↩︎

We replace standard MSA with the adaptive Window-Specific MoH (WS-MoH), detailed in Fig. 3. MSA (Eq. 4 ) activates all \(h\) heads uniformly, regardless of the input context. \[\label{eq:msa} \text{MSA}(X) = \sum_{i=1}^{h} H^i(X) W_O^i\tag{4}\]

MoH [16] introduces a router to dynamically weight heads, denoted as \(g_i\). We propose WS-MoH, which adapts this routing to the window level for efficiency in 3D data.

First, standard Q, K, V projections are computed for the input features within a window \(X_W\). \[\label{eq:qkv} Q = X_W W_Q, \quad K = X_W W_K, \quad V = X_W W_V\tag{5}\]

Crucially, instead of routing each token, we generate a window representative feature \(X_{rep}\) by average pooling the features \(X_W\).

The router uses this single representative feature \(X_{rep}\) to determine the routing scores \(g_i\) for the entire window. The final WS-MoH output is the weighted sum of the outputs from the selected heads: \[\label{eq:w-moh} \text{WS-MoH}(X_W) = \sum_{i=1}^{h} g_i(X_{rep}) H^i(X_W) W_O^i\tag{6}\]

This enables spatially adaptive feature extraction, as different windows activate different combinations of heads based on their local context.

We employ the two-stage routing strategy [16]. Heads are divided into Shared Heads, denoted as \(H_S\), which capture common knowledge and are always active, and Routed Heads, denoted as \(H_R\), which handle specialized patterns.

A router network processes \(X_{rep}\) to produce probabilities for the routed heads, and the Top-K heads are selected. \[\label{eq:router95probs} P_R = S(W_{r}X_{rep})\tag{7}\]

The routing score \(g_i\) (Eq. 8 ) is determined by balancing the contributions of \(H_S\) and the Top-K selected \(H_R\) via coefficients \(\alpha_1, \alpha_2\), where \([\alpha_{1},\alpha_{2}] = S(W_{h}X_{rep})\) and \(S(\cdot)\) denotes the Softmax function. \[\label{eq:routing} g_{i}(X_{rep})=\begin{cases} \alpha_{1} S(W_{s}X_{rep})_{i}, & \text{if } i \in H_S \\ \alpha_{2} (P_R)_{i}, & \text{if } i \in \text{Top-K}(H_R) \\ 0, & \text{otherwise}. \end{cases}\tag{8}\]

2.3.3 Load Balance Loss↩︎

In mixture models utilizing routing, training often leads to an imbalance where a few heads process the majority of inputs, leaving others undertrained [19]. To mitigate this in WS-MoH and ensure effective utilization of all Routed Heads for robust PNI prediction, we incorporate a load balance loss \(\mathcal{L}_{LB}\). This loss encourages an even distribution of windows across the Routed Heads. \[\label{eq:load95balance} \mathcal{L}_{LB} = \sum_{i \in H_R} P_i \cdot f_i\tag{9}\] where \(P_i\) is the average routing probability for head \(i\), and \(f_i\) is the fraction of windows that selected head \(i\).

2.4 Xpace Block↩︎

The Cross-Spatial Attention (Xpace) block fuses features across encoder stages. First, it downsamples previous-stage features. The fusion mechanism operates through bidirectional cross-attention between downsampled and current-stage features. Spatial attention is applied to highlight important regions in both representations. The resulting features are then concatenated along the channel dimension and forwarded to the next stage.

3 Experiments and Results↩︎

3.1 Experimental Setup↩︎

All experiments were performed on NVIDIA A100 GPUs. We employed stratified 5-fold cross-validation across the 168 cases. Models were trained using the AdamW optimizer [22] with a learning rate of \(8\times10^{-5}\) and a batch size of 4. The total loss function combines the Weighted BCE Loss, denoted as \(\mathcal{L}_{WCE}\), utilizing a positive weight of 1.5 to address class imbalance, and the load balance loss \(\mathcal{L}_{LB}\), weighted by a factor \(\beta=0.01\). \[\label{eq:total95loss} \mathcal{L}_{Total} = \mathcal{L}_{WCE} + \beta\mathcal{L}_{LB}\tag{10}\]

3.2 PNI Prediction Performance↩︎

Table 1 summarizes the performance comparison. MMA-Former achieved the highest mean AUC of 0.752. It significantly outperformed the best-performing CNN, ResNet (which achieved an AUC of 0.708), and the standard 3D Swin Transformer (which achieved an AUC of 0.681).

3.3 Ablation Study↩︎

We conducted ablation studies to validate the design choices of MMA-Former (Table 2). The default configuration uses 2 Shared heads (S=2) and a Top-K ratio of 75%.

Impact of Input Strategy (A): Training MMA-Former on uncropped images decreased the AUC by 0.047. This confirms that localization to the tumor region is crucial for detecting the subtle features of PNI, validating our preprocessing strategy.

Impact of Core Components (B): Replacing WS-MoH with standard MSA decreased the AUC by 0.021. This highlights that the adaptive, window-specific feature extraction provided by WS-MoH contributes to performance. Removing the CFT structure also reduced performance, resulting in an AUC of 0.690. Additionally, removing the Xpace Block resulted in degraded performance, lowering the AUC to 0.709.

Impact of CFT Configuration (C): The parallel configuration achieved the best performance, outperforming both the Fine path only (AUC of 0.704) and the Coarse path only (AUC of 0.718). This validates the CFT structure, confirming that synergizing local and global contexts outperforms single-scale approaches.

Impact of WS-MoH Shared Heads (D): The configuration with 2 Shared heads provided the optimal balance. Relying solely on Routed heads failed to capture common features effectively, while using only Shared heads lacked adaptive specialization.

Impact of WS-MoH Top-K Ratio (E): Activating 75% of the routed heads yielded the best result. Using fewer heads reduced capacity, while using more likely reintroduced redundancy.

3.4 Qualitative Results and Visualization↩︎

We utilized 3D Grad-CAM [23] to visualize the regions most influential to the model’s predictions (Fig. 1). The visualizations confirm that MMA-Former focuses precisely on the tumor and the immediate peritumoral interface, the areas most critical for identifying PNI.

4 Discussion and Conclusion↩︎

The MMA-Former introduces a novel approach to 3D medical image analysis by enabling spatially adaptive feature extraction within a hierarchical transformer framework. The effectiveness of the approach is strongly supported by the necessity of input localization.

The core innovation is the WS-MoH within the CFT framework. The CFT structure provides multi-scale context. WS-MoH significantly outperforms CFT+MSA via window-level adaptation. Dynamically routing window representations to specialized heads captures nuanced local context. This contrasts with MSA’s uniform processing and avoids the computational burden of token-level MoH.

The primary limitation is the single-institution nature of the dataset, comprising 168 cases. Although the cohort size is comparable to related PNI studies, external validation on multi-center datasets is required to ensure generalizability.

In conclusion, we proposed the MMA-Former, an adaptive 3D Transformer for PNI prediction. By synergizing the CFT structure with a novel WS-MoH, MMA-Former effectively captures multi-scale context while enabling efficient, spatially adaptive feature extraction. Our experiments demonstrate the superiority of MMA-Former (AUC of 0.752) over existing 3D backbones (best AUC of 0.708), highlighting the benefit of window-level adaptive attention for complex 3D medical imaging tasks.

5 Acknowledgments↩︎

This work was supported by the Institute of Information & Communications Technology Planning & Evaluation (IITP), funded by the Korea government (MSIT), under the Artificial Intelligence Semiconductor Support Program to nurture the best talents (IITP-2023- RS-2023-00256081) and the grant for the Development of an AI Deep Learning Processor and Module for a 2,000 TFLOPS Server (No. 2020-0-01305)

References↩︎