LAW & ORDER: Adaptive Spatial Weighting for Medical Diffusion and Segmentation
This work addresses the problem of spatial imbalance in medical image analysis for researchers and practitioners working with medical diffusion and segmentation models, offering strong specific gains for these tasks.
This paper tackles spatial imbalance in medical image analysis by introducing two network adapters: LAW, which improves diffusion model generation by 20% FID, and ORDER, which enhances segmentation by 6.0% Dice coefficient while being significantly smaller than standard models. The synthetic data generated by LAW also improved downstream segmentation by 4.9% Dice coefficient.
Medical image analysis relies on accurate segmentation, and benefits from controllable synthesis (of new training images). Yet both tasks of the cyclical pipeline face spatial imbalance: lesions occupy small regions against vast backgrounds. In particular, diffusion models have been shown to drift from prescribed lesion layouts, while efficient segmenters struggle on spatially uncertain regions. Adaptive spatial weighting addresses this by learning where to allocate computational resources. This paper introduces a pair of network adapters: 1) Learnable Adaptive Weighter (LAW) which predicts per-pixel loss modulation from features and masks for diffusion training, stabilized via a mix of normalization, clamping, and regularization to prevent degenerate solutions; and 2) Optimal Region Detection with Efficient Resolution (ORDER) which applies selective bidirectional skip attention at late decoder stages for efficient segmentation. Experiments on polyp and kidney tumor datasets demonstrate that LAW achieves 20% FID generative improvement over a uniform baseline (52.28 vs. 65.60), with synthetic data then improving downstream segmentation by 4.9% Dice coefficient (83.2% vs. 78.3%). ORDER reaches 6.0% Dice improvement on MK-UNet (81.3% vs. 75.3%) with 0.56 GFLOPs and just 42K parameters, remaining 730x smaller than the standard nnUNet.