PathRWKV: Enhancing Whole Slide Image Inference with Asymmetric Recurrent Modeling

Fig. 2a illustrates the overall PathRWKV pipeline. Following existing works [18], each WSI $S$ is first loaded at a target resolution (e.g., 0.5 microns per pixel (mpp)) and partitioned into a non-overlapping grid of tiles $\{T_{\text{i,j}}\}$ of size $T_{\text{size}}$. To ensure data quality, a two-stage filtering protocol is employed. First, tiles with tissue coverage below a predefined threshold (e.g., $<50\%$ of the tile area) are discarded. Second, tiles with pixel variance falling below a quantitative cutoff (e.g., $Var(I)<0.01$, where $I$ denotes the normalized pixel intensity in [0,1]) are removed. This ensures that only informative tiles are retained for downstream tasks. After pre-processing and filtering, each tile is embedded into a dense semantic feature vector (Fig. 2a) using a pathological foundation model (e.g., Prov-GigaPath [33]). This embedding process enhances both the training efficiency and performance of the slide-level MIL. During training, PathRWKV processes a randomly shuffled subset of tiles to facilitate efficient multi-slide learning; conversely, during inference, it utilizes the complete sequential tile sequence from the WSI. Finally, the slide-level output features from PathRWKV are passed to the MTL module to generate predictions for each task (e.g., cancer subtyping, tumor grading, overall survival).

PathRWKV serves as the backbone for slide-level feature modeling, consisting of 2 blocks with a hidden dimension of 768, and 12 heads. The input tile features are first combined with 2D sinusoidal Position Encoding (2D PE) to restore spatial relationships disrupted during the sampling process. Subsequently, these features are processed by the PathRWKV blocks (Fig. 2b). Each block comprises a TimeMix module and a ChannelMix module, integrated with layer normalization and residual connections. The TimeMix module dynamically captures multi-scale inter-tile dependencies via the temporal dimension, while the ChannelMix module focuses on intra-tile feature interactions across the channel dimension. Output features from the final PathRWKV block are aggregated via max pooling to produce the final slide-level representation.

Fig. 2c illustrates the key mathematical operations within the TimeMix and ChannelMix modules. The TimeMix module is specifically designed to capture multi-scale temporal dependencies among tiles. It effectively integrates fine-grained local interactions via token shifting and interpolation with coarse-grained global context via time-decayed linear attention, combining the strengths of Transformers [7] and RNNs [11].

To capture short-range, fine-grained dependencies between adjacent tiles, the module first employs a token-shift mechanism, TimeShift. A shifted version of the input $x$, denoted as $x_{last}$, is generated using zero-padding:

The current input $x$ is then mixed with $x_{last}$ via data-dependent linear interpolation (ddlerp). Unlike static interpolation, ddlerp dynamically computes the mixing coefficient $\mu$ using a Low-Rank Adaptation (LoRA) mechanism:

This mechanism allows the model to adaptively aggregate local information from the immediate predecessor based on the current context, ensuring that high-frequency local variations are preserved before global processing.

Following local aggregation, the interpolated features are projected into five vectors: receptance $r$, key $k$, value $v$, time-decay $w$, and gate $g$. To capture long-range, coarse-grained dependencies across the entire slide sequence, we employ a time-decayed linear attention mechanism. Notably, the decay rate $w$ is modulated by a LoRA projection, allowing for data-dependent decay speeds that can adaptively focus on relevant historical context. The global linear attention is computed efficiently as:

Crucially, this formulation can be switched to a recurrent structure, which underpins our asymmetric design. By maintaining a recurrent state $S_{t}$, the model propagates global context sequentially:

Finally, the output is gated by $g$ and projected by $W_{o}$:

The ChannelMix module focuses on intra-tile feature interactions. Specifically, it employs learnable linear projections ($r,k,v$) to blend information across the channel dimension $D$ for each tile independently. This is coupled with a Squared ReLU activation:

which induce robust non-linear transformations, enabling the extraction of complex morphological features within each tile.

As illustrated in Fig. 1, distinct from the standard token-by-token processing in original RWKV, we propose a novel hybrid set-by-set recurrent architecture tailored for high-resolution WSI analysis. While leveraging the mathematical efficiency of linear attention, our core innovation lies in the asymmetric formulation of slide-level modeling to resolve the memory bottlenecks inherent in existing SSMs.

