Computationally efficient decoupled momentum optimization algorithm for medical imaging models

The goal of this study is to empirically evaluate the Decoupled Momentum Optimizer (DeMo) in medical image segmentation while demonstrating its extensibility to applications outside LLMs. We aim to characterize the behavior of each parameter group and their adherence to conjectures underlying DeMo's function. DeMo leverages spatial redundancy in gradients through a spatially partitioned frequency decomposition compression algorithm, reducing network traffic and smoothing gradient noise. DeMo provides up to a 150x traffic reduction and 1.6x wall-time speedup on lung segmentation of COPDGene CTs. Analysis of gradients support the conjectures that the primary components of the gradient exhibited higher spatial autocorrelation and lower temporal variance. We find that these conjectures are not uniformly true across all parameters, but rather are predominantly observed in a small subset of them. We also introduce DeMoDropout, a modification to the algorithm that selectively compresses only the largest gradients to significantly reduce computational overhead while maintaining effective overall compression. Using the Beyond the Cranial Vault dataset, we demonstrate potential speed-ups at bandwidths of 1 Gb/s and 100 Mb/s (1.6x vs 1.5x and 6.151 vs 6.31x for DeMoDropout and DeMo, respectively).