SNR-ST-Mix: Sample-specific Neighborhood Regression Mixup for Augmented Spatial Transcriptomics Imputation with Deep Neural Network

Technical Summary: SNR-ST-Mix

Problem Statement
Spatial transcriptomics (ST) enables gene expression profiling within tissue context but is often constrained by noisy, low-resolution, and sparsely sampled data. While deep neural networks (DNNs) have emerged as powerful tools for imputing gene expression from histology images, their performance is limited by small sample sizes and a lack of biologically informed data augmentation. Existing augmentation strategies, such as Mixup and CutMix, are primarily designed for classification tasks. When applied to regression tasks like gene expression prediction, these methods often generate biologically implausible interpolations by mixing random samples. This neglects the inherent spatial and transcriptomic relationships in ST data, leading to artificial spatial discontinuities that hinder prediction performance and generalization, particularly in sample-specific training regimes where data scarcity is acute.

Methodology
The authors propose SNR-ST-Mix (Sample-specific Neighborhood Regression Mixup), a geometry- and expression-aware data augmentation framework tailored for ST regression. The method operates on the premise that neighboring tissue regions exhibit correlated morphology and expression gradients. The framework consists of two primary conditioning mechanisms:

1. Spatial Neighborhood Selection: Instead of random mixing, SNR-ST-Mix restricts the selection of mixing partners to the $k$-nearest spatial neighbors (KNN) of an anchor spot based on Euclidean distance in coordinate space. This preserves local anatomical context and ensures that mixed pairs exhibit limited target discrepancy, adhering to a Lipschitz prior where expression varies smoothly across tissue.
2. Expression Similarity Weighting: Within the selected spatial neighborhood, the probability of mixing with a specific neighbor is adaptively weighted by a Gaussian kernel based on the distance between their gene expression profiles. This ensures that augmented samples are derived from transcriptionally similar regions, avoiding interpolation between dissimilar or heterogeneous tissue compartments.

The mixing process generates convex combinations of both input histology patches ($x_{mix}$) and target gene expression vectors ($y_{mix}$) using a mixing ratio $\lambda \sim Beta(\alpha, \alpha)$.

To further enforce biological consistency, the training objective incorporates a composite loss function:

• Vicinal Risk ($L_{vic}$): Standard Mean Squared Error (MSE) on mixed pairs.
• Mixup Consistency ($L_{cons}$): A penalty ensuring the model's prediction on a mixed input matches the linear mixture of predictions on the original endpoints, controlling unnecessary curvature along the tissue manifold.
• Edge Loss ($L_{edge}$): Aligns predicted expression differences with true expression differences across spatial edges, weighted by sampling probabilities. This prevents oversmoothing at genuine biological boundaries while promoting smoothness where targets are similar.
• Correlation Loss ($L_{corr}$): A mild prior biasing optimization toward higher Pearson correlation without destabilizing the MSE fit.

The method is architecture-agnostic, integrating seamlessly with existing backbones (e.g., Vision Transformers) without requiring additional network modules or significant computational overhead.

Key Contributions

1. First Regression-Based Mixup for ST: SNR-ST-Mix is introduced as the first mixup framework specifically designed for spatial transcriptomic regression, enforcing both spatial locality and expression similarity to generate biologically meaningful augmentations.
2. Bridging Paradigms: The method bridges externally trained and sample-specific paradigms, enhancing robustness to data scarcity while maintaining strong in-distribution generalization. It mitigates overfitting and reduces dependence on large external datasets.
3. Efficiency and Integration: The approach requires no additional network complexity and is easily integrable into standard pipelines. It demonstrates high adaptability across diverse tissue types and experimental conditions.

Experimental Results
The authors evaluated SNR-ST-Mix on 10x Genomics Visium datasets covering breast cancer, bowel cancer, healthy heart, ovarian cancer, and prostate cancer.

• Performance: Across eight datasets, SNR-ST-Mix consistently achieved the lowest Mean Squared Error (MSE) and Mean Absolute Error (MAE), and the highest Pearson Correlation Coefficient (PCC) compared to baselines (no augmentation, basic geometric augmentation, vanilla Mixup, and CutMix). For instance, in the TENX13 breast cancer dataset, SNR-ST-Mix reduced MSE from 0.2074 (baseline) to 0.1886 and improved PCC to 0.4748.
• Visual Quality: Qualitative analysis of spatial expression maps (e.g., genes ISG15, MT-ND3) showed that SNR-ST-Mix produced sharper spatial boundaries and more coherent regional intensities compared to other methods, effectively recovering high-expression regions while suppressing noise.
• Ablation Studies: Experiments confirmed that both the KNN spatial constraint and label-similarity weighting are essential and complementary. Combining both constraints yielded synergistic improvements over using either alone. Furthermore, the proposed composite loss function (consistency, edge, and correlation losses) was shown to further enhance performance by providing effective regularization.
• Robustness: Hyperparameter sensitivity analysis indicated that the method is robust to variations in the number of neighbors ($K$) and the bandwidth parameter ($\sigma$), maintaining stable performance across a wide range of settings.

Significance and Claims
The paper claims that SNR-ST-Mix provides an effective and biologically principled augmentation strategy for spatial transcriptomics regression. By explicitly leveraging spatial geometry and transcriptomic similarity, the method expands the effective training manifold and improves predictive performance without increasing model complexity. The authors emphasize that their approach addresses the fundamental limitation of existing augmentation methods failing to respect the unique spatial structure of ST data. They posit that SNR-ST-Mix offers a practical solution for improving data efficiency and prediction stability, particularly in sample-specific training scenarios where data availability is constrained. The work suggests that biologically informed augmentation is critical for advancing spatial transcriptomic analysis, especially as technologies move toward higher resolutions.