Three Common Model Aggregation Methods in Federated Learning (FL) with TensorFlow Examples

Source：DeepHub IMBA

This article is approximately 1200 words, recommended reading for 7 minutes

Federated Learning (FL) is an excellent ML method that enables multiple devices (such as Internet of Things (IoT) devices) or computers to collaborate in model training without sharing their data.

The “clients” are the computers and devices used in FL, which can be completely isolated from each other and have their own distinct data that can apply the same privacy policies and are owned by different organizations, with no access to each other’s data.

With FL, models can learn from a broader data source without accessing the data. The widely used areas of FL are as follows:

• Healthcare
• Internet of Things (IoT)

• Mobile Devices

Since data privacy is a significant concern for many applications (such as medical data), FL is primarily used to protect clients’ privacy without sharing their data with any other clients or parties. Clients in FL share their model updates with a central server to aggregate the updated global model. The global model is sent back to the clients, who can use it for predictions or other operations on local data.

Key Concepts of FL

Data Privacy: Applicable to sensitive or private data applications.

Data Distribution: Training is distributed across a large number of devices or servers; the model should be able to generalize to new data.

Model Aggregation: Updating models across different clients and aggregating to generate a single global model. The aggregation methods are as follows:

• Simple Average: Average across all clients.
• Weighted Average: Weighting based on the quality of the model or the quantity of training data before averaging each model.
• Federated Averaging: Useful in reducing communication overhead and helps improve the convergence of the global model considering the local data differences used for model updates.

• Hybrid Methods: Combining various model aggregation techniques mentioned above.

Communication Overhead: The transmission of model updates between clients and the server, requiring consideration of communication protocols and the frequency of model updates.

Convergence: A key factor in FL is the model’s convergence to a good solution concerning the distributed nature of the data.

Simple Steps to Implement FL

1. Define the model architecture.
2. Partition the data into client datasets.
3. Train the model on the client datasets.
4. Update the global model.

5. Repeat the learning process above.

TensorFlow Code Example

First, we set up a simple server:

 import tensorflow as tf  # Set up a server and some client devices server = tf.keras.server.Server() devices = [tf.keras.server.ClientDevice(worker_id=i) for i in range(4)]  # Define a simple model and compile it inputs = tf.keras.Input(shape=(10,)) outputs = tf.keras.layers.Dense(2, activation='softmax')(inputs) model = tf.keras.Model(inputs=inputs, outputs=outputs) model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])  # Define a federated dataset and iterate over it federated_dataset = tf.keras.experimental.get_federated_dataset(devices, model, x=X, y=y) for x, y in federated_dataset:     # Train the model on the client data     model.fit(x, y

Then we implement the model aggregation steps:

1. Simple Average

 # Average the updated model weights model_weights = model.get_weights() for device in devices:     device_weights = device.get_weights()     for i, (model_weight, device_weight) in enumerate(zip(model_weights, device_weights)):         model_weights[i] = (model_weight + device_weight) / len(devices)  # Update the model with the averaged weights model.set_weights(model_weights)

2. Weighted Average

 # Average the updated model weights using weights based on the quality of the model or the amount of data used to train it     model_weights = model.get_weights()     total_weight = 0     for device in devices:         device_weights = device.get_weights()         weight = compute_weight(device)  # Replace this with a function that returns the weight for the device         total_weight += weight         for i, (model_weight, device_weight) in enumerate(zip(model_weights, device_weights)):             model_weights[i] = model_weight + (device_weight - model_weight) * (weight / total_weight)  # Update the model with the averaged weights     model.set_weights(model_weights)

3. Federated Averaging

 # Use federated averaging to aggregate the updated models model_weights = model.get_weights() client_weights = [] for device in devices:     client_weights.append(device.get_weights()) server_weights = model_weights for _ in range(num_rounds):     for i, device in enumerate(devices):         device.set_weights(server_weights)         model.fit(x[i], y[i])         client_weights[i] = model.get_weights()     server_weights = server.federated_average(client_weights)  # Update the model with the averaged weights model.set_weights(server_weights)

The above are the three basic model aggregation methods in federated learning, which I hope will be helpful to you.

Author: Dr. Roushanak Rahmat, PhD

Editor: Huang Jiyan