Sparse Mixture-of-Experts for Multi-Channel Imaging: Are All Channel Interactions Required?

TL;DR

The paper tackles the high computational cost of cross-channel attention in ViTs applied to multi-channel imaging. It introduces MoE-ViT, a sparse mixture-of-experts architecture that treats each channel as an expert and uses a lightweight router to activate only the most relevant channels per patch, reducing attention complexity from $O\left(N^2 C^2 D\right)$ to $O\left(N^2 C k D\right)$ where $k \ll C$. Empirical results on JUMP-CP and So2Sat show that MoE-ViT achieves substantial FLOPs savings with accuracy comparable to or better than strong channel-wise baselines, and provides a tunable efficiency-accuracy trade-off via the Top-$k$ parameter. The work demonstrates the practicality of scalable, efficient cross-channel modeling in ViTs and outlines future directions for hardware-aware optimization and broader multi-channel applications.

Abstract

Vision Transformers ($\text{ViTs}$) have become the backbone of vision foundation models, yet their optimization for multi-channel domains - such as cell painting or satellite imagery - remains underexplored. A key challenge in these domains is capturing interactions between channels, as each channel carries different information. While existing works have shown efficacy by treating each channel independently during tokenization, this approach naturally introduces a major computational bottleneck in the attention block - channel-wise comparisons leads to a quadratic growth in attention, resulting in excessive $\text{FLOPs}$ and high training cost. In this work, we shift focus from efficacy to the overlooked efficiency challenge in cross-channel attention and ask: "Is it necessary to model all channel interactions?". Inspired by the philosophy of Sparse Mixture-of-Experts ($\text{MoE}$), we propose MoE-ViT, a Mixture-of-Experts architecture for multi-channel images in $\text{ViTs}$, which treats each channel as an expert and employs a lightweight router to select only the most relevant experts per patch for attention. Proof-of-concept experiments on real-world datasets - JUMP-CP and So2Sat - demonstrate that $\text{MoE-ViT}$ achieves substantial efficiency gains without sacrificing, and in some cases enhancing, performance, making it a practical and attractive backbone for multi-channel imaging.

Figures (3)

• Figure 1: Conceptual illustration of multi-channel imaging. Each patch contains multiple channels represented by different colors.
• Figure 2: Overview of MoE-ViT. (a) While the vanilla ViT adopts a standard Multi-Head Attention module in the transformer encoder block, MoE-ViT replaces it with a specialized Multi-Head Channel MoE module tailored for multi-channel imaging. (b) Inside the Multi-Head Channel MoE module, input patches pass through the channel router, which allocates each patch to the most relevant channels for interaction (top-2 in this example). Once routing is completed, based on patch and expert indices, we loop each channel to perform patch–channel cross-attention. Specifically, patches routed to a given channel (Channel 4 in this example) serve as the source matrix (Query), while patches belonging to that channel serve as the target matrix (Key and Value) for the cross-attention operation.
• Figure 3: Ablation study comparing the Baseline (ChAda-ViT) and Ours (MoE-ViT) for varying top-$k$ values (i.e., channels) with patch sizes of 8 and 16 on both the JUMP-CP and So2Sat datasets.