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

arXiv: https://arxiv.org/abs/2603.04795 · HTML: https://mlsanity.com/paper/2603.04795/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 43 of its 49 references that are in the map.
- First posted on arXiv 5 Mar 2026.

## Builds on (most cited first, 30 of 43)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Adding Conditional Control to Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.05543/) (2023)
- [MobileNetV2: Inverted Residuals and Linear Bottlenecks](https://mlsanity.com/paper/1801.04381/) (2018)
- [Squeeze-and-Excitation Networks](https://mlsanity.com/paper/1709.01507/) (2017)
- [MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications](https://mlsanity.com/paper/1704.04861/) (2017)
- [CBAM: Convolutional Block Attention Module](https://mlsanity.com/paper/1807.06521/) (2018)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [UNet++: A Nested U-Net Architecture for Medical Image Segmentation](https://mlsanity.com/paper/1807.10165/) (2018)
- [Attention U-Net: Learning Where to Look for the Pancreas](https://mlsanity.com/paper/1804.03999/) (2018)
- [T2I-Adapter: Learning Adapters to Dig out More Controllable Ability for Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.08453/) (2023)
- [Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation](https://mlsanity.com/paper/2105.05537/) (2021)
- [PraNet: Parallel Reverse Attention Network for Polyp Segmentation](https://mlsanity.com/paper/2006.11392/) (2020)
- [Score-based diffusion models for accelerated MRI](https://mlsanity.com/paper/2110.05243/) (2021)
- [Diffusion Models for Medical Anomaly Detection](https://mlsanity.com/paper/2203.04306/) (2022)
- [The state of the art in kidney and kidney tumor segmentation in contrast-enhanced CT imaging: Results of the KiTS19 Challenge](https://mlsanity.com/paper/1912.01054/) (2019)
- [UNeXt: MLP-based Rapid Medical Image Segmentation Network](https://mlsanity.com/paper/2203.04967/) (2022)
- [DiffuMask: Synthesizing Images with Pixel-level Annotations for Semantic Segmentation Using Diffusion Models](https://mlsanity.com/paper/2303.11681/) (2023)
- [Polyp-PVT: Polyp Segmentation with Pyramid Vision Transformers](https://mlsanity.com/paper/2108.06932/) (2021)
- [ControlNet++: Improving Conditional Controls with Efficient Consistency Feedback](https://mlsanity.com/paper/2404.07987/) (2024)
- [Dataset Diffusion: Diffusion-based Synthetic Dataset Generation for Pixel-Level Semantic Segmentation](https://mlsanity.com/paper/2309.14303/) (2023)
- [Shallow Attention Network for Polyp Segmentation](https://mlsanity.com/paper/2108.00882/) (2021)
- [Fast Unsupervised Brain Anomaly Detection and Segmentation with Diffusion Models](https://mlsanity.com/paper/2206.03461/) (2022)
- [HiFormer: Hierarchical Multi-scale Representations Using Transformers for Medical Image Segmentation](https://mlsanity.com/paper/2207.08518/) (2022)

## Abstract (arXiv)

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. On held-out diffusion test sets, LAW lowers FID from 158.13±0.15 to 108.43±0.71 on Polyps, from 144.13±0.31 to 89.51±0.96 on KiTS19, and from 139.22±0.38 to 112.58±0.68 on BRISC, while improving held-out mask-recovery Dice from 0.681±0.013 to 0.825±0.003 on Polyps. When the resulting images are added to nnUNet training, downstream Polyps mDice rises from 71.7±0.4 to 74.1±0.8. On the cleaned Polyps segmentation protocol, the reported ORDER configuration reaches 76.3±1.9 mDice and 67.2±2.0 mIoU at 42K parameters and 0.11 GFLOPs, versus 70.3±1.5 mDice and 59.9±1.7 mIoU for matched MK-UNet. On BRISC under the same training recipe, ORDER reaches 77.4±0.8 mDice and 68.1±0.7 mIoU. These results position adaptive spatial weighting as a practical design idea for both medical diffusion and efficient segmentation.