PFedMA: A Federated Learning Algorithm with Pruned Matched Averaging

Federated aggregation-based methods such as FedAvg and FedProx perform aggregation across the same parameter indices of locally trained client models. This assumes that the clients represent comparable features at the same indices, which is not necessarily true, as fully connected and convolutional layers exhibit permutation invariance with respect to neurons or filters within a layer. FedMA was introduced to address this mismatch by probabilistically matching neurons or channels layer by layer before performing aggregation. However, unmatched units from client models are appended to the global model. This causes layer-wise growth, which is particularly difficult to accommodate in networks such as ResNets, where channel dimensions and orderings are constrained by cross-layer connections. We introduce Pruned Federated Matched Averaging (PFedMA), a variant of FedMA that preserves a fixed model topology. Rather than retaining every unmatched unit, PFedMA discards unmatched non-donor units and fills vacant global positions using unmatched units from a randomly selected donor client. We further extend matched aggregation to be compatible with residual networks. We do this by enforcing consistent channel assignment across Batch Normalization layers and residual connections. We evaluate PFedMA on four image-classification benchmarks under IID and multiple Dirichlet non-IID partitions, using both sequential CNNs and ResNet-18. FedMA obtains higher accuracy in some sequential-CNN settings. On CIFAR-10 with ResNet-18, PFedMA substantially outperforms the evaluated parameter-wise baselines across the tested partitions. These results suggest that residual networks can also benefit from permutation-aware aggregation.