TallyTrain: Communication-Efficient Federated Distillation

2026-06-30

AI Summary

This work addresses the communication bottleneck in federated learning caused by large model sizes and numerous classes, particularly under non-IID data distributions where soft-label distillation suffers from low efficiency and high susceptibility to noise. To overcome these limitations, the authors propose a hard-label federated distillation framework that transmits only the index of the highest-confidence predicted class (argmax) per round, reducing communication overhead to ⌈log₂C⌉ bits. By integrating sparse parameter synchronization with a majority-vote consensus mechanism, the method effectively filters noisy updates and enhances robustness. Combining strengths from FedAvg, FedProx, and FedDF, the approach matches or surpasses state-of-the-art soft-label distillation performance on standard benchmarks using less than one-thousandth of the communication cost, achieving Pareto optimality between communication efficiency and model accuracy across all evaluated configurations.

Abstract

Federated learning is bandwidth-bound on two orthogonal axes: model size, which limits how often parameter-averaging methods can afford to merge, and class count, which makes per-probe soft-label distillation prohibitive at large vocabularies. Both ceilings tighten as modern systems scale. We collapse the class-count axis to $\lceil \log_2 C \rceil$ bits per probe by transmitting only each peer's $\arg\max$ class index, where $C$ is the number of output classes. The resulting protocol, TallyTrain, is not merely compressed: under non-IID training it can be preferable to soft-label distillation, because under-trained peers are confidently wrong and majority voting filters this noise where soft-label averaging amplifies it. Across standard benchmarks, TallyTrain matches or beats soft-label distillation at up to three orders of magnitude less communication. We also relax the model-size axis: we compose the cheap hard-label consensus with sparse parameter merges to obtain a bandwidth-bridge variant, which Pareto-dominates every tested operating point of the standard FedAvg, FedProx and FedDF baselines.

Problem

• federated learning
• communication efficiency
• model size
• class count
• federated distillation

Innovation

• federated distillation
• communication efficiency
• hard-label consensus
• non-IID learning
• sparse parameter merging