ChanSFormer: A Channel Agnostic Vision Transformer for Multi-Channel Cell Painting Images

AI summary

Overview

Research area: Computer vision / biomedical image analysis — specifically architecture design for multi-channel microscopy (Cell Painting) and other heterogeneous multi-channel imagery such as satellite data.

Technical level: Intermediate. The paper assumes familiarity with Vision Transformers, self-attention complexity, and self-supervised learning frameworks (DINOv2, iBOT), but the core architectural idea can be grasped without deep mathematics. The GPU-kernel section is Advanced.

Scope (one sentence): The paper proposes ChanSFormer, a Vision Transformer whose attention is split into separate channel and spatial steps with per-channel class tokens, so that a single model can handle any number and ordering of input channels without retraining, and evaluates it on CHAMMI, JUMP-CP and So2Sat in supervised and self-supervised regimes.

What This Paper Is About

Cell Painting assays produce images with multiple fluorescent channels (typically five or more), each revealing a different cellular component, and experimental pipelines constantly change which channels exist. Standard channel-adaptive Vision Transformers handle this by adding a learnable "channel embedding" to every token and running one global self-attention over all channels, which means the model is effectively hard-coded to the channel configuration it was trained on. The authors' goal is a backbone that is genuinely channel-agnostic — invariant to the number and order of channels — while still letting information flow between channels and remaining computationally efficient.

Key Contributions

1. Disentangled channel-spatial attention. Each encoder replaces monolithic global attention with a channel attention (across channels, at each spatial position) followed by a spatial attention (across spatial positions, within each channel), so cross-channel interaction is preserved without any learned channel embedding.
2. Individual channel CLS tokens. A single learnable CLS token is copied to all C channels at the input and allowed to diverge through the stack, so each channel carries its own class token instead of one averaged global token. This also serves as multiple independent views for self-supervised learning and as the signal for channel sampling.
3. Reduced complexity and higher throughput. The decomposition cuts cost from O((CN)²) to O(CN² + NC²), and a custom Triton channel-attention kernel that batches many short sequences and falls back to CUDA cores for small channel counts improves throughput by 55%–313% over the global-attention baseline.
4. Representation-based channel sampling. Instead of sampling channels using distances between static channel embeddings (as DiChaViT does), the model maintains an exponential moving average matrix S of pairwise cosine similarities between channel CLS tokens and uses it to bias which channels are sampled during training.

Main Findings

