receiving, at the client device, a global model from a server;

performing model inversion on the global model to generate synthetic data for one or more classes of data in the global model, wherein performing model inversion comprises performing at least one of image distillation or zero-shot data generation (ZSDG) on the global model;

augmenting collected data in a dataset of the client device with the synthetic data to generate an augmented dataset comprising the collected data and the synthetic data, the augmented dataset having a more uniform distribution of data across classes than the dataset;

training the global model on the augmented dataset to generate a client model; and

transmitting the client model from the client device to the server.