ImmuVis: Hyperconvolutional Foundation Model for Imaging Mass Cytometry

Technical Summary: ImmuVis

Problem Statement
Imaging Mass Cytometry (IMC) enables high-throughput, subcellular spatial profiling of tissues by measuring metal-tagged antibodies. However, a fundamental challenge in modeling IMC data is the lack of a fixed channel space. Unlike natural images with a static RGB channel set, IMC panels vary significantly across studies due to physical channel interference and cohort-specific experimental designs. This heterogeneity violates the core assumption of standard vision backbones, which require a fixed input dimension. Existing approaches, such as U-Nets, are restricted to fixed marker panels, while recent Transformer-based models (e.g., VirTues, EVA) attempt to handle variable inputs by concatenating marker-specific spatial tokens with identity encodings. These token-based methods often ignore crucial local marker dependencies and incur high computational costs as marker counts and spatial resolutions increase. Furthermore, virtual staining (predicting unmeasured markers) is an ill-posed problem when evidence is weak; current models often lack calibrated uncertainty estimates, making it difficult to distinguish between biologically plausible reconstructions and erroneous predictions.

Methodology
The authors introduce ImmuVis, a family of foundation models designed to unify variable-input learning, computational efficiency, and uncertainty quantification. The core innovation is a marker-adaptive hyperconvolution architecture.

1. Hyperconvolutional Architecture:

   • Marker-Agnostic Encoding: Input marker channels are first processed independently by a shared convolutional stem ($Enc_{ma}$) to capture local spatial patterns.
   • Conditional Hyperconvolution: Instead of fixed weights, the model employs a hypernetwork that generates convolutional kernels conditioned on learned marker embeddings. For an input set of markers $I_e$, a hyperkernel $H^e_{I_e}$ is constructed by concatenating learned kernel generators $\phi_e(i)$ for each marker $i \in I_e$. This allows the model to dynamically adapt its feature extraction operators to the specific subset of observed markers without architectural changes.
   • Pan-Marker Backbone: The output of the hyperconvolution is a unified, fixed-size embedding that can be processed by standard computer vision backbones (e.g., ConvNeXt v2, Swin Transformer, ViT).
   • Symmetric Decoder: The decoder mirrors the encoder, using a marker-conditioned hyperconvolution to map the latent representation to specific target markers, predicting both the mean expression and log-variance (uncertainty).

2. Training Strategy:

   • Dataset: The models are pretrained on IMC17M, the largest IMC dataset to date, comprising 24,405 images, 265 markers, and over 17 million patches across 28 cohorts.
   • Objective: The models are trained using a masked-channel reconstruction task with a Gaussian heteroscedastic regression loss. This objective models the target pixels as independent Gaussian distributions, allowing the network to predict both the mean ($\mu$) and the variance ($\sigma^2$). The variance prediction provides a calibrated uncertainty estimate that correlates with reconstruction error.
   • Preprocessing: The authors employ a specific preprocessing pipeline involving arcsinh transformation (to handle ion count dynamics), Butterworth low-pass filtering (to reduce IMC-specific noise), and min-99th percentile scaling.

Key Contributions

• Hyperconvolutional Design: A novel architecture that generates convolutional kernels from learned marker embeddings, enabling a single model to process arbitrary marker subsets without retraining or architectural modification.
• Efficient Foundation Model Family: Implementation of three variants (ImmuVisConv, ImmuVisSwin, ImmuVisViT) that leverage established backbones. These models achieve state-of-the-art performance with significantly lower computational costs compared to token-concatenation-based Transformer alternatives.
• Uncertainty-Aware Learning: Integration of a heteroscedastic likelihood objective that produces pixel-wise uncertainty estimates, addressing the reliability issues inherent in virtual staining for rare phenotypes or atypical patterns.
• Large-Scale Pretraining: The curation and utilization of IMC17M, facilitating the learning of robust, transferable representations across diverse histologies and marker panels.

Results

• Virtual Staining: ImmuVis variants outperform state-of-the-art baselines (VirTues, EVA) in virtual staining tasks. Specifically, ImmuVisSwin and ImmuVisConv achieved the lowest Mean Squared Error (MSE) and highest Pearson correlation, demonstrating superior preservation of spatial structure compared to token-based methods.
• Uncertainty Calibration: The heteroscedastic head produces uncertainty estimates that strongly correlate with reconstruction error (Pearson correlation > 0.92). The models provide meaningful per-pixel confidence signals, with ImmuVisSwin achieving the lowest Area Under Sparsification Error (AUSE).
• Representation Learning: In downstream tasks, ImmuVis embeddings yielded state-of-the-art performance in cell typing (ImmuVisConv achieved a Macro-F1 of 0.709) and clinical prediction (ImmuVisSwin achieved a Macro-F1 of 0.641), surpassing VirTues and EVA variants.
• Computational Efficiency: ImmuVis models operate at a substantially lower compute cost. ImmuVisConv requires approximately 89.9 GFLOPs per forward pass, which is an order of magnitude lower than VirTues (1049.3 GFLOPs) and nearly 40 times lower than EVA (3590 GFLOPs). Inference times are correspondingly faster (e.g., 13.4 ms vs. 28.5 ms for VirTues).

Significance
The paper positions ImmuVis as a practical framework for real-world IMC modeling that addresses the critical constraints of panel heterogeneity and computational efficiency. By enabling a single foundation model to operate on arbitrary marker subsets without retraining, ImmuVis facilitates large-scale spatial tissue profiling and retrospective upgrading of legacy cohorts. The integration of calibrated uncertainty offers a reliability-aware approach to virtual staining, which is crucial for clinical applications where false positives or low-confidence predictions must be flagged. The authors suggest that this operator-level channel adaptivity provides a general recipe for building foundation models in multiplex imaging and beyond, combining flexibility, efficiency, and reliability within a single architecture.