CR-MSA: Cross-Region Multi-head Self-Attention

• CR-MSA is a localized self-attention mechanism that re-embeds spatially contiguous image patches to enable effective cross-region feature fusion.
• It incorporates an Embedded Positional Encoding Generator (EPEG) via 1D convolutions to capture relative positional information without explicit encodings.
• As a key component of the R²T architecture, CR-MSA enhances performance in multiple instance learning by aggregating both local details and global context in digital pathology tasks.

The Re-embedded Regional Transformer (R²T) is an architectural framework for multiple instance learning (MIL) that targets the efficient and effective aggregation of local and global information in high-dimensional, large-scale datasets, particularly gigapixel whole-slide histopathology images. R²T leverages local regional self-attention, hierarchical or cross-region feature fusion, and region-focused inference to achieve superior classification and prognostic performance, especially in digital pathology and computational pathology tasks. It is designed to operate as a modular, online re-embedding block that can be integrated into standard MIL pipelines, offering consistent accuracy gains across diverse backbone architectures and datasets.

1. Structural Overview and Core Components

R²T builds on the MIL paradigm, which processes large bags of instance features (e.g., patches extracted from a whole-slide image) to produce a single bag-level/slide-level prediction using weak supervision. Its defining architectural sequence consists of:

• Instance Feature Extractor: An offline backbone, such as a ResNet-18 (Cersovsky et al., 2023) or ResNet-50 (Tang et al., 27 Feb 2024), is used to convert raw image patches $x$ of fixed spatial size (typically $224 \times 224$ pixels) into high-dimensional embeddings $h_i \in \mathbb{R}^D$.
• Regional Transformer Modules: Patches are organized into spatially contiguous, non-overlapping regions. For each region, embeddings are fused via local multi-head self-attention (regional MSA), optionally including a learnable class token for pooling.
• Hierarchical or Cross-Region Fusion: The outputs of regional modules can be hierarchically aggregated (stacked) (Cersovsky et al., 2023) to form coarse-scale region tokens or, in an alternative structure (Tang et al., 27 Feb 2024), via a lightweight cross-region MSA to propagate dependencies across more distant regions.
• Global Transformer Aggregator: All region-level tokens are globally aggregated using an additional Transformer block, whose class token produces the final slide representation.
• Slide-Level Classifier: A linear classification head transforms the global representation into logits, which may be passed through either a sigmoid function (for binary classification) or Cox loss (for survival analysis).

This modular construction enables R²T to replace conventional MIL aggregators while retaining global context and computational tractability. The architecture is distinguished by the lack of explicit positional encodings, relying instead on spatial partitioning and hierarchical structure.

2. Regional Multi-Head Self-Attention and Re-Embedding Strategy

R²T's central innovation is the regional, locally-constrained multi-head self-attention mechanism:

• Patch-to-Region Assignment: Given a set of embeddings $H = \{h_1, ..., h_I\} \in \mathbb{R}^{I \times D}$, these are reshaped into an approximately square grid of $L \times L$ regions, each region comprising $M \times M$ patches ($I \approx (L M)^2$).
• Regional MSA (R-MSA): Within each region $\ell$, the per-region subset $S^{\ell} \in \mathbb{R}^{M^2 \times D}$ undergoes LayerNorm and multi-head self-attention:

$$Q^\ell = S^\ell W^Q,\quad K^\ell = S^\ell W^K,\quad V^\ell = S^\ell W^V$$

$$e_{ij}^\ell = \frac{Q^\ell_i \cdot K^\ell_j}{\sqrt{d_k}}, \quad \widetilde{e}_{ij}^\ell = e_{ij}^\ell + \mathrm{Conv}_{1D}(e_{ij}^\ell)$$

$$\alpha_{ij}^\ell = \mathrm{softmax}_j(\widetilde{e}_{ij}^\ell), \quad R\!-\!\mathrm{SA}^\ell_i = \sum_{j=1}^{M^2} \alpha_{ij}^\ell V^\ell_j$$

These are concatenated and projected to produce region-wise re-embedded representations, and reassembled into the instance order via reshape.

• Residual Connections and LayerNorm: Each attention block is wrapped with LayerNorm and residual addition; crucially, feed-forward networks (FFNs) are omitted as their inclusion degrades performance and increases model size.
• Embedded Positional Encoding Generator (EPEG): R²T augments attention logits with a 1D convolution over the attention scores to incorporate relative position information, enabling the model to exploit local spatial coherence while maintaining translation invariance.

3. Cross-Region Aggregation and Hierarchical Stacking

Following the local re-embedding, R²T performs further cross-region context fusion:

• Cross-Region MSA (CR-MSA) (Tang et al., 27 Feb 2024):
  • Each region $\ell$ is summarized via a softmax-weighted pooling, producing $K$ representative vectors per region:

  $$W_a^\ell = \mathrm{softmax}_{p=1\ldots M^2} (\hat{Z}_p^\ell \Phi), \quad R^\ell = (W_a^\ell)^\top \hat{Z}^\ell \in \mathbb{R}^{K \times D}$$ - Concatenated regional representatives are fused via standard self-attention (MSA), then redistributed to each instance with suitable attention-weighted mapping using MinMax normalization and softmax. - This process enables propagation of information across distant, spatially separated regions.