During the training phase, we adopt a set-based parallel strategy to maximize throughput. Instead of processing tiles sequentially, the model ingests a fixed number of sampled tiles as a dense batch. By utilizing a parallel CUDA kernel, we compute cumulative states and gradients simultaneously. This design fully exploits the massive parallelism of modern GPUs, facilitating rapid backpropagation and stable convergence compared to pure recurrent training.

During the inference phase, we introduce a streaming recurrent mechanism to achieve constant $\mathcal{O}(1)$ spatial complexity. A critical limitation in previous SSM-MIL methods (e.g., S4MIL, MambaMIL) is their reliance on Attention-based aggregation, which necessitates storing all tile features ($\mathcal{O}(N)$) for the final calculation, causing memory overflows on gigapixel slides. To overcome this, our architecture decomposes the WSI into sequential chunks. A latent state S (of size HeadSize$\times$HeadSize) acts as a memory carrier, propagating context across chunks via a recurrent kernel.

Furthermore, a critical bottleneck remains in the aggregation mechanism of existing frameworks, including recent SSMs. While SSM backbones theoretically allow constant spatial complexity ($\mathcal{O}(1)$) during inference, they typically adapt the Gated Attention mechanism from ABMIL [13] for slide-level aggregation. This approach necessitates computing a global Softmax normalization term across all tiles:

Consequently, the feature vectors of all $N$ tiles must be retained in GPU memory to calculate the denominator until the entire slide is processed, forcing the inference memory complexity to scale linearly with the slide size ($\mathcal{O}(N)$). This disrupts the memory efficiency gained by the SSM backbone.

To achieve a fully memory-efficient pipeline, we integrate the backbone’s sequential processing with a streaming aggregation strategy. We decompose the WSI $\mathcal{X}$ into bags $\mathcal{B}_{1},\mathcal{B}_{2},\ldots,\mathcal{B}_{M}$ aligned with the inference chunks, and identify the feature-wise max operation as the optimal choice due to its recursive update property. Given the tile encoder $\phi_{\theta}$, we define the local summary $z_{i}$ and the combination rule as:

We then update the slide-level representation sequentially:

where $h_{i}$ represents the summary of the first $i$ bags. Unlike attention mechanisms, this design decouples memory usage from sequence length, achieving true $\mathcal{O}(1)$ space complexity during inference. This enables PathRWKV to process arbitrarily large slides on edge devices while preserving consistency with the ideal global computation. Comprehensive theoretical analysis, including proofs for the unbiased nature of the gradients and memory complexity comparisons, are provided in the supplementary materials.

Given the high dimensionality of WSIs and the sparse supervision signal, models are prone to overfitting. To mitigate this, we employ the random sampling strategy during training. Instead of processing the entire slide or a fixed region, we randomly sample a subset of tiles from the WSI in each iteration. This approach acts as a strong data augmentation technique, preventing the model from memorizing specific tile sequences and enhancing its generalization capabilities.

However, a critical side effect of random sampling is the disruption of the intrinsic 2D spatial structure and the anatomical adjacency of the tissue microenvironment. This is the main advantage of recent permutation-variant methods (e.g., MmabaMIL [34]) compared to early permutation-invariant methods (e.g., ABMIL [13]). To compensate for this loss of spatial context and enable PathRWKV to model geometry-aware dependencies, we introduce a 2D PE. Formally, let $\mathbf{z}_{i}\in\mathbb{R}^{D}$ denote the feature embedding of the $i$-th tile, and $(x_{i},y_{i})$ represent its normalized spatial coordinates. We employ sinusoidal functions to map these coordinates into a continuous embedding space:

where $\Omega$ is a scaling factor and $p\in\{x_{i},y_{i}\}$. The final spatial embedding $\mathbf{P}_{i}$ is constructed by concatenating the encodings of horizontal and vertical coordinates and injected into the tile features via addition: $\hat{\mathbf{z}}_{i}=\mathbf{z}_{i}+\mathbf{P}_{i}$. This ensures that geometric relationships are restored regardless of the sampling order.

