[论文解读] Channel Vision Transformers: An Image Is Worth 1 x 16 x 16 Words

Yujia Bao, Srinivasan Sivanandan+1 more|arXiv (Cornell University)|Sep 28, 2023

Cell Image Analysis Techniques被引用 6

一句话总结

ChannelViT 提出了一种视觉Transformer的变体，将每个输入通道视为独立的序列维度，并使用可学习的通道嵌入来增强跨通道推理。它在ImageNet、JUMP-CP和So2Sat上均优于标准ViT，尤其在通道稀疏条件下表现更优，且分层通道采样（HCS）显著提升了仅部分通道在测试时可用情况下的鲁棒性。

ABSTRACT

Vision Transformer (ViT) has emerged as a powerful architecture in the realm of modern computer vision. However, its application in certain imaging fields, such as microscopy and satellite imaging, presents unique challenges. In these domains, images often contain multiple channels, each carrying semantically distinct and independent information. Furthermore, the model must demonstrate robustness to sparsity in input channels, as they may not be densely available during training or testing. In this paper, we propose a modification to the ViT architecture that enhances reasoning across the input channels and introduce Hierarchical Channel Sampling (HCS) as an additional regularization technique to ensure robustness when only partial channels are presented during test time. Our proposed model, ChannelViT, constructs patch tokens independently from each input channel and utilizes a learnable channel embedding that is added to the patch tokens, similar to positional embeddings. We evaluate the performance of ChannelViT on ImageNet, JUMP-CP (microscopy cell imaging), and So2Sat (satellite imaging). Our results show that ChannelViT outperforms ViT on classification tasks and generalizes well, even when a subset of input channels is used during testing. Across our experiments, HCS proves to be a powerful regularizer, independent of the architecture employed, suggesting itself as a straightforward technique for robust ViT training. Lastly, we find that ChannelViT generalizes effectively even when there is limited access to all channels during training, highlighting its potential for multi-channel imaging under real-world conditions with sparse sensors. Our code is available at https://github.com/insitro/ChannelViT.

研究动机与目标

• 为解决显微镜和卫星成像等领域的多通道成像挑战，其中通道携带语义上不同的信息，且可能稀疏可用。
• 通过实现空间位置与输入通道之间的联合推理，提升视觉Transformer（ViT）的性能。
• 在推理时仅部分输入通道可用这一常见现实场景下，提升模型的鲁棒性。
• 引入一种正则化技术，确保模型在未见通道组合上的泛化能力。
• 证明通道特定的表征可提升注意力图中的可解释性和特征定位能力。

提出的方法

• ChannelViT从每个输入通道独立构建图像块标记，将每个通道视为图像块嵌入过程中独立的序列。
• 它引入了一种可学习的通道嵌入，添加到每个图像块标记上，类似于位置嵌入，以保留通道特定的信息。
• 该模型在所有通道上共享线性投影和位置嵌入，保持架构效率的同时支持跨通道注意力。
• 分层通道采样（HCS）被引入作为正则化技术：首先采样通道数量，然后均匀采样特定的通道配置，从而提升对通道稀疏性的鲁棒性。
• HCS与通道丢弃不同，其为所有通道数量分配均匀权重，减少了对特定通道组合的偏差。
• 该架构在ImageNet、JUMP-CP（显微镜）和So2Sat（卫星成像）上进行了评估，并通过通道子集泛化消融研究进行了分析。

Figure 1: Illustration of Channel Vision Transformer (ChannelViT). The input for ChannelViT is a cell image from JUMP-CP, which comprises five fluorescence channels (colored differently) and three brightfield channels (colored in B&W). ChannelViT generates patch tokens for each individual channel, u

实验结果

研究问题

• RQ1 视觉Transformer能否在成像任务中有效适应对多个语义上不同的输入通道进行推理？
• RQ2 与标准ViT相比，所提出的ChannelViT架构在多通道成像场景中如何提升性能和鲁棒性？
• RQ3 当测试时仅可用部分输入通道时，分层通道采样（HCS）在多大程度上增强了泛化能力？
• RQ4 可学习的通道嵌入能否提升注意力图中的可解释性和特征定位能力？
• RQ5 当模型在不完整通道集上进行训练并在未见通道组合上进行测试时，其性能是否仍保持强劲？

主要发现

• 在ImageNet上，ChannelViT在使用HCS时达到66.22%的top-1准确率，优于ViT的52.06%。
• 在JUMP-CP显微镜基准上，当测试4个通道时，ChannelViT-S/16相比ViT-S/16实现了14.28%的相对提升；当在全部8个通道上训练并在7个通道上测试时，提升达35.13%。
• 当仅使用红色通道时，HCS将ImageNet验证准确率从29.39%提升至68.86%，表明其对通道稀疏性具有极强的鲁棒性。
• HCS在ViT和ChannelViT上均持续提升泛化性能，FAN-S/16（线性图像块）在使用HCS训练后，在部分通道测试中优于其非HCS对应模型。
• ChannelViT中的注意力可视化揭示了有意义的、通道特定的特征相关性，增强了模型在空间和光谱维度上的可解释性。
• ChannelViT中学习到的通道嵌入对应于可解释的语义角色，表明模型能有效区分并利用不同的通道信息。

Figure 2: Correlation patterns among image channels (left) and the learned channel embeddings (right) for ImageNet, JUMPCP, and So2Sat. ImageNet displays a strong correlation among the three RGB input channels while JUMPCP and So2Sat show minimal correlation between different signal sources (Fluores