• Hierarchical Regional Aggregation (Cersovsky et al., 2023):

  • Regional Transformer modules may be stacked in multiple levels: tokens from small-scale regions are recursively grouped and re-embedded via attention.
  • At each level $l$, tokens are assembled with a freshly learnable class token, and local window attention is repeated, exponentially increasing the spatial coverage up to the global scale. The recursion is:

  $$E^{(l)}_j = [c^{(l)}; r^{(l-1)}_{j,1}; \dots; r^{(l-1)}_{j,S^2}],\;\; r^{(l)}_j = [\mathrm{softmax}(Q^{(l)} K^{(l)\top}/\sqrt d)V^{(l)}]_0$$ - Optional combination of multi-level global tokens can be performed by either concatenation or summation.

4. Integration into Multiple Instance Learning Pipelines

R²T is designed as a plug-in, online re-embedding module. In a standard MIL workflow, its place is as follows:

1. Instance Feature Extraction: $H = F(X)$, where $X$ is the set of slide patches and $H$ the matrix of feature vectors.

2. Online Regional Re-embedding: $Z = R(H)$, with all R²T parameters ($W^Q, W^K, W^V, W^O, \Phi$) trained jointly.

3. Instance Aggregation: MIL aggregator $e = A(Z)$, e.g., attention pooling.

4. Bag Classification: $\hat{Y} = C(e)$, via a linear or Cox head.

Training is end-to-end, using cross-entropy loss for classification or Cox loss for survival analysis (Tang et al., 27 Feb 2024, Cersovsky et al., 2023), with Adam (or AdamW) optimizers and standard data augmentation protocols. Region size (L), representative count (K), and pooling windows (S) are robust hyperparameters.

5. Inference Strategies and Attention-based Region Selection

R²T supports an inference-time refinement targeting datasets with highly sparse or localized informative regions (e.g., small metastases):

• Two-Pass High-Attention Rerun (Cersovsky et al., 2023):

1. First Pass: Compute attention weights $$a_i = \exp(c_{\mathrm{glob}}^\top e_i) / \sum_j \exp(c_{\mathrm{glob}}^\top e_j)$$ for all instance embeddings using the global class token.
2. Selection: Fit a two-component 1D clustering (e.g., k-means or GMM) to $\{a_i\}$, designating each as ‘high’ or ‘low’ attention ($\mu_i \in \{0, 1\}$).
3. Second Pass: Set low-attention embeddings to zero ($\tilde{e}_i = \mu_i e_i$) and rerun the aggregation pipeline, focusing learning capacity only on high-saliency patches.

This technique increases prediction sharpness on slides where the label is determined by rare morphologies, as confirmed by empirical gains in CAMELYON16.

6. Empirical Benchmarks, Ablation Studies, and Computational Considerations

Extensive experiments on public pathology datasets validate R²T's contribution:

• Binary Classification Tasks:

  • On CAMELYON16 (lymph-node metastasis detection), using ResNet-50 + AB-MIL, R²T-MIL achieves mean AUC 97.32% (+2.78 pp over baseline); TCGA-BRCA 93.17% (+2.07 pp); TCGA-NSCLC 96.40% (+1.12 pp) (Tang et al., 27 Feb 2024).
  • With foundation model features (PLIP), R²T-MIL adds 0.26%–1.37% AUC.
• Survival Analysis:
  • On TCGA-LUAD, R²T-MIL improves C-index from 58.78% (baseline AB-MIL) by +8.41% to 67.19%.
• Comparative Performance (Region Strategies):
  • Regional MSA (local) outperforms global Nystrom-approximated MSA and standard vanilla Transformer approaches.
  • Inclusion of CR-MSA and EPEG further improves metric gains, especially on survival analysis.
• Ablations:
  • FFN blocks, positional encodings not only fail to raise performance but sometimes substantially degrade it.
  • The locality/region grid parameter L is robust, with best results at L=8 given patch counts per slide in tested datasets.
• Runtime and Memory:
  • Training R²T-MIL on whole slides (CAMELYON-16, single GPU) incurs modest computational costs: 6.5 s/epoch, 10.1 GB memory, surpassing TransMIL's 13.2 s/epoch and 10.6 GB, with ∼3× throughput and comparable parameter sizes.
  • Parameter overhead is moderate (+2.7M parameters over 26M backbone), offering favorable efficiency on large datasets compared to global attention methods.

Method	Epoch Time (s)	Memory (GB)	FPS
AB-MIL	3.1	2.3	1250
TransMIL	13.2	10.6	76
R²T-MIL (full)	6.5	10.1	236

7. Significance and Application Context

R²T demonstrates a systematic improvement to MIL accuracy and robustness by leveraging local regional attention and cross-region/global feature propagation, with negligible additional training and inference burden. Its plug-in nature enables application across established MIL architectures, providing consistent metric improvements whether features are sourced from standard CNN encoders or advanced vision-language foundation models. The approach is particularly beneficial in scenarios with patch-level label sparsity—common in digital pathology—where the two-pass attention strategy significantly sharpens discriminative signal localization.

A plausible implication is that, given its efficiency and representational power, R²T may generalize to domains beyond histopathology where data comes as unstructured sets with rich local correlations. The design choices, notably the exclusion of explicit positional encodings and FFNs, are empirically justified in this context and reflect the unique properties of MIL on high-resolution images.

Further research directions include transfer to higher-resolution feature extractors, adaptation to non-rectilinear region topologies, and extension to multi-modal data. The methodology is openly available and reproducible (Tang et al., 27 Feb 2024), supporting broad adoption and adaptation in computational pathology and beyond.