Finally, to maximize the utility of limited annotated data and enhance training efficiency, we incorporate a multi-task learning (MTL) module. It comprises multiple prediction heads, allowing the model to learn from diverse clinical objectives simultaneously. By leveraging task-wise correlations, the model extracts more discriminative features, further reducing the risk of overfitting on the feature distribution of a single task. We employ Cross-Entropy, Cox proportional hazards, and L1 losses for classification, survival analysis, and regression tasks, respectively. The total loss is aggregated as $\mathcal{L}_{total}=\sum_{i=1}^{T}\mathcal{L}_{i}$, where $T$ represents the number of task heads. Crucially, gradients are computed only for tasks with available labels, enabling flexible training on partially annotated datasets.

We evaluated PathRWKV on 11 datasets across 9 downstream tasks covering diverse diagnostic scenarios, as shown in Fig. 3, to demonstrate its performance and generalizability. The PANDA [2] dataset with the ISUP Grade task assesses prostate cancer aggressiveness; CAMELYON16 [1] with the Breast Metastasis task classifies lymph nodes as normal or tumorous; IMP-CRS-2024 [22] with the CRC-Tumor task identifies tumor tissues in colorectal images. The TCGA [4] datasets cover multiple cancer types and tasks: TCGA-BRCA with the IHC-HER2 task predicts HER2 receptor status from H&E-stained slides; TCGA-GBM with the Histological Diagnosis task classifies glioblastoma subtypes based on morphology; TCGA-LGG with the Tumor Stage task predicts the WHO grade of lower-grade gliomas; TCGA-CESC with the Lymphovascular Involvement task detects the presence of lymphovascular invasion; TCGA-ESCA with the Cancer Grade task assesses the histological differentiation grade of esophageal carcinoma; and TCGA-LIHC, TCGA-BLCA, and TCGA-LUNG, the combination of TCGA-LUAD and TCGA-LUSC, with the Overall Survival task predict patient survival time in months from liver, bladder, and lung tissue morphology, respectively.

We use the UnPuzzle framework [18] for preprocessing, where each WSI is tessellated into $224\times 224$ patches at 0.5 mpp and embedded using the tile-level encoder of Prov-GigaPath [33]. All models are initialized from scratch and trained for 100 epochs using the AdamW optimizer and a cosine decay scheduler with a final learning rate factor of 0.1. We conduct a grid search over ten learning rates ($1\times 10^{-6}$ to $1\times 10^{-3}$) and employ early stopping with a patience of 10 epochs based on validation loss. During training, we use a batch size of 4 and randomly sample a maximum of 2,000 tiles per WSI. For evaluation, we select the checkpoint with the lowest validation loss, use a batch size of 1, and process all tiles per WSI. We report the average performance of the top-3 results for each metric. For statistical reporting, significance was evaluated using Welch’s t-test [32] across these top-3 performing learning rate configurations to compare each method against PathRWKV. All experiments were conducted on 4 NVIDIA RTX4090 GPUs using Python 3.12.12, PyTorch 2.9.1, and CUDA 12.8.

To demonstrate the effectiveness of PathRWKV for slide-level WSI modeling, we compared it against 11 state-of-the-art (SOTA) methods, including SlideAve and SlideMax from MINNs [30], ABMIL [13], CLAM [19], DSMIL [16], DTFD-MIL [36], TransMIL [28], Prov-GigaPath [33], S4MIL [9], MambaMIL [34], and MamMIL [8].

As presented in Tab. 1, PathRWKV demonstrates superior performance and robust generalizability, achieving SOTA results on 10 out of 11 datasets across 9 distinct downstream tasks. Specifically, on standard classification benchmarks such as CAMELYON16 and IMP-CRS-2024, most deep learning-based methods achieve high metrics ($>90\%$ Accuracy, AUC, and F1), with the exception of simple pooling strategies, SlideAve and SlideMax, that lack the representational capacity for comprehensive slide-level modeling. Among high-performing models, PathRWKV consistently secures the highest scores across all three metrics.

