[MM 2024] FedBCGD: Breakthrough in Communication-Efficient FL via Accelerated Block Coordinate Descent

TL;DR

Training large-scale models like Vision Transformers (ViT) in Federated Learning (FL) is notoriously slow due to the massive size of model updates being sent over asymmetric networks. FedBCGD (Federated Block Coordinate Gradient Descent) solves this by partitioning the model into $N$ blocks and only requiring clients to upload one block plus a tiny shared block. Unlike previous attempts that froze parameters locally, FedBCGD updates everything but transmits fractionally, achieving up to 11.5x faster convergence on complex datasets like Tiny ImageNet.

The "Freezing" Trap: Why Naive Sparsification Fails

In the quest to reduce communication overhead, a natural instinct is to have each client train only a small portion of the model.

• The Intent: Reduce upload traffic by $1/N$.
• The Reality: "FedBCGD_freezing" (updating only the active block and freezing the rest) leads to disastrous convergence.

As the authors identify, freezing parameters locally creates a massive drift between blocks. Because the frozen blocks don't adapt to the changes in the active block, the global model becomes a fragmented collection of weights that haven't "learned" to work together. This creates a high-variance gradient environment that stalls training.

Methodology: The FedBCGD Insight

The core philosophy of FedBCGD is "Update All, Upload One."

1. Model Partitioning

The model $x$ is divided into $N$ parameter blocks plus a shared block $x_s$.

• $x_s$ (The Shared Small Block): Usually the last layer (classifier). In ResNet-18, this represents only 0.01% of parameters but is crucial for alignment.
• $x_{(j)}$ (The Assigned Block): Each client is assigned a specific block to "responsible" for uploading.

2. Full Local Optimization

Clients do NOT freeze parameters. They perform SGD on the entire model. This ensures that the dependencies between layers are maintained.

3. Server-Side Momentum Compensation

Since the server only receives updates for specific blocks from specific clients, it uses a Block Momentum term ($v^r$) to "remember" and smooth out the update directions of the blocks that weren't updated in the current round.

Figure 1: The FedBCGD framework showing client-block allocation.

FedBCGD+: Taming Heterogeneity

Data in FL is rarely IID (Independent and Identically Distributed). To handle "Non-IID" data, the authors propose FedBCGD+, which integrates:

• Client Drift Control: Inspired by SCAFFOLD, using control variates ($c$) to correct the local update direction towards the global optimal.
• Variance Reduction: Using SVRG-like techniques to handle the noise introduced by stochastic gradients.

The local update rule in FedBCGD+ becomes: $x_{k,j}^{r,t+1}=extSGDTerm+extClientDriftCorrection+extVarianceReduction$

Experimental Showdown

The performance gains are most visible when looking at the Communication Floats required to reach a target accuracy.

1. Speedup on Vision Transformers (ViT)

On the Tiny ImageNet dataset using a ViT-Base model, FedBCGD reached the target accuracy over 11 times faster than standard FedAvg.

Figure 2: Test accuracy vs. Communication floats for ViT-Base.

2. SOTA Complexity Comparison

As shown in Table 1 of the paper, FedBCGD+ achieves the lowest non-convex communication complexity compared to heavyweights like SCAFFOLD and FedLin.

Critical Perspective: Why Does This Work?

The physical intuition here is that model updates are redundant. In a high-dimensional parameter space, you don't need every weight update in every round to progress toward the manifold of the global minimum. However, you do need the local model to remain "internally consistent" during training. By allowing the client to update all parameters but only communicate a block, FedBCGD preserves the Jacobian structure of the network locally while drastically thinning the pipe back to the server.

Limitations

• Partitioning Sensitivity: The paper uses fixed-layer partitioning. However, optimal partitioning might depend on the specific architecture (e.g., Attention heads vs. MLP layers in Transformers).
• Memory Overhead: Since clients still update the full model, there is no reduction in local computation or memory—only in communication.

Conclusion

FedBCGD provides a powerful new toolkit for Federated Learning on large models. By moving the "sparsity" from the training phase to the communication phase, it bypasses the accuracy degradation typical of sparsified FL while delivering the bandwidth savings promise of Block Coordinate Descent.