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

Abstract: Medical image analysis depends on accurate segmentation and controllable synthesis, but both tasks face severe spatial imbalance: lesions occupy small regions against large backgrounds. We study adaptive spatial weighting as a task-level design principle and instantiate it in two adapters. LAW learns per-pixel loss weights for mask-conditioned diffusion by modulating a ratio prior with a feature-dependent delta map, with normalization, clamping, and Dice regularization for stability. ORDER improves lightweight segmentation by adding selective bidirectional skip attention with stage-wise confidence gating. \rev{On held-out diffusion test sets, LAW lowers FID from 158.13$\pm$0.15 to 108.43$\pm$0.71 on Polyps, from 144.13$\pm$0.31 to 89.51$\pm$0.96 on KiTS19, and from 139.22$\pm$0.38 to 112.58$\pm$0.68 on BRISC, while improving held-out mask-recovery Dice from 0.681$\pm$0.013 to 0.825$\pm$0.003 on Polyps.} \rev{When the resulting images are added to nnUNet training, downstream Polyps mDice rises from 71.7$\pm$0.4 to 74.1$\pm$0.8.} \rev{On the cleaned Polyps segmentation protocol, the reported ORDER$[0,1]$ configuration reaches 76.3$\pm$1.9 mDice and 67.2$\pm$2.0 mIoU at 42K parameters and 0.11 GFLOPs, versus 70.3$\pm$1.5 mDice and 59.9$\pm$1.7 mIoU for matched MK-UNet. On BRISC under the same training recipe, ORDER$[0,1]$ reaches 77.4$\pm$0.8 mDice and 68.1$\pm$0.7 mIoU.} These results position adaptive spatial weighting as a practical design idea for both medical diffusion and efficient segmentation.

Submission Type: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission: Described in response to reviewers. Will update after reviewer discussion is over.

Assigned Action Editor: ~Jose_Dolz1

Submission Number: 7770