• Supervised state of the art on JUMP-CP and So2Sat. ChanSFormer reaches 73.46% (Full, 8 channels) and 61.91% (Flu., 5 fluorescence channels) top-1 accuracy on JUMP-CP, beating DiChaViT (69.19% / 57.98%) by 3.93%–4.27%. On So2Sat it reaches 64.76% (Full, 18 channels) and 49.75% (Sen-1, 8 channels) versus DiChaViT's 63.36% / 47.76%, an improvement of 1.40%–1.99%.
• Slightly behind on CHAMMI. ChanSFormer scores 68.83% macro F1 versus DiChaViT's 69.68%, a gap of 0.85%. The authors attribute this to omitting DiChaViT's embedding-dependent Channel Diversification Loss, and note DiChaViT falls to 68.07% without that loss — below ChanSFormer.
• Smaller backbone. The gains come with 20.1M parameters versus 21.5M for the recent ViT baselines, achieved by reducing embedding dimension from 384 to 320 and using 5 attention heads.
• Self-supervised transfer. With DINOv2 pretraining on CHAMMI, ChanSFormer scores 50.24 (CHAMMI Avg), 40.82 (JUMP-CP Full), 39.14 (JUMP-CP Flu.), 68.08 (So2Sat Full) and 46.23 (So2Sat Sen-1), beating BoC (48.92 / 39.36 / 36.42 / 67.50 / 46.46) and ChannelViT (44.29 / 23.94 / 24.71 / 66.65 / 44.54). In-domain linear probing beats BoC by 1.32% on CHAMMI and 2.91% on JUMP-CP, and cross-dataset transfer leads BoC by up to 2.72%.
• ChannelViT collapses when channels are unseen. During cross-dataset transfer, ChannelViT lags ChanSFormer by up to 16.88% because its learned representations are tied to the pretraining channel configuration and unseen channels must be zero-padded.
• Novel-channel generalization (leave-one-source-out). Pretraining on CHAMMI with HPA held out, ChanSFormer scores 46.00 on CHAMMI, 38.38 on JUMP-CP Full and 37.25 on JUMP-CP Flu., versus BoC's 35.72 / 19.63 / 19.65 and ChannelViT's 33.92 / 19.54 / 19.53. With Allen held out, ChanSFormer scores 44.38 / 38.58 / 37.58 versus BoC's 46.76 / 37.75 / 35.67. ChanSFormer wins five of the six settings; ChannelViT sits near 19.5% on JUMP-CP.
• Robustness to missing channels. On JUMP-CP with k of 8 channels available, ChanSFormer leads DiChaViT by 0.20% to 3.70% for 2–7 channels, and is comparable at k=1 (20.65% ± 1.3 versus 20.90% ± 1.6, where DiChaViT is marginally higher).
• Throughput. On 8-channel 224×224 inputs, inference throughput goes from 72.2 images/sec (ChannelViT) to 112.0 (ChanSFormer) at patch size 16, and from 7.7 to 22.9 at patch size 8. Training goes from 16.3 to 33.9, and from 1.5 to 6.2. This corresponds to 55% (patch 16) and 197% (patch 8) inference gains, and 108%–313% training gains.
• Kernel-level speedups. On tensors of shape (B, C, N+1, D) = (64, 8, 197, 384), the custom channel attention takes 4.9 ms / 634 MB at inference and 12.1 ms / 877 MB in training, versus naive attention at 7.1 ms / 1029 MB and 26.3 ms / 1117 MB, and memory-efficient attention at 6.0 ms / 797 MB and 18.6 ms / 2283 MB. That is a 35%–54% time reduction and 20%–62% memory reduction.
• Channel embeddings are not needed. Adding channel embeddings to ChanSFormer gives only marginal changes (68.83 → 69.03 on CHAMMI, 73.46 → 73.78 on JUMP-CP Full), whereas removing them from ChannelViT causes a 14.53-point drop on CHAMMI (to 50.37) — the configuration used by CA-MAE.
• Cross-channel interaction is essential. A "Spatial-only" variant that processes channels independently (BoC-style with the same attention pooling) drops to 61.73 on CHAMMI and 41.56 on JUMP-CP Full.
• Individual CLS tokens matter. Replacing representation-based sampling with uniform HCS sampling costs 2.66% on CHAMMI, 0.65% on JUMP-CP Full and 0.13% on JUMP-CP Flu.; additionally reverting to a single TimeSformer-style averaged CLS token drops scores further to 64.58 / 69.77 / 58.00.
• Channel-level supervision helps SSL. Channel-level tokens beat image-level attention pooling in all six DINOv2 probe settings (e.g., 50.24 vs 48.61 on CHAMMI).

Methodology in Plain English

The model takes an image with C channels. Each channel is cut into patches and passed through the same shared patch-embedding layer, plus a spatial positional embedding that is shared across channels — so nothing in the input stage depends on which channel it is. A single learnable class token is duplicated once per channel and prepended to that channel's patch tokens.

Each transformer block then does two attentions instead of one. First, at every spatial position, tokens from all channels look at each other (channel attention), which lets the model compare what different stains show at the same pixel location. Second, within each channel, tokens look at each other across space (spatial attention), which is the usual image-processing step. Because the channel tokens start identical and each one only mixes with its own channel's spatial tokens, they drift apart over layers and end up representing individual channels. No channel index or ordering is ever encoded, so the model is permutation-equivariant along the channel axis.

For supervised classification, the per-channel CLS tokens are pooled into one image token by a single-query attention pool. For DINOv2 self-supervised pretraining, each channel's CLS token is treated as a separate view feeding the DINO loss, while patch tokens feed the iBOT loss (with KoLeo disabled and an iBOT weight of 0.6 to stabilize small models).

Training also randomly drops channels: a subset size is drawn uniformly from 1 to C, and an exponential moving average of pairwise cosine similarities between channel CLS tokens biases which channels get selected. This replaces the fixed channel-embedding lookup that prior work used for the same purpose. Sampling is switched off at inference.

Finally, because channel attention operates over very short sequences (typically fewer than 16 channels), standard attention kernels perform badly — they are tuned for long sequences and require extra tensor permutations. The authors fuse the whole channel-attention computation into one Triton kernel, batch many short instances together, and switch between CUDA cores and TensorCores depending on channel count.

Why This Matters