In contrast, the TCGA datasets present a significantly more challenging scenario, where the average performance of most methods drops to approximately 70% due to the intrinsic complexity and heterogeneity of the samples. Despite these challenges, PathRWKV establishes its efficacy on the majority of TCGA datasets, validating the capability of the proposed TimeMix and ChannelMix modules to capture complex pathological dependencies. However, we observe a performance gap on the TCGA-LUNG overall survival task, where MambaMIL achieves the leading performance. We attribute this to the inherent trade-off of our max pooling strategy. While it ensures $\mathcal{O}(1)$ inference memory, it focuses on the most salient features and may inadvertently discard global contextual cues (e.g., total tumor burden) that are beneficial for specific prognostic predictions, which Attention mechanism in MambaMIL preserve better. Nevertheless, the strong performance of both PathRWKV and MambaMIL underscores the structural advantage of SSMs over conventional pooling models and Transformers in modeling multi-scale relationships within WSIs.

A core innovation of PathRWKV is its asymmetric design, which fundamentally optimizes GPU memory utilization during inference. Fig. 5 compares the memory consumption profiles of various methods as the number of input tiles increases. Conventional attention-based methods (e.g., ABMIL, CLAM), Transformers (e.g., TransMIL, Prov-GigaPath), and even recent linear-complexity models (e.g., MambaMIL) exhibit linear memory growth ($\mathcal{O}(N)$). Although these SSM-based approaches utilize linear attention mechanisms similar to PathRWKV, they typically rely on the Gated Attention mechanism from ABMIL [13] for final aggregation. This strategy requires matrix multiplication between the complete input and output tensors to compute attention scores, necessitating the retention of the entire input tensor in GPU memory until the output is generated. This dependency causes memory usage to scale linearly with sequence length, effectively negating the inherent efficiency advantages of the SSM backbone and severely constraining applicability to large-scale WSIs in resource-limited environments. In contrast, by equal contribution of linear attention and max aggregation strategy, PathRWKV achieves constant memory consumption ($\mathcal{O}(1)$), as evidenced by the flat trajectory in Fig. 5. This efficiency confirms that our recurrent inference formulation enables sequential iteration over WSI tiles without caching historical states, successfully resolving the trade-off between training efficiency and inference scalability. Notably, while PathRWKV exhibits linear time complexity ($\mathcal{O}(N)$) same with other methods, it maintains a competitive inference speed, effectively balancing its relatively large parameter size for complex feature modeling with computational efficiency.

To further assess the interpretability of our framework, Fig. 4 visualizes the Class Activation Maps (CAMs) for representative samples from the CAMELYON16 dataset. Note that, as standard attention maps are not directly obtainable from Prov-GigaPath and S4MIL, we visualize their saliency maps instead. Consistent with quantitative findings, SlideAve and SlideMax fail to generate meaningful activation patterns due to their simplistic aggregation logic. Compared with the ground truth, attention-based methods (e.g., ABMIL, CLAM) tend to assign uniform weights across the entire region of interest, demonstrating a limited capacity to distinguish fine-grained intratumoral heterogeneity. Furthermore, they exhibit relatively higher attention scores in normal regions. Among transformers, TransMIL misdirects attention to incorrect areas, while Prov-GigaPath focuses on normal regions in the second sample; this is likely attributable to overfitting on the small-scale dataset. In contrast, all SSMs successfully detect the accurate regions. Notably, PathRWKV not only accurately highlights global tumor regions but also delineates local feature intensity variances through its heatmap distribution. This visualization validates the effectiveness of the multi-scale modeling facilitated by the Time Mix and Channel Mix modules, demonstrating PathRWKV’s ability to extract hierarchically significant pathological features.

A core innovation of PathRWKV is the asymmetric structure, designed to resolve the conflict between training throughput and inference memory efficiency. This design is predicated on two critical hypotheses: first, that a model trained on short sequences can effectively generalize to full-length WSIs during inference; and second, that max pooling serves as a sufficient and efficient aggregator for slide-level features. To validate these premises, we conducted ablation studies on each component.

To address the overfitting risks on small-scale datasets and the loss of spatial structure due to sampling, we introduced random sampling paired with 2D PE.

Finally, we analyzed the components designed to enhance model generalization under data-constrained conditions.

To efficiently implement the mathematical duality of the Linear Attention mechanism, we implemented two distinct sets of CUDA kernels, corresponding to the parallel (training) and recurrent (inference) modes.

The code implementation orchestrates the interaction between the CUDA kernel and the WSI data to minimize memory footprint. Let $N$ be the total number of tiles and $B_{chunk}$ be the chunk size. The inference process proceeds as follows: