LAW & ORDER: Adaptive Spatial Weighting for Medical Diffusion and Segmentation

TL;DR

This paper introduces adaptive spatial weighting techniques, LAW and ORDER, to improve medical image diffusion and segmentation by focusing computational resources on lesion regions, resulting in better generative and segmentation performance.

Contribution

It proposes two novel network adapters, LAW and ORDER, that enhance diffusion and segmentation tasks through adaptive spatial weighting, addressing spatial imbalance issues.

Findings

20% FID improvement in diffusion quality

4.9% Dice coefficient increase in segmentation

ORDER is 730x smaller than nnUNet

Abstract

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.