Kwark et al. COMB framework for label-free whole-slide virtual staining at UIUC, USA; Significant gains in fidelity and tiling consistency.

Kwark et al. COMB framework for label-free whole-slide virtual staining at UIUC, USA; Significant gains in fidelity and tiling consistency.

Seamless Whole Slide Label-Free Virtual Staining

Seamless Whole Slide Label-Free Virtual Staining Dou Hoon Kwark1, Kianoush Falahkheirkhah1, Ji-Hun Oh1, Shirui Luo2, Volodymyr Kindratenko1,2⋆, and Rohit Bhargava1⋆ 1 University of Illinois Urbana-Champaign, Urbana, IL, USA 2 National Center for Supercomputing Applications, Urbana, IL, USA dkwark2@illinois.edu Abstract. Label-free virtual staining offers a compelling, non-destruct- ive alternative to standard histopathology; however, its clinical adoption is hindered by the computational bottlenecks inherent to processing gi- gapixel Whole Slide Images (WSIs). Current deep learning approaches re- quire patch-based inference to avoid memory constraints, which disrupts global tissue continuity and introduces tiling artifacts—displaying visi- ble seams and color shifts. To address this, we introduce the Consistency Memory Bank (COMB), a novel label-free virtual staining framework that enforces spatial and channel consistency across tiles without memory bottlenecks. COMB decouples context storage from computation, utiliz- ing a dynamic retrieval mechanism to fetch feature representations from adjacent tiles. This enables a retrieval-based context integration strategy that adopts local padding to resolve spatial discontinuities and neighbor- aware channel attention to stabilize statistical drift. Further optimized with a sliding window schedule to ensure minimal memory overhead, our method demonstrates superior performance over state-of-the-art base- lines, achieving significant improvements in both perceptual fidelity and tiling consistency, while suggesting its downstream utility in tumor seg- mentation. Code is available at https://github.com/dou0000/COMB. Keywords: Label-free virtual staining · Tiling artifact mitigation 1 Introduction Pathology diagnosis relies on physical staining to make tissue structures visible under a microscope. While effective, this standard workflow is labor-intensive, inconsistent due to chemical variations, and consumes tissue samples required for downstream molecular assays. Label-free imaging [6] offers a compelling al- ternative by visualizing tissue structures using their inherent signals, eliminating the need for destructive chemical dyeing process. The advent of virtual stain- ing—using deep learning to translate raw data into histological images—offers a powerful way to optimize pathology workflows across diverse imaging modalities [22,24,7,8,9] However, clinical adoption faces a major computational hurdle: memory con- straints prevent the single holistic inference of gigapixel Whole-Slide Images ⋆Contributed equally arXiv:2609.10914v1 [eess.IV] 9 Sep 2026

2 D. H. Kwark et al. (WSIs). To that end, recent models have been developed to process the image in independent patches. This fragmentation, however, disrupts global tissue conti- nuity, introducing tiling artifacts—spatial and photometric discontinuities that artificially cleave biological structures—which confound the visual continuity re- quired for accurate pathological review. Prior solutions for these tiling artifacts, widely adopted in virtual H&E-to- IHC staining, rely on Instance Normalization (IN) [28] augmented with global [10] or adjacent statistics [12,13]. While IN works well for standard RGB-to- RGB tasks, it can potentially underperform in label-free generation (e.g., spec- tral signal-to-RGB). Since it normalizes statistics based on the local statistics (in contrast to Batch Normalization (BN), which relies on learned global statistics), IN can sometimes lead to unstable predictions, particularly in label-free source modalities where signal distributions can be sparse and heterogeneous. Conse- quently, BN [11] remains the standard for label-free virtual staining to ensure generation quality and stability [17,27]. Despite the stability of BN, mechanisms to enforce channel and spatial con- sistency across patches with respect to BN remain underdeveloped. The recent CC-WSI-Net [19] attempted to address this by building upon the VSGD-Net architecture [18]—a strong generator integrating ResNet encoders and attention mechanisms—and adding global histogram matching to guide consistency and pixel-adjacency loss. However, this approach has limitations: global constraints often suppress local details [13], and simple pixel-adjacency losses are ineffec- tive when the input domain (e.g., spectral data) lacks direct visual correlation with the output (H&E). Therefore, there is currently no single BN-based frame- work capable of delivering seamless, artifact-free generation for label-free virtual staining. In this work, we introduce the COnsistency Memory Bank (COMB), a novel framework designed to address the limitations of patch-based inference in label- free whole-slide virtual staining. COMB dynamically aligns adjacent patches to ensure visual continuity. Specifically, our contributions are as follows: (i) We present a novel seamless whole-slide generation framework tailored specifically for label-free virtual staining. (ii) We introduce a retrieval-based context integra- tion strategy that integrates local padding [1] and neighbor-aware Convolutional Block Attention Module to enforce spatial and channel consistency with minimal memory overhead. (iii) We validate that our method achieves superior perfor- mance on virtual staining quality and tiling artifacts compared to state-of-the-art baselines, with demonstrated utility in downstream tumor segmentation. 2 Method Problem Formulation. Let x ∈RCin×H×W be a multi-spectral input and y ∈R3×H×W be the corresponding H&E slide. Our goal is to learn a generator Gθ such that ˆy = Gθ(x) is locally faithful (preserving cellular morphology and staining) and globally seamless. To handle gigapixel WSIs, we partition x into a grid of smaller tiles {xi,j}i,j of size T × T (e.g. T = 256 or 512).

