BIRD+: Design of a Lightweight Communication Compressor for Resource-Constrained Distribution Learning Platforms

The Top-K sparsification-based compression framework is extensively explored for reducing communication costs in distributed learning. However, we identified several issues with existing Top-K sparsification-based compression methods: (<i>i</i>) The limited compressibility of the Top-K parameter's indexes critically restricts the overall communication compression ratio; (<i>ii</i>) Several time-consuming compression operations significantly offset the benefits of communication compression; (<i>iii</i>) The use of error feedback techniques to maintain model quality results in a high memory footprint consumption. To solve these issues, we propose BIRD, a lightweight tensor-wise <i>Bi-Random sampling</i> strategy with an expectation invariance property. Specifically, BIRD applies a tensor-wise <i>index sharing</i> mechanism that reduces the index proportion by allowing multiple tensor elements to share a single index, thus improving the overall compression ratio. Additionally, BIRD replaces the time-consuming Top-K sorting with a faster <i>Bi-Random sampling</i> strategy based on the aforementioned <i>index sharing</i> mechanism, significantly reducing compression overheads; Moreover, BIRD establishes an <i>expectation invariance</i> property into the <i>Bi-Random sampling</i> to ensure an approximate unbiased representation for the <inline-formula><tex-math notation="LaTeX">$L_1$</tex-math></inline-formula>-norm of the sampled tensors, effectively maintaining the model quality without incurring extra memory costs. We further optimize BIRD to BIRD+ by introducing the uniform distribution-based sampling and Gamma correction on the tensor-wise sampling process, achieving a more flexibly adjustment of the sparsity with better convergence performance. Experimental evaluations across multiple conventional distributed learning tasks demonstrate that compared to state-of-the-art approaches, BIRD+ achieves higher communication compression ratios up to 36.2<inline-formula><tex-math notation="LaTeX">$\times$</tex-math></inline-formula> and higher computation throughput up to 149.6<inline-formula><tex-math notation="LaTeX">$\times$</tex-math></inline-formula> while maintaining the model quality without incurring extra memory costs.