Re-embedded Regional Transformer (R²T)

• Re-embedded Regional Transformer (R²T) is a neural architecture that enhances multiple instance learning by fusing local patch-level features using regional and cross-region self-attention.
• It integrates an Embedded Position Encoding Generator (EPEG) and a hierarchical attention mechanism to efficiently capture both intra- and inter-region dependencies in high-resolution images.
• Empirical evaluations demonstrate notable AUC improvements and faster inference compared to standard MIL models, proving its efficacy in computational pathology.

The Re-embedded Regional Transformer (R²T) is a neural architecture developed for advancing multiple instance learning (MIL) in computational pathology and related high-resolution computer vision tasks. R²T is designed to enhance slide-level predictions from patch-level representations by fusing local context through regionalized Transformer-based self-attention, combining both intra-region and inter-region information. As a drop-in module, R²T enables any existing MIL model to achieve gains comparable to, or exceeding, those observed with foundation model features, operating efficiently even on gigapixel whole slide images (WSIs) (Cersovsky et al., 2023, Tang et al., 27 Feb 2024).

1. Architectural Overview

R²T operates on a set of instance features extracted from image patches. The dominant use-case involves WSIs where an upstream encoder (such as ResNet-50 or a vision-LLM like PLIP) provides frozen patch embeddings $H = \{h_1, ..., h_I\} \in \mathbb{R}^{I\times D}$, with $I$ patches and $D$-dimensional features. The R²T block consists of two main modules:

1. Regional Multi-Head Self-Attention (R-MSA):
   • Features are reshaped into a 2D grid, split into $L \times L$ non-overlapping regions of $M \times M$ patches ($L \cdot M \approx \sqrt{I}$).
   • Within each region, standard scaled dot-product multi-head self-attention is applied, using learned linear projections for queries, keys, and values. The local context is fused to yield a re-embedded region feature. An Embedded Position Encoding Generator (EPEG) is used, adding a learned 1D convolution to the attention logits before the softmax, leveraging relative spatial information.
   • After attention, output features are reshaped back to the original patch order.
2. Cross-Region Multi-Head Self-Attention (CR-MSA):
   • Each region is summarized into $K$ "representatives" using a linear attention mechanism ($$W_a^\ell = \operatorname{softmax}(\hat{Z}^\ell \Phi)$$, with $\Phi \in \mathbb{R}^{D \times K}$). The resulting representatives are stacked and passed through a standard Transformer MSA block, capturing non-local dependencies across regions.
   • The processed cross-region information is redistributed to the instance level using attention-weighted combinations and addition of residuals.

The overall architecture for a forward pass is summarized as:

• $$Z = (\operatorname{CR\text{-}MSA} \circ \operatorname{LN} \circ \operatorname{R\text{-}MSA} \circ \operatorname{LN})(H) + 2H$$.

No feed-forward sub-layers are used within R²T, as ablation shows their presence degrades performance and increases parameter count.

2. Integration with Multiple Instance Learning Pipelines

R²T is inserted immediately after feature extraction in standard MIL pipelines:

• $$X \xrightarrow{\operatorname{F}} H \xrightarrow{\operatorname{R}} Z \xrightarrow{\operatorname{A}} e \xrightarrow{\operatorname{C}} \hat{Y}$$

Here, $F$ is a frozen or pretrained patch encoder, $R$ is the R²T module, $A$ an MIL aggregator (e.g., attention pooling), and $C$ a classifier head. The addition of R²T enables end-to-end training of the aggregator and classifier with efficient contextualization of local and global patterns, all parameters of $R$ being learned jointly with $A$ and $C$ via standard losses (cross-entropy or Cox, depending on prediction target).

A standard forward pass (PyTorch-style) is as follows:

def forward(slide_patches): H = FeatureExtractor(slide_patches) # [I, D] Hn = LayerNorm(H) Z1 = regional_MSA(Hn) + H Z1n = LayerNorm(Z1) Z = cross_region_MSA(Z1n) + Z1 bag_embedding = Aggregator(Z) Y_hat = Classifier(bag_embedding) return Y_hat

R²T is architecturally agnostic to the choice of aggregator, and experimental results demonstrate robust improvements when integrated with diverse frameworks such as AB-MIL, CLAM, DSMIL, TransMIL, DTFD-MIL, IBMIL, and MHIM-MIL.

3. Inference and Performance Characteristics

R²T includes optional strategies to further sharpen predictions, particularly beneficial in tasks with sparse discriminative signals. In (Cersovsky et al., 2023), a two-pass inference is described:

• A first pass computes attention weights $a_i$ for each patch using the global class token.
• A clustering algorithm is used to select high-attention patches ($\mu_i$ binary mask).
• In the second pass, features of low-attention patches are zeroed ($\tilde{e}_i = \mu_i e_i$), and MIL aggregation is rerun.

This approach enhances discrimination in cases like small metastases (CAMELYON16), yielding up to 15-percentage-point AUC improvements over the global attention MIL baseline.

4. Empirical Evaluation and Ablation

R²T was evaluated across multiple public datasets with weak slide-level labeling, including CAMELYON-16 (metastasis detection), TCGA-BRCA (IDC vs. ILC), and TCGA-NSCLC (LUAD vs. LUSC), as well as survival analysis (C-index on TCGA-LUAD, -LUSC, -BLCA).

Key empirical results:

Dataset	Baseline (AB-MIL)	+R²T-MIL (ResNet-50)	+R²T-MIL (PLIP)
CAMELYON-16 AUC	94.54%	97.32% (+2.78%)	+1.37%
TCGA-BRCA AUC	91.10%	93.17% (+2.07%)	+0.26%
TCGA-NSCLC AUC	95.28%	96.40% (+1.12%)	+0.83%
LUAD C-index	58.78%	67.19% (+8.41%)	+0.43%

Ablations revealed:

• R²T (native per-region MSA + CR-MSA + EPEG) achieves greater AUC gains compared to global-only (TransMIL +1.28%), Nystrom-approximation (N-MSA +1.66%), and local N-MSA per region (+1.95%).
• The EPEG positional encoding is beneficial compared to alternatives (PEG₇×₇, PPEG); omitting it reduces performance.
• Adding feed-forward sublayers increases parameter count (+7–10M) and reduces performance.
• Region partitioning grid size L is robust; optimal at $L=8$ for typical slide sizes.

5. Computational Efficiency and Resource Considerations

R²T achieves its performance with moderate computational overhead:

• R-MSA complexity: $O(L^2 (M^2)^2 D)$ (quadratic in region size, avoiding full quadratic in patch count).
• CR-MSA complexity: $O((L^2 K)^2 D)$, with $K \ll M^2$.
• Additional parameter count: +2.70M over a 26M backbone.
• Runtime per epoch (CAMELYON-16, single GPU): AB-MIL: 3.1 s, TransMIL: 13.2 s, R²T-MIL: 6.5 s.
• R²T is approximately 3× faster per epoch and has a smaller memory footprint than TransMIL.

No special positional embeddings are used, as region membership and hierarchical structure already encode locality. All operations are differentiable and amenable to minibatch SGD.

6. Comparison to Related Regional Transformer Frameworks

Both (Cersovsky et al., 2023) and (Tang et al., 27 Feb 2024) propose regional Transformer methodologies for MIL; each takes a distinctive approach to regional aggregation:

• (Cersovsky et al., 2023) employs a hierarchical stacking of regional Transformer modules and a global Transformer at the top level, allowing for directly interpretable aggregation pathways from local regions to whole-slide predictions.
• (Tang et al., 27 Feb 2024) generalizes R²T as a portable re-embedding module, suitable for any instance set, and formalizes detailed R-MSA and CR-MSA operations with EPEG, demonstrating broad applicability and performance across multiple MIL benchmark models and backbone encoders.

7. Implementation Parameters and Training Protocols

Experimental setups for R²T universally employ standard training practices:

• Optimizers: AdamW or Adam, learning rate $2 \times 10^{-5}$ to $2 \times 10^{-4}$, cosine schedule, early stopping.
• Batch size: 1–2 slides per batch, typical for WSI processing due to memory constraints.
• Data augmentation: patch-level flips, contrast/sharpness, Gaussian blur.
• Regularization: only standard dropout in Transformers, weight sharing in regional modules.
• Losses: binary cross-entropy for classification, Cox for survival outcomes.
• No pre-training of vision Transformers is needed; R²T modules are trained end-to-end.

A plausible implication is that R²T's properties (modularity, efficient locality-aware attention, and ease of integration) make it practical for scaling MIL to large images and for transfer across feature domains, with hyperparameters (region count $L$, representatives $K$) robust to selection within broad ranges.

---

For further technical details, code, and reproduction instructions, see (Cersovsky et al., 2023, Tang et al., 27 Feb 2024).