Seamless Whole Slide Label-Free Virtual Staining 3 Fig. 1. COMB Architecture. A sliding-window memory bank efficiently manages global context. Described for the context radius R = 1 While such partition enables gigapixel WSI processing, it introduces two crit- ical artifacts: (i) spatial discontinuities (seams) due to truncated receptive fields at tile borders, and (ii) channel-wise drift caused by the discrepancy between global training statistics and local inference windows. 2.1 Consistency Memory Bank (COMB) To resolve these, building upon VSGD-Net [18], we introduce COMB, a novel label-free virtual staining framework that decouples context storage from compu- tation, enabling efficient retrieval of consistent spatial and channel information to ensure seamless generation. Specifically, as shown in the model overview in Fig. 1, features from neighboring tiles are extracted by the encoder and decoder and stored in COMB. This memory bank is dynamically maintained using a sliding- window mechanism, retaining only the necessary features to avoid redundancy. During inference on a given patch, COMB retrieves the relevant neighboring features as contextual guidance, allowing the current generation to leverage con- sistent spatial and channel information and ensuring spatial coherence across tiles. Spatial-Wise Consistency. Standard frameworks typically rely on explicit padding (e.g., reflection) or implicit zero-padding (e.g., the default zero-padding in standard convolution) to maintain feature map dimensions. However, these in- troduce synthetic boundary signals that are disconnected from adjacent patches. To address this, we implement local padding [1] via our proposed retrieval mech- anism. As shown in Fig. 2(i), we explicitly pad the margins of size p by querying COMB for valid context—specifically, retrieving the boundary features from the full spatial neighborhood, utilizing stored history for processed tiles and pre- fetched representations for future context. This ensures the subsequent convolu- tion operates on real, continuous tissue signals rather than synthetic artifacts. Channel-Wise Consistency (Neighbor-Aware Attention). While spatial padding ensures morphological continuity, it does not correct the statistical drift

4 D. H. Kwark et al. Fig. 2. Retrieved features enforce (i) spatial consistency via local padding and (ii) channel consistency via neighbor-aware CBAM. caused by the discrepancy between the learned global training statistics and the current local inference windows. To mitigate this, we propose a neighbor-aware Convolutional Block Attention Module (CBAM) [30] integrated within the skip connections. Standard CBAM enhances virtual staining by preserving local de- tails during encoder-to-decoder transfer [18]. However, it calculates attention weights based solely on the isolated patch, making it blind to regional inconsis- tencies. Consequently, it often amplifies local anomalies (e.g., stain variations) rather than correcting them. To address this, we reformulate the module to per- form dynamic feature recalibration (Fig. 2(ii)). By aligning the attention mech- anism with the retrieved neighbor context, we enforce regional consistency while retaining authentic tissue signals. First, we compute the standard local channel attention map ˜gij using the isolated tile features Fij via Global Average Pooling (GAP) and Global Max Pooling (GMP): ˜gij = σ (MLP (GAP(Fij)) + MLP (GMP(Fij))) , (1) where σ denotes the sigmoid function. To harmonize the local representation with its surroundings, we compute a regional attention map ¯gij from the neighbor set Nij (tiles within the context radius R of (i, j)) and use it to stabilize the final attention weights ˆgij: ¯gij = 1 |Nij| X (p,q)∈Nij ˜gpq, ˆgij = (1 −α)˜gij + α¯gij, (2) where α is a hyperparameter. This fusion ensures that if the local signal ˜gij is weak or anomalous, the model leverages the stable neighborhood map ¯gij to produce consistent inference. Concurrent Prefetching and Inference Scheduling. A naive implemen- tation of COMB requires a two-pass strategy: first caching the global context, then performing inference. However, this effectively doubles the computational cost. To overcome this, we adopt a concurrent prefetching schedule [13], enabling COMB to populate contexts in a single continuous scan. Processing in column- major order, we construct a compound mini-batch at each step containing the

Seamless Whole Slide Label-Free Virtual Staining 5 VSGD-Net Input Ours CC-WSI-Net Ground-Truth 0 1 250 𝝁m 150 𝝁m Fig. 3. Visual comparison on SRS (top) and IR (middle) datasets. The third row represents segmentation prediction on the second row. current inference tile at (i, j) and a look-ahead prefetch tile at (i+R+1, j+R), as illustrated in Fig. 1. In this setup, the prefetch tile computes and caches features into the bank, while the inference tile retrieves these pre-populated contexts to generate the final output ˆyij. This ensures the necessary neighbor contexts are available in memory just-in-time for the inference step. However, this concurrent scheduling necessitates initializing the entire global context map to accommo- date the sparse updates, which is significantly expensive for gigapixel WSIs. To address this, we further optimize the memory footprint by implementing COMB as a sliding window buffer rather than a global registry. We maintain only the active columns required by the context radius R, evicting stale columns as the inference front progresses. This reduces the memory complexity from O( H T · W T ) to O(R · H T ) in the number of cached tiles (and thus in cached feature memory), making high-resolution WSI processing feasible on standard consumer hardware. 3 Experiments Dataset. We validated our method on a cohort of frozen prostate surgical re- section samples (collected and IRB approved at Mayo Clinic), comprising two label-free modalities: multispectral mid-infrared (IR) and Stimulated Raman Scattering (SRS). The IR data was acquired using Laser Scanning Confocal Mi- croscopy (LSCM) [4] across ten discrete wave