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

TL;DR

This paper introduces MoE-ViT, a sparse mixture-of-experts approach for multi-channel vision transformers that reduces computational costs by selectively modeling channel interactions, maintaining or improving performance.

Contribution

It proposes a novel MoE-based architecture for multi-channel ViTs that efficiently models only essential channel interactions, addressing computational bottlenecks.

Findings

Achieves significant efficiency improvements in multi-channel ViTs.

Maintains or improves performance on real-world datasets.

Reduces FLOPs and training costs without sacrificing accuracy.

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.