PaperlayerView exact revision on arXiv

Source-linked AI summary

Fed-ADApt: Federated Anytime Depth Adaptation for Resource-Aware Medical Image Segmentation

Abhijeet Parida, Zhifan Jiang, Pooneh Roshanitabrizi, Austin Tapp, Maria J. Ledesma-Carbayo, Syed Muhammad Anwar, Ziyue Xu, Marius George Linguraru, Holger R. Roth

arXiv:2610.03474v1cs.CV

TL;DR

Federated segmentation must include sites whose hardware cannot train or deploy full-capacity models, while UNet depth adaptation must preserve skip-connected encoder–decoder paths. Fed-ADApt combines budget-aware training with multi-depth supervision and hierarchical aggregation; across 2D retinal and 3D brain tumor tasks, it matches FedAvg Dice performance in 3D and achieves competitive 2D performance with lower resource costs.

  • Problem

    Existing federated approaches assume similar compute resources across sites, while depth adaptation for UNet segmentation must preserve encoder–decoder compatibility, skip connections, and full-resolution prediction.

    Source: §I · p. 2

  • Method

    Fed-ADApt combines multi-depth supervision and hierarchical depth-wise aggregation so sites train according to local budgets and select inference depth at deployment.

    Source: §I · p. 2, C. Federated Anytime UNet · p. 4

  • Results

    Fed-ADApt matched FedAvg at a WT Dice of 0.82 in 3D and achieved a Dice of 0.81 in 2D versus 0.85 for FedAvg, while federated collaboration outperformed non-collaborative training across both tasks.

    Source: §III · pp. 5, 7, §IV · p. 7

  • Takeaways & Limitations

    Fed-ADApt supports participation by low-compute sites and produces a global multi-depth model for deployment-time depth selection without retraining.

    Source: §V · p. 8

    Evaluation used simulated compute assignments rather than deployment on real IoMT or edge hardware.

    Source: §IV · p. 7

Abstract

from arXiv · show

Federated learning (FL) enables collaborative training of medical image segmentation models without sharing raw patient data, yet existing approaches assume a homogeneous compute budget across institutions, limiting participation of low-resource sites. We propose Fed-ADApt, a depth-adaptive federated framework for UNet-based segmentation that jointly addresses low-compute training and inference. Fed-ADApt integrates multi-depth supervision with hierarchical depth-wise aggregation, allowing each site to train according to its local compute budget while contributing to a global model that supports dynamic depth selection at deployment. We evaluated Fed-ADApt on multi-site 2D retinal fundus disc segmentation and 3D brain tumor segmentation. Across both tasks, federated collaboration substantially improves robustness under domain shift. Fed-ADApt matched the full-resource FedAvg performance in 3D and achieved competitive 2D performance with a 4.7% average Dice reduction, while reducing average inference cost by 19.5% in 3D and 34.5% in 2D and substantially reducing training cost by 98% at the most constrained sites. Importantly, Fed-ADApt enables low-resource institutions that cannot train full-capacity models to participate in federations while maintaining competitive global performance under a favorable accuracy to efficiency trade-off. By considering training and inference compute budgets, Fed-ADApt provides a practical and equitable solution for federated medical image segmentation across heterogeneous clinical and edge-enabled imaging environments.

I. INTRODUCTION

Federated medical segmentation must accommodate sites with different training and deployment budgets, but existing approaches largely assume similar resources. Fed-ADApt combines budget-aware training with depth-adaptive inference for UNet segmentation.

I. INTRODUCTION

  • Existing federated learning approaches assume sites train and deploy identical models with similar compute resources, limiting participation by low-resource institutions.The need is especially relevant where local hardware cannot train or deploy full-capacity segmentation models.Source: §I · p. 2
  • Training-time and inference-time heterogeneity have typically been treated separately, leaving no unified framework for UNet-based medical image segmentation.UNet depth adaptation must preserve encoder–decoder compatibility, skip-connected feature transfer, and full-resolution prediction.Source: §I · p. 2
  • Fed-ADApt adds multi-depth supervision and hierarchical depth-wise aggregation so sites train to local capacity while contributing to a global model with dynamic depth selection.The framework also supports lower communication costs and anytime inference at deployment.Source: §I · p. 2
  • Across two multi-site datasets, Fed-ADApt matches FedAvg Dice performance in 3D and offers a favorable accuracy-efficiency trade-off in 2D while reducing resource costs.The evaluations cover retinal fundus and brain tumor segmentation.Source: §I · p. 2

II. METHOD

The method formulates federated segmentation across clinical sites with different compute budgets. Sites optimize a shared task while accessing only the model parameters their training budgets permit.

II. METHOD

  • Each clinical site has a local dataset of input images and segmentation masks, and sites differ in their normalized training budgets.A budget of 5 denotes full-capacity training; deployment inference may have a separate budget because forward-only inference is cheaper than training.Source: §II · p. 2
  • The objective is to learn a federated segmentation model that supports heterogeneous training and inference budgets while maintaining performance across sites.All sites solve the same segmentation task and share a common label space.Source: §II · p. 2
  • The proposed framework lets each site train only depths allowed by its local budget while contributing to a global model with dynamic depth selection at deployment.This jointly addresses the training and inference constraints illustrated for Fed-ADApt.Source: Figure 1 · p. 3

A. Multi-depth UNet

The multi-depth UNet creates prediction paths from different encoder depths, each decoding its available features to a full-resolution segmentation. These paths can be selected independently at inference to adapt computation.

A. Multi-depth UNet

  • A five-stage UNet is extended to produce predictions from multiple encoder depths, with each selected representation decoded to full resolution.For depth l, the representation z_l is formed by the encoder blocks through E_l.Source: A. Multi-depth UNet · p. 3
  • Each depth-specific path uses the corresponding decoder blocks and skip connections from available encoder features; depth 5 is the standard full UNet.Shallower paths require less computation.Source: A. Multi-depth UNet · p. 3
  • Unlike conventional auxiliary deep-supervision heads, every path produces a full-resolution output that can be selected independently at inference.These outputs define compute-adaptive inference paths rather than serving only as training signals.Source: A. Multi-depth UNet · p. 3

B. DepthFL UNet

DepthFL UNet adapts federated training by having each site train a single model depth set by its compute budget. The server aggregates parameters by depth, but inference remains fixed at full depth.

B. DepthFL UNet

  • Each site trains a single-depth UNet whose maximum encoder depth is determined by its training budget, using corresponding decoder blocks to produce one prediction.Predictions at other depths are unavailable to that site.Source: B. DepthFL UNet · p. 3
  • Supervision is applied only at the site's maximum supported depth, and updates to each depth's parameters come only from sites with sufficient training budgets.Deeper models contain the shallower model's parameters plus additional encoder and decoder components.Source: B. DepthFL UNet · p. 3
  • The server aggregates parameters independently by depth, using the sites eligible to train each depth.Eligibility is determined by the training budget relative to that depth.Source: B. DepthFL UNet · p. 4
  • DepthFL addresses heterogeneous training capacity but not compute-adaptive inference, because sites use the global model at maximum depth after training.This inference behavior is independent of each site's deployment-time inference budget.Source: B. DepthFL UNet · p. 4

C. Federated Anytime UNet

Federated anytime UNet jointly optimizes predictions at multiple depths, then lets each site select an inference depth according to its compute budget. Training all depths with FedAvg is resource-intensive.

C. Federated Anytime UNet

  • All encoder depths produce supervised predictions, jointly optimizing multiple depths during federated training.The model includes prediction paths at depths 1 through 5.Source: C. Federated Anytime UNet · p. 4
  • At inference, each site selects the depth with the highest validation score that fits its compute budget, without modifying or retraining the model.Source: C. Federated Anytime UNet · p. 4
  • Training every depth requires all sites to train the full model and update all parameters using FedAvg, making optimization resource-intensive.Source: C. Federated Anytime UNet · p. 4
  • The supervision encourages rough predictions at early depths and refinement by deeper depths.Source: C. Federated Anytime UNet · p. 4

D. Federated Anytime Depth Adaptation (Fed-ADApt)

Fed-ADApt limits each site's training depth to fit its local compute budget while aggregating depth-specific parameters hierarchically. Its global model supports depth selection at inference under a separate deployment budget.

D. Federated Anytime Depth Adaptation (Fed-ADApt)

  • Sites instantiate only the maximum network depth their training budget permits, limiting the number of supervised depths.Source: D. Federated Anytime Depth Adaptation… · p. 4
  • Each site trains encoder and decoder blocks with prediction paths through its permitted depth, supervising all available depths with the multi-depth loss.Source: D. Federated Anytime Depth Adaptation… · p. 4
  • Depth-wise aggregation lets sites with sufficient training budgets contribute corresponding global parameters, transferring shallow-layer knowledge from resource-rich to low-resource sites.Deeper layers are specialized using updates from high-budget institutions.Source: D. Federated Anytime Depth Adaptation… · p. 4
  • At deployment, each site selects a prediction depth within its inference budget, which may permit a deeper path than its training budget.The inference choice is based on the best validation score within that budget.Source: D. Federated Anytime Depth Adaptation… · p. 4

III. EXPERIMENTS AND RESULTS

Fed-ADApt was evaluated on multi-site 2D retinal disc and 3D brain tumor segmentation, with Dice performance compared against federated and local baselines. It preserved full-resource FedAvg performance in 3D and remained competitive in 2D while reducing compute costs.

III. EXPERIMENTS AND RESULTS

  • Evaluation covered two multi-site tasks, compared FedAvg and local UNets, and measured site-averaged Dice and hardware-agnostic GMAC compute cost.The 2D task used five clinical sites; the 3D task used five clinical centers.Source: §III · pp. 4–5
  • At the lowest-resource 2D site, Fed-ADApt improved Dice to 0.86 from 0.20 with local training and 0.45 with DepthFL.At the smallest site, training cost fell 96% versus FedAvg and 98% versus federated anytime UNet; inference cost fell 34.5% versus FedAvg.Source: §III · pp. 5–6
  • Fed-ADApt reached a Dice of 0.81 versus 0.85 for FedAvg in 2D, and matched FedAvg's 3D WT Dice at 0.82.The reported comparisons are for retinal disc Dice and whole-tumor (WT) Dice, respectively.Source: §III · pp. 5, 7, Table II · p. 6, Table III · p. 6
  • As federation budgets increased, Dice rose from 0.80 to 0.82 and 0.84; Site C rose from 0.86 to 0.87 and 0.89, while Site E remained largely stable.The federation-level trend shows that Fed-ADApt could exploit added client compute capacity.Source: §III · pp. 6–7
  • In 3D, Fed-ADApt reduced inference cost by 22% at the smallest site and communication cost by 25% versus FedAvg.At Site T, training cost decreased 72% versus FedAvg, from 477 to 133 GMAC; average estimated training VRAM was 2.43 GB versus 3.76 GB for federated anytime UNet.Source: §III · p. 7

IV. DISCUSSION

Fed-ADApt addresses training and inference heterogeneity in federated segmentation, with competitive performance and reduced resource costs. Its results also expose site-dependent capacity gains and deployment boundaries.

IV. DISCUSSION

  • Federated collaboration outperformed non-collaborative training across both tasks, while Fed-ADApt maintained competitive worst-site performance versus full-resource FL at reduced compute and communication cost.The discussion reports stronger cross-site robustness, especially for retinal segmentation, while brain-tumor data showed fewer shifts due to high-quality preprocessing.Source: §IV · p. 7
  • Low-resource sites performed poorly with shallow networks, while the largest site's high-resource model generalized poorly under cross-site variability, especially in retinal segmentation.Brain-tumor data showed fewer shifts, which the authors attribute to high-quality preprocessing.Source: §IV · p. 7
  • Multi-depth supervision propagates gradients to shape early representations for truncated inference, unlike DepthFL's supervision of only each site's maximum instantiated depth.The authors say this gives shallower layers consistent guidance from higher-compute sites.Source: §IV · p. 7
  • Restricting network depth reduces estimated memory for parameters, activations, gradients, and optimizer states, addressing a practical barrier at constrained sites.The authors identify reduced estimated VRAM as relevant to edge and IoMT deployment, while noting that realized hardware benefits remain untested.Source: §IV · p. 7
  • Increasing the federation's compute budget improved Dice, but gains plateaued at some sites, including the 2D task at Site E.The discussion suggests strategically allocating compute rather than increasing it uniformly.Source: §IV · p. 7
  • Fed-ADApt assumes predefined budgets; evaluation used simulated assignments, selected inference depth by validation performance, and has not established performance on other hierarchical architectures.The authors identify real IoMT deployment and extensions to architectures such as UNETR and SwinUNETR as future work.Source: §IV · p. 7

V. CONCLUSION

Fed-ADApt adapts federated UNet segmentation to heterogeneous training budgets through depth-specific prediction paths and hierarchical aggregation. Across retinal-disc and brain-tumor tasks, it preserves competitive performance while reducing resource requirements and costs.

V. CONCLUSION

  • Fed-ADApt defines depth-specific full-resolution prediction paths, preserves encoder–decoder skip connections, and hierarchically aggregates parameters according to each client's training budget.This UNet-compatible design allows low-compute sites to participate without full-capacity model optimization.Source: §V · pp. 7–8
  • Across 2D retinal disc and 3D brain tumor segmentation, Fed-ADApt preserved competitive performance while lowering resource requirements for constrained clients and inference and communication costs versus full-model FL.The conclusion characterizes this as a practical step toward equitable federated segmentation across heterogeneous clinical and edge-enabled environments.Source: §V · p. 8
  • The global multi-depth model supports deployment-time depth selection, letting sites balance inference cost and segmentation accuracy without retraining.The framework is presented as a practical step toward equitable federated medical image segmentation across heterogeneous clinical and edge-enabled settings.Source: §V · p. 8
Loading 2610.03474v1…