SAE-Swin: Sparsity Aware Efficient Swin Transformer for Virtual Histopathological Staining
摘要
Virtual histopathological staining by convolutional neural networks (CNNs) and generative adversarial networks (GANs) has difficulty with the global context due to localized receptive fields. This limitation leads to artifacts in fine-grained tissue texture and inadequate modeling of subtle morphological details. Vision transformers (ViTs) offer an alternative to such limitations by modeling images as sequences of patches and applying self-attention to capture global context. But ViTs incur high computational cost and oversmoothing during training and inference. This inefficiency is a major obstacle for deploying ViTs on high resolution histology images. In this paper, we propose SAE-Swin, a sparsity aware efficient Swin transformer that reduces computational redundancy and mitigates over-smoothing via three core modules i.e., sparsity attention mechanism, adaptive feature recalibration, and hierarchical residual refinement. Trained on a paired E-Staining DermaRepo dataset, SAE-Swin achieves superior image fidelity (SSIM = 0.912, PSNR = 28.4 dB) and perceptual quality (FID = 15.2, LPIPS = 0.068) compared to a standard Swin Transformer (SSIM = 0.887, PSNR = 26.9 dB; FID = 23.5, LPIPS = 0.091). SAE-Swin reduces inference latency (~33%), peak GPU memory (~29%), and training parameters (~7.7%). Ablation experiments also confirm each module’s contribution to both fidelity and efficiency, while sparsity ratio experiments (50%, 30%, and 10% token retention) allow flexible efficiency–accuracy tradeoffs. These results demonstrate that SAE-Swin advances virtual staining by delivering high fidelity outputs under strict resource constraints.