TLDR: Kai Yi’s dissertation introduces advanced strategies to enhance communication efficiency in distributed and federated learning. It presents EF-BV, a unified framework for biased and unbiased gradient compression; Scafflix, an algorithm for doubly accelerated federated learning through personalization and local training; FedP3, a privacy-friendly pruning framework for heterogeneous models; SPPM-AS, a method for reducing communication costs in cross-device federated learning by allowing multiple local communication rounds; and SymWanda and R2-DSnoT, novel post-training compression and training-free fine-tuning techniques for large language models.
In the rapidly evolving world of machine learning, training complex models often requires vast amounts of data and computational power. This has led to the rise of collaborative training methods like distributed learning (DL) and federated learning (FL). While these approaches are essential for handling large datasets and preserving privacy, they face a significant hurdle: communication overhead. Imagine trying to coordinate millions of devices or numerous data centers; the sheer volume of data exchanged can slow everything down.
A recent dissertation by Kai Yi, titled Strategies for Improving Communication Efficiency in Distributed and Federated Learning: Compression, Local Training, and Personalization, tackles this very challenge. The research offers a comprehensive exploration of innovative strategies to make these learning systems more efficient, focusing on three key areas: model compression, local training, and personalization.
Unifying Compression Techniques
One of the core contributions of this work is a new, unified theoretical framework for understanding and applying gradient compression. In distributed and federated learning, sending full model updates or gradients can be very costly. Compression techniques reduce the size of this information. Traditionally, there have been two distinct types of compressors: unbiased (like ‘rand-k’ which randomly selects and scales elements) and biased (like ‘top-k’ which keeps only the largest elements). Each had its own algorithms and theoretical underpinnings, such as DIANA for unbiased and EF21 for biased compression.
Kai Yi’s research introduces a new algorithm called EF-BV (Error Feedback with Bias-Variance decomposition). This algorithm successfully unifies both types of compressors under a single framework. It’s a significant step because it allows for the use of powerful biased compressors, while also benefiting from faster convergence rates when many workers are involved, a feature previously unique to unbiased methods. This means more flexibility and potentially faster training for a wider range of scenarios.
Accelerating Learning with Personalization and Local Training
Another major focus is on enhancing communication efficiency through a combination of local training and personalization. Local training involves clients performing multiple updates on their local data before communicating with a central server. This reduces how often data needs to be sent, but it can also lead to ‘client drift’ where local models diverge too much from the global goal, especially with diverse data.
Personalization, on the other hand, tailors the global model to better fit individual client data, which is crucial in federated learning where data distributions can vary widely (non-IID data). The dissertation introduces Scafflix, a novel algorithm that effectively integrates explicit personalization with accelerated local training. Scafflix is designed to achieve ‘double communication acceleration,’ meaning it speeds up training by leveraging both personalization and local computation. This is a significant advancement, as it allows for faster convergence while adapting models to individual client needs, outperforming existing algorithms in real-world learning setups.
Privacy-Friendly Pruning for Heterogeneous Models
The research also addresses the practical challenge of model heterogeneity, where clients might have different memory, processing power, or network bandwidth. This often means a ‘one-size-fits-all’ global model architecture isn’t ideal. The proposed FedP3 (Federated Personalized and Privacy-friendly network Pruning) framework offers a versatile solution. It incorporates both global pruning (from server to client) and local pruning (client-specific) strategies. Crucially, FedP3 is designed to be privacy-friendly, transmitting only selected segments of the global model back to the server after local training. This not only reduces communication costs but also helps conceal the full model structure, enhancing privacy. The framework also includes a local differential privacy variant, LDP-FedP3, with strong privacy guarantees and improved communication efficiency.
Optimizing Communication Rounds in Cross-Device Federated Learning
In cross-device federated learning, where millions of mobile devices participate, clients often operate in a stateless regime, meaning they can’t store information between communication rounds. Current methods typically involve a single communication round per client cohort. This research challenges that primitive, introducing SPPM-AS (Stochastic Proximal Point Method with Arbitrary Sampling). This method allows for multiple local communication rounds within a selected cohort before global aggregation. Surprisingly, this approach can lead to a significant reduction in total communication costs – up to 74% in some scenarios – to reach a desired accuracy. It also supports various client sampling strategies, including novel clustering-based techniques that further enhance performance.
Also Read:
- FediLoRA: Enhancing Federated Multimodal Fine-tuning with Adaptive LoRA Ranks and Missing Modality Handling
- Collaborative AI for Education: Addressing Privacy and Personalization with Federated Foundation Models
Symmetric Post-Training Compression for Large Language Models
Finally, the dissertation delves into the compression of large language models (LLMs), which are often too large for practical deployment. The work introduces SymWanda (Symmetric Weight And Activation), a novel formulation for post-training pruning (PTP). SymWanda minimizes the impact of pruning on both input activations and output influences of weights, providing theoretical backing for the empirical successes of methods like Wanda and RIA. It also proposes innovative pruning strategies, including a stochastic approach for relative importance that significantly reduces sampling costs while maintaining performance.
Furthermore, the research presents R2-DSnoT, a novel training-free fine-tuning method. This approach leverages relative weight importance and a regularized decision boundary within a pruning-and-growing framework, achieving state-of-the-art results in zero-shot performance without requiring additional retraining. This is particularly valuable for maintaining model robustness under high sparsity.
Overall, Kai Yi’s dissertation offers a comprehensive suite of strategies and algorithms that collectively push the boundaries of communication efficiency in distributed and federated learning. By addressing challenges in compression, local training, personalization, and model heterogeneity, this work contributes significantly to making advanced machine learning more scalable, efficient, and accessible across diverse, resource-constrained environments.