Impact on research. The paper reframes channel-agnosticism as an architectural property rather than a lookup table. It shows that the two common workarounds — dropping channel embeddings entirely (CA-MAE) or processing each channel in isolation (BoC) — are both worse than decomposing attention, with the embedding-free ChannelViT dropping 14.53 points on CHAMMI and the spatial-only variant losing more than 11 points. The channel-CLS-as-views idea also gives self-supervised learning more supervision per image, which is useful for the small-model regime typical of academic microscopy work.

Real-world applications:

• Drug discovery screening with Cell Painting, where assays are run with varying dye panels and a model trained on one panel should transfer to another without retraining.
• Protein localization and cell-perturbation profiling from Human Protein Atlas or JUMP-CP-style datasets, where fluorescence and brightfield channels are mixed and often partially missing.
• Remote sensing with heterogeneous sensor stacks (Sentinel-1 SAR plus Sentinel-2 multispectral in So2Sat), where satellites and sensors change over time.
• Any pipeline where channels are semantically distinct and unordered — the authors explicitly motivate the design by contrasting it with video frames (time-ordered) and multi-view images (camera-geometry-related).

Industry relevance. Pharmaceutical and biotech companies running high-content imaging generate channel configurations that change between experiments; a backbone that does not need channel-specific retraining lowers the cost of deploying a shared foundation model across assays. The throughput gains (55%–313%) and memory savings (20%–62% on the channel-attention kernel) matter for large-scale pretraining budgets. Srinivasan Sivanandan is employed by Insitro, and Insitro has filed a patent application related to the preliminary work.

Future Directions

• Channel-agnostic diversification loss. The authors' stated limitation is that dropping fixed channel embeddings precludes DiChaViT-style channel-diversity auxiliary losses, which is why CHAMMI is 0.85% behind. They note that directly forcing data-dependent CLS tokens apart may suppress meaningful cross-channel correlations, so a new objective is needed.
• Scaling to foundation-model size. The paper states that scaling ChanSFormer to the data and model size of BoC is left to future work due to resource constraints.
• Extending the evaluation. So2Sat is excluded from self-supervised pretraining because of its small 32×32 resolution; whether the channel-agnostic benefit holds for other multi-sensor remote-sensing benchmarks is untested here.
• Generalizing the kernel. The appendix detailing the channel-attention kernel is truncated in the provided content (it begins describing channel sizes ranging from 1), so the full design space of the short-sequence attention kernel is only partially described; whether similar kernels help other short-axis attention patterns is an open question.

Target Audience

Researchers and engineers working on multi-channel or multi-modal imaging who need a backbone that survives changes in channel count and semantics: computational biology and microscopy groups, remote-sensing practitioners, and ML engineers building foundation models for scientific imaging. It is also relevant to ViT architecture researchers interested in factorized attention and to systems-minded readers interested in attention kernels for short sequence lengths. Readers without transformer background will need to consult the cited ViT, DINOv2 and iBOT papers first.

Authors’ abstract

High-content multi-channel imaging, from Cell Painting assays to remote sensing, is central to modern scientific pipelines. Because these channels are highly heterogeneous and experimental configurations frequently evolve, a practical vision backbone must be fundamentally channel-agnostic. Current channel-adaptive vision transformers rely on global self-attention with learnable channel embeddings, which can be computed only for previously known channels. Attempts to overcome this limitation of new channels either discard the embeddings, losing channel identity, or process each channel independently, thereby eliminating critical cross-channel interactions. To address this, we introduce ChanSFormer, an efficient Vision Transformer that replaces rigid embeddings with disentangled spatial-channel attention and individual channel CLS tokens. This architecture natively preserves both distinct channel identities and cross-channel information flow. Furthermore, the individual channel class token design unlocks more informative feature representations for self-supervised learning and enables representation-based channel sampling while remaining channel-agnostic. We evaluate ChanSFormer on two biology multi-channel datasets, CHAMMI, JUMP-CP and a satellite dataset, So2Sat. Experimental results show that ChanSFormer outperforms state-of-the-art methods by 3.93%-4.27% on JUMP-CP in classification accuracy. It also demonstrates exceptional cross-dataset transferability under a self-supervised setting, outperforming previous methods by up to 2.72% in accuracy in the same classification tasks. Furthermore, the disentangled attention reduces the quadratic complexity of channel-by-spatial sequence length of the global attention baseline, thus improving throughput by 55%-313%. Code is available at https://github.com/cvlab-stonybrook/ChanSFormer.