Icon Legend

This session is not in your schedule.

This session is in your schedule. Click again to remove it.

Bioinformatics, Computational and Systems Biology

Bioinformatics, Computational and Systems Biology - Poster Session E

M16 - Evaluation of Decoupled Momentum Optimization for Distributed Medical Image Segmentation

Saturday, October 11, 2025

10:00 AM - 11:00 AM PDT

Location: Exhibit Hall F, G & H

Presenting Author(s)

Jasmine N. Jung

Undergraduate Research Assistant
University of Texas at Austin
Del Valle, Texas, United States

Co-Author(s)

Aaron Luong

Undergraduate Researcher
University of Texas at Austin, United States
Joshua R. Joseph

Undergraduate Researcher
University of Texas at Austin
Coppell, Texas, United States
DC

Debarghya Chaki

Undergraduate Research Assistant
University of Texas at Austin, United States
YL

Yi-Kuan Liu

PhD Candidate
The Department of Biomedical Engineering, The University of Texas at Austin, United States
AN

Amanda Nowacki

Research Assistant
University of Texas at Austin, United States
BE

Blake Evans

Graduate Student
The University of Texas at Austin
Austin, Texas, United States

Last Author(s)

Edward Castillo, PhD

Associate Professor & Associate Chair
The University of Texas at Austin
Austin, Texas, United States

Introduction: : Distributed training is essential for many AI applications, including 3D medical image segmentation. However, communication overhead remains a major bottleneck. Optimizers like AdamW synchronize full gradient and momentum tensors at each step, which is inefficient in bandwidth-constrained environments. The Decoupled Momentum (DeMo) optimizer addresses this issue by separating momentum accumulation from synchronization, transmitting only fast-moving components identified via frequency decomposition.

Although theoretically promising, DeMo has not been empirically evaluated in high-dimensional imaging tasks. This study benchmarks DeMo against AdamW, SGD, and DiLoCo for lung lobe segmentation on the COPDGene(R) dataset using the SwinUNETR architecture.

We empirically assess three theoretical conjectures underlying DeMo: (1) that fast-moving momentum components exhibit higher spatial autocorrelation (2) and show lower temporal variance than slow-moving components, and (3) that local accumulation of slow-moving components improves convergence.

Our contributions are threefold: 1) implement DeMo for high-dimensional segmentation tasks using SwinUNETR 2) evaluate DeMo’s performance across various compression settings and 3) validate its core theoretical properties using gradient diagnostics, communication metrics, and segmentation outcomes.

Our results demonstrate that DeMo improves communication efficiency without compromising accuracy, suggesting practical utility in bandwidth-sensitive or federated training settings.

Materials and

Methods: : We trained SwinUNETR models for 3D lung lobe segmentation using 8,702 breath-hold CT scans from the COPDGene(R) dataset downsampled to 128³ voxel resolution. Training was performed using PyTorch’s DistributedDataParallel (DDP) module across two compute nodes, each with dual high-memory NVIDIA GPUs and high-speed interconnects (40 Gb/s InfiniBand), suitable for communication-intensive workloads.

We compared AdamW, SGD, DiLoCo, and a custom DeMo implementation. DeMo uses segment-wise Discrete Cosine Transform (DCT) to extract k dominant frequencies from s-sized chunks of the momentum tensor. We evaluated 16 (k, s) combinations with k ∈ {1, 4, 16, 64} and s ∈ {4, 16, 64, 128}. AdamW was applied with a learning rate of 5e-4 and weight decay of 0.01; DeMo employed the same learning rate and no weight decay.

Metrics recorded included: 1) Dice Similarity Coefficient (DSC) on a held-out validation set, 2) training loss as a function of iteration, 3) network communication (bytes sent and received), 4) wall time per training step, and 5) gradient diagnostics, including Moran’s I (spatial autocorrelation) and per-component temporal variance.

Gradients were saved every 40 iterations to assess the theoretical conjectures. Fast and slow components were compared in terms of spatial autocorrelation and temporal variance. To evaluate Conjecture 3, we trained DeMo with and without slow component accumulation and compared final convergence metrics. All visualizations were generated using matplotlib and Seaborn libraries.

Results, Conclusions, and Discussions:: DeMo achieved comparable segmentation performance and improved computational efficiency relative to AdamW while significantly reducing communication overhead and training time. The best-performing DeMo runs (e.g., k=16, s=4) reached Dice scores above 0.976—on par with AdamW—while cutting average network traffic per step by over 50-fold and decreasing wall time by 1.6-fold.

Training loss and Dice curves showed that DeMo converged more quickly than AdamW in 5 out of 16 configurations. Notably, DeMo avoided the sharp loss oscillations observed in AdamW, maintaining smoother training dynamics. Final segmentation accuracy was preserved in 12 of 16 (k, s) settings, with Dice drops under 0.002 compared to AdamW. Overcompressed configurations (e.g., k=1, s=64 or 128) led to slower convergence, confirming a tradeoff between compression and performance.

Communication metrics measured via NVIDIA Collective Communications Library (NCCL) confirmed a sharp reduction in transmitted data without convergence degradation when appropriate (k, s) values were selected.

Gradient analysis supported all three conjectures:
Conjecture 1: Fast components consistently exhibited higher Moran’s I than both original and slow components, indicating stronger spatial autocorrelation.
Conjecture 2: Fast components showed lower temporal variance than slow components throughout training.
Conjecture 3: Removing the slow component accumulation resulted in significantly worse convergence (higher loss, lower final Dice), confirming their long-term importance.

This study demonstrates that DeMo is a communication-efficient and high-performing optimizer for 3D medical image segmentation in distributed environments. Compared to AdamW, it can reduce network traffic by up to 150-fold and accelerate training by up to 1.6-fold while maintaining comparable segmentation accuracy, with slight Dice improvements in optimal settings.

By directly measuring spatial autocorrelation, temporal variance, and convergence behavior, we empirically validated DeMo’s theoretical foundations. These results establish DeMo as a practical choice for bandwidth-constrained or federated learning scenarios. DeMo requires no architectural changes and integrates seamlessly with existing training pipelines, making it well-suited for privacy-sensitive and multi-institutional applications. Future work will explore scaling DeMo to larger imaging models and incorporating dynamic compression schedules that adapt (k, s) during training to optimize performance and communication jointly.

Acknowledgements and/or References (Optional): : 1. Peng, B.; Quesnelle, J.; Kingma, D.P. DeMo: Decoupled Momentum Optimization, 2024, [arXiv:cs.LG/2411.19870].
2. Hatamizadeh, A.; Nath, V.; Tang, Y.; Yang, D.; Roth, H.; Xu, D. Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images, 2022, [arXiv:eess.IV/2201.01266].
3. Wang, Z.; Wen, M.; Xu, Y.; Zhou, Y.; Wang, J.H.; Zhang, L. Communication compression techniques in distributed deep learning: A survey. Journal of Systems Architecture 2023, 142, 102927. https://doi.org/https://doi.org/10.1016/j.sysarc.2023.102927.
4. Bernstein, J.; Wang, Y.; Azizzadenesheli, K.; Anandkumar, A. signSGD: compressed optimisation for non-convex problems. CoRR 2018, abs/1802.04434, [1802.04434].
5. Wen, W.; Xu, C.; Yan, F.; Wu, C.; Wang, Y.; Chen, Y.; Li, H. TernGrad: Ternary Gradients to Reduce Communication in Distributed Deep Learning, 2017, [arXiv:cs.LG/1705.07878].
6. DeepSeek-AI.; Liu, A.; Feng, B.; Xue, B.; Wang, B.; Wu, B.; Lu, C.; Zhao, C.; Deng, C.; Zhang, C.; et al. DeepSeek-V3 Technical Report, 2025, [arXiv:cs.CL/2412.19437].
7. Zhao, J.; Zhang, Z.; Chen, B.; Wang, Z.; Anandkumar, A.; Tian, Y. GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection, 2024, [arXiv:cs.LG/2403.03507].
8. Renggli, C.; Alistarh, D.; Hoefler, T. SparCML: High-Performance Sparse Communication for Machine Learning. CoRR 2018, abs/1802.08021, [1802.08021].
9. Bereyhi, A.; Liang, B.; Boudreau, G.; Afana, A. Regularized Top-k: A Bayesian Framework for Gradient Sparsification, 2025, [arXiv:cs.LG/2501.05633].
10. Li, Y.; He, Z.; Gu, X.; Xu, H.; Ren, S. AFedAvg: Communication-efficient federated learning aggregation with adaptive communication frequency and gradient sparse. Journal of Experimental & Theoretical Artificial Intelligence 2024, 36, 47–69.
11. Han, S.; Mao, H.; Dally, W.J. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiv:1510.00149 2015.
12. Kingma, D.P.; Ba, J. Adam: A Method for Stochastic Optimization, 2017, [arXiv:cs.LG/1412.6980].
13. Loshchilov, I.; Hutter, F. Decoupled Weight Decay Regularization, 2019, [arXiv:cs.LG/1711.05101].
14. Yamada, H. Moran’s I for Multivariate Spatial Data. Mathematics 2024, 12. https://doi.org/10.3390/math12172746.