Learning Generalized Medical Image Segmentation from Decoupled Feature Queries

Authors

  • Qi Bi Jarvis Research Center, Tencent YouTu Lab School of Remote Sensing and Information Engineering, Wuhan University
  • Jingjun Yi Jarvis Research Center, Tencent YouTu Lab School of Remote Sensing and Information Engineering, Wuhan University
  • Hao Zheng Jarvis Research Center, Tencent YouTu Lab
  • Wei Ji Department of Electrical and Computer Engineering, University of Alberta
  • Yawen Huang Jarvis Research Center, Tencent YouTu Lab
  • Yuexiang Li Medical AI ReSearch (MARS) Group, Guangxi Medical University
  • Yefeng Zheng Jarvis Research Center, Tencent YouTu Lab

DOI:

https://doi.org/10.1609/aaai.v38i2.27839

Keywords:

CV: Medical and Biological Imaging, CV: Segmentation, ML: Transfer, Domain Adaptation, Multi-Task Learning

Abstract

Domain generalized medical image segmentation requires models to learn from multiple source domains and generalize well to arbitrary unseen target domain. Such a task is both technically challenging and clinically practical, due to the domain shift problem (i.e., images are collected from different hospitals and scanners). Existing methods focused on either learning shape-invariant representation or reaching consensus among the source domains. An ideal generalized representation is supposed to show similar pattern responses within the same channel for cross-domain images. However, to deal with the significant distribution discrepancy, the network tends to capture similar patterns by multiple channels, while different cross-domain patterns are also allowed to rest in the same channel. To address this issue, we propose to leverage channel-wise decoupled deep features as queries. With the aid of cross-attention mechanism, the long-range dependency between deep and shallow features can be fully mined via self-attention and then guides the learning of generalized representation. Besides, a relaxed deep whitening transformation is proposed to learn channel-wise decoupled features in a feasible way. The proposed decoupled fea- ture query (DFQ) scheme can be seamlessly integrate into the Transformer segmentation model in an end-to-end manner. Extensive experiments show its state-of-the-art performance, notably outperforming the runner-up by 1.31% and 1.98% with DSC metric on generalized fundus and prostate benchmarks, respectively. Source code is available at https://github.com/BiQiWHU/DFQ.

Published

2024-03-24

How to Cite

Bi, Q., Yi, J., Zheng, H., Ji, W., Huang, Y., Li, Y., & Zheng, Y. (2024). Learning Generalized Medical Image Segmentation from Decoupled Feature Queries. Proceedings of the AAAI Conference on Artificial Intelligence, 38(2), 810-818. https://doi.org/10.1609/aaai.v38i2.27839

Issue

Section

AAAI Technical Track on Computer Vision I