Bi-level Learning of Task-Specific Decoders for
Joint Registration and One-Shot Medical Image Segmentation

CVPR 2024

Xin Fan1, Xiaolin Wang1, Jiaxin Gao1, Jia Wang1, Zhongxuan Luo1, Risheng Liu1,

1School of Software Technology, Dalian University of Technology, Dalian, China   

Paper

Code

BibTex

Abstract

One-shot medical image segmentation (MIS) aims to cope with the expensive, time-consuming, and inherent hu-man bias annotations. One prevalent method to address one-shot MIS is joint registration and segmentation (JRS)with a shared encoder, which mainly explores the voxel-wise correspondence between the labeled data and unlabeled data for better segmentation. However, this method omits underlying connections between task-specifc decoders for segmentation and registration, leading to unstable train-ing. In this paper, we propose a novel Bi-level Learning of Task-Specifc Decoders for one-shot MIS, employing a pretrained fxed shared encoder that is proved to be more quickly adapted to brand-new datasets than existing JRS without fxed shared encoder paradigm. To be more spe-cifc, we introduce a bi-level optimization training strategy considering registration as a major objective and segmenta-tion as a learnable constraint by leveraging inter-task cou-pling dependencies. Furthermore, we design an appear-ance conformity constraint strategy that learns the back-ward transformations generating the fake labeled data used to perform data augmentation instead of the labeled image,to avoid performance degradation caused by inconsistent styles between unlabeled data and labeled data in previ-ous methods. Extensive experiments on the brain MRI task across ABIDE, ADNI, and PPMI datasets demonstrate that the proposed Bi-JROS outperforms state-of-the-art one-shot MIS methods for both segmentation and registration tasks.


Highlight

  • We propose a Bi-level optimization-based framework for Joint registration and One-shot Segmentation, termed as Bi-JROS, which precisely characterizes the coupling constraints between decoders specifc to registration and segmentation tasks.
  • We design an iterative Gradient Response (GR) algo- rithm to tackle the nested bi-level optimization challenge. It leverages the gradient response of the segmentation decoder to the registration decoder during each step of the optimization process, ensuring more effective and stable training compared to simple alternating learning strategy.
  • We propose an Appearance Conformity Constraint (ACC) to avoid the texture gap between target and atlas images and increase the diversity of the data. This is inte- grated into the segmentation task to strengthen the inter- connection between registration and segmentation.


Motivation



(a) The bi-level optimization framework which establishes the coupling dependencies between registration- and segmentation-specifc decoders. (b) Illustrating the bi-level optimization process with the feedback of segmentation optimization to the registration learning process.

Network



Overall framework of the proposed Bi-JROS. (a) demonstrates the pretraining process of the shared encoder, (b) and (c) together constitute the bi-level optimization learning phase and (d) illustrates the mechanism of gradient updating.

Qualitative Results

Segmenattion

Registration


Quantitative Results

Segmentation

Registration

Reg & Seg


Citation

@inproceedings{fan2024bi,
  title={Bi-level learning of task-specific decoders for joint registration and one-shot medical image segmentation},
  author={Fan, Xin and Wang, Xiaolin and Gao, Jiaxin and Wang, Jia and Luo, Zhongxuan and Liu, Risheng},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={11726--11735},
  year={2024}
}