KANResDiff

Official pytorch implementation for our MICCAI 2026 paper "KANResDiff: Learning Local Residual Diffusion via Kolmogorov-Arnold Network for Ambiguous Medical Image Segmentation"

Introduction

Ambiguous medical image segmentation aims to provide a series of diverse but plausible segmentation hypotheses. However, existing methods introduce stochasticity in a fixed and pre-defined manner, failing to form a progressive semantic modeling process. To address these challenges, we propose KANResDiff to learn local residual diffusion with Kolmogorov-Arnold Network, thereby assigning distinct roles across stages for ambiguity modeling. Specifically, we propose Independent Time Encoding that offers spline-based time embeddings instead of linear ones from MLPs, which enhances the independence across inference stages and assigns progressive semantic roles to different stages. We propose Residual Schr"odinger Bridge that injects deterministic residual prior with learnable weights by constructing local Schr"odinger Bridge instead of following manually settings, achieving a flexible deterministic–stochastic interaction and stage-aware ambiguity modeling thanks to local optimal diffusion path. Extensive experimental results on two public datasets demonstrate that KANResDiff achieves SOTA performance on GED and HM-IoU, with maximum improvements of 16.8% and 7.7%, respectively, while maintaining competitive performance on the MDM metric.

Overview

Datasets

Two public datasets: LIDC and ISIC Subset are implemented in this work. You can download the datasets from the following links:

• LIDC Dataset as preprocessed by @Stefan Knegt
• ISIC Subset as preprocessed by @killanzepf

Requirements

You can build the dependencies by executing the following command.

conda create -n kanresdiff python=3.10
conda activate kanresdiff

pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1+cu116 \
  --extra-index-url https://download.pytorch.org/whl/cu116
pip install -r requirement.txt

Prepare template/template_global_config.yml or your global config so the LIDC data path points to preprocessed .pickle files.

Training Procudure

1. Pretrain Probabilistic U-Net

This model is used to precompute Target-Rough (TR) priors.

python pretrain_PUNet.py \
  --batch_size 8 \
  --num_epochs 1600 \
  --output_path /path/to/outputs

2. Precompute TR Priors

Run once for each split.

python precomp_tr.py \
  --resume_filepath_pre /path/to/PreT_ProbUNet/best.pth \
  --mode train \
  --num_samples 16 \
  --output_path /path/to/TR

python precomp_tr.py \
  --resume_filepath_pre /path/to/PreT_ProbUNet/best.pth \
  --mode val \
  --num_samples 16 \
  --output_path /path/to/TR

This produces files such as TR_train_d0.8_n16.pt and TR_val_d0.8_n16.pt.

3. Train Backbone

python train_stage0_backbone.py \
  --batch_size 8 \
  --num_epochs 600 \
  --output_path /path/to/outputs

4. Train RSB

Recommended combined entry point:

python train_rsb.py \
  --resume_filepath /path/to/backbone/Best_xxx.pth \
  --tr_filepath /path/to/TR/TR_train_d0.8_n16.pt \
  --tr_filepath_val /path/to/TR/TR_val_d0.8_n16.pt \
  --rsb_enabled \
  --num_epochs 600 \
  --output_path /path/to/outputs

Two-step entry point:

python train_stage1_path.py \
  --resume_filepath /path/to/backbone/Best_xxx.pth \
  --num_epochs 600 \
  --output_path /path/to/stage1

python train_stage2_trnet.py \
  --resume_filepath /path/to/backbone/Best_xxx.pth \
  --stage1_filepath /path/to/stage1/K1PUPD_SB_stage1/YYYY-MM-DD_HH-MM-SS/paths_stage1.pth \
  --tr_filepath /path/to/TR/TR_train_d0.8_n16.pt \
  --tr_filepath_val /path/to/TR/TR_val_d0.8_n16.pt \
  --num_epochs 600 \
  --output_path /path/to/stage2

Stage 2 freezes the diffusion backbone and trains TRNet with the learned path models.

Testing Procedure

python test_rsb.py \
  --resume_filepath /path/to/checkpoint/Best_xxx.pth \
  --tr_filepath_test /path/to/TR/TR_test_d0.8_n16.pt \
  --sampling_mode rsb \
  --num_samples 32 \
  --sampling_times 40 \
  --output_path /path/to/outputs

Key Arguments

Argument	Meaning
--resume_filepath	Checkpoint path. For RSB inference it must contain both model_state_dict and paths.
--tr_filepath, --tr_filepath_val, --tr_filepath_test	Precomputed TR cache files.
--rsb_enabled	Enable RSB path learning in train_rsb.py.
--lambda_bridge	Forward/backward path consistency weight.
--lambda_ddpm_anchor	Weak DDPM anchor weight for forward paths.
--lambda_alpha	TRNet alpha regularization weight.
--sampling_mode	standard or rsb.
--sampling_times	Number of reverse sampling steps.

Main Files

train_stage0_backbone.py   # DDPM/KANResDiff backbone training
train_stage1_path.py       # Path initialization
train_stage2_trnet.py      # TRNet + path training with frozen backbone
train_rsb.py               # Combined RSB training
test_rsb.py                # Inference and comparison
validate.py                # Shared validation utilities

models/kan_resdiff_sb.py   # Main U-Net + TRNet model
models/path_sb/            # Path models and bridge loss
models/pretrain/           # Probabilistic U-Net and TR cache tools

Acknowledgements

• This work is proposed based on our previous CVPR 2025 open-source work. We thank @Qiu-xy for releasing the implementation.
• We thank @Stefan Knegt for the preprocessed dataset.
• We thank @killanzepf for the preprocessed dataset and the GTR baseline.

Citations

If you find this work useful, please cite KANResDiff and our previous LDSB:

@inproceedings{li2026kanresdiff,
  title={KANResDiff: Learning Local Residual Diffusion via Kolmogorov-Arnold Network for Ambiguous Medical Image Segmentation},
  author={Li, Fanding and Wang, Chenglin and Li, Xiangyu and Qiu, Xingyu and Ma, Xinghua and Yin, Xiangming and Li, Haiyang and Dong, Suyu and Wang, Wei and Wang, Kuanquan and Luo, Gongning and Li, Shuo},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  year={2026},
  organization={Springer}
}
@inproceedings{qiu2025finding,
  title={Finding local diffusion schrodinger bridge using kolmogorov-arnold network},
  author={Qiu, Xingyu and Yang, Mengying and Ma, Xinghua and Li, Fanding and Liang, Dong and Luo, Gongning and Wang, Wei and Wang, Kuanquan and Li, Shuo},
  booktitle={Proceedings of the Computer Vision and Pattern Recognition Conference},
  pages={23227--23236},
  year={2025}
}