[論文レビュー] LAW & ORDER: Adaptive Spatial Weighting for Medical Diffusion and Segmentation

Anugunj Naman, Ayushman Singh+2 more|arXiv (Cornell University)|Mar 5, 2026

Advanced Neural Network Applications被引用数 0

ひとこと要約

要約: 本研究は LAW と ORDER を導入し、拡散とセグメンテーションの適応的空間重み付けを提案。拡散忠実度と下流のセグメンテーションを改善し、コンパクトで効率的な設計を実現する。

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.

研究の動機と目的

• 病変が大きな背景の中で占める小さな領域における医用画像の空間的不均衡を是正する。
• 拡散ベースの合成とセグメンテーションの両方に適用可能な統一的適応重み付けフレームワークを開発する。
• 領域認識型リソース配分によって生成忠実度を高め、下流のセグメンテーションを改善する。
• 実用的な医用画像処理パイプラインに適した軽量アーキテクチャで効率を維持する。

提案手法

• LAW（Learnable Adaptive Weighter）は、中間拡散特徴とコンディショニングマスクからピクセル単位の損失調整を予測し、正規化・クランプ・Dice正則化を用いて訓練を安定化させる。
• ORDER（Optimal Region Detection with Efficient Resolution）は、軽量な MK-UNet ボトムアップバックボーンの後期デコーダ段で選択的な双方向スキップ注意を追加し、注のコストを削減するために共有類似度マトリクスを使用する。
• LAW はマスク条件付き拡散と比率ベースの事前分布をベースに、デルタマップを加えて適応的な重みを生成し、最終的な正規化ステップで重みを bounds 内に収める。
• ORDER は MK-UNet のマルチカーネルエンコーダを維持し、後期段階の双方向注意機構で残差更新をスケールする信頼度ゲートを導入する。
• 訓練は LAW ベースの重み付き拡散損失、教師あり監督、任意の蒸留損失の組み合わせを含み、セグメンテーションは selective attention と融合を用いた BCE+Dice を用いる。

Figure 7 : Qualitative segmentation results from ORDER on challenging polyp cases. Each row shows: input image, ground truth mask, and predicted segmentation. ORDER produces accurate boundaries even on difficult cases with varying polyp sizes and shapes, demonstrating the effectiveness of selective

実験結果

リサーチクエスチョン

• RQ1 拡張的空間重み付けは医用画像合成の拡散モデルにおける病変焦点の除去をどのように改善するか。
• RQ2 学習された特徴認識型の重み付けは、拡散訓練における領域強調の固定事前分布よりも優れているか。
• RQ3 軽量セグメンテーションモデルで選択的双方向注意は、空間的に不確かな領域を過剰計算なしに改善できるか。
• RQ4 合成データで訓練した場合、下流セグメンテーション性能に対する適応的重み付けの影響は何か。
• RQ5 超コンパクトなセグメンテーションバックボーンに適用した場合、実用的な効率と精度のバランスはどうなるか。

主な発見

• LAW はポリープデータセットで均一ベースラインに対して FID を 20% 改善（65.60 → 52.28）。
• LAW によるデータで訓練した下流セグメンテーションで Dice が 6% 向上（78.3 → 83.2）。
• ORDER は MK-UNet と比較して Dice が 6.0% 向上（81.3% vs 75.3%）を、パラメータ数 42k で実現。
• ORDER は 0.56 GFLOPs と 0.56G パラメータを達成し、nnUNet の約 730 倍小さくなりつつ Dice を 81.3%、IoU を 81.7% に改善。
• LAW と ORDER は Polyps および KiTS19 データセットの両方で一貫した改善を示し、生成タスクと識別タスクの両方で適応的空間重み付けの共通原理を検証した。

Figure 8: Visualization of learned attention maps from LAW during training. Top row : Conditioning masks. Bottom row : Learned attention maps (delta maps). At step 0, attention is uniform. By step 250, the model begins focusing on lesion regions. At step 500, attention strongly emphasizes lesion bou