TallyTrain: Communication-Efficient Federated Distillation
2026-06-30 • Machine Learning
Machine Learning
AI summaryⓘ
The authors address communication challenges in federated learning caused by large models and many output classes. They propose TallyTrain, which reduces communication by sending only the class with the highest score from each peer instead of full soft labels. This method can perform better when data is unevenly distributed across devices since majority voting helps filter errors. They also improve the approach by combining it with sparse model updates, outperforming common federated learning methods while using much less bandwidth.
federated learningmodel sizeclass countsoft-label distillationnon-IID trainingargmaxmajority votingsparse parameter mergeFedAvgFedProx
Authors
Radhakrishna Achanta, Will Reed
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.