Robust Federated Learning Through Representation Matching and Adaptive
  Hyper-parameters
  Federated learning is a distributed, privacy-aware learning scenario which
trains a single model on data belonging to several clients. Each client trains
a local model on its data and the local models are then aggregated by a central
party. Current federated learning methods struggle in cases with heterogeneous
client-side data distributions which can quickly lead to divergent local models
and a collapse in performance. Careful hyper-parameter tuning is particularly
important in these cases but traditional automated hyper-parameter tuning
methods would require several training trials which is often impractical in a
federated learning setting. We describe a two-pronged solution to the issues of
robustness and hyper-parameter tuning in federated learning settings. We
propose a novel representation matching scheme that reduces the divergence of
local models by ensuring the feature representations in the global (aggregate)
model can be derived from the locally learned representations. We also propose
an online hyper-parameter tuning scheme which uses an online version of the
REINFORCE algorithm to find a hyper-parameter distribution that maximizes the
expected improvements in training loss. We show on several benchmarks that our
two-part scheme of local representation matching and global adaptive
hyper-parameters significantly improves performance and training robustness.
