Federated Learning with TensorFlow Federated

Last Updated : 23 Jul, 2025

Federated learning is a distributed machine learning technique that allows multiple devices to collaboratively train a shared model while keeping their data locally. TensorFlow Federated (TFF) is an open-source framework developed by Google for machine learning on decentralized data. It enables developers to implement and simulate federated learning algorithms with TensorFlow. This article provides an in-depth introduction to TensorFlow Federated, exploring its architecture, applications, advantages, and how to get started with it.

Table of Content

What is Federated Learning?

Federated learning is a decentralized learning approach that enables model training across multiple edge devices, such as smartphones or IoT devices, without transferring the data to a central server. It ensures data privacy by keeping personal data on the devices themselves.

Instead of sending raw data to a central server for model training, each device trains the model locally and sends the updated weights or parameters to a central server. The server aggregates these updates and improves the global model, which is then distributed back to the devices.

Key Features of Federated Learning

  • Data Privacy: Since the data never leaves the device, federated learning maintains privacy by design.
  • Efficiency: By minimizing the need to transmit large datasets over the network, federated learning reduces communication costs.
  • Personalization: Local training enables models to learn from unique user behavior or conditions, leading to more personalized models.

Overview of TensorFlow Federated

TensorFlow Federated (TFF) is a framework built on TensorFlow for building federated learning systems. TFF enables developers to define machine learning algorithms that can be executed across multiple devices in a distributed fashion.

Key Components of TFF

  • Federated Computation: TFF introduces a new programming model to express computations that can be distributed across devices. Federated computations are written in Python and define how local computations (on devices) and global computations (on the server) are composed.
  • Federated Learning API: The TFF framework includes high-level APIs for defining and training federated models. It allows for easy integration with existing TensorFlow models.
  • Simulation Environment: TFF comes with a built-in simulation environment, enabling developers to simulate federated learning scenarios locally before deployment.

Architecture of TensorFlow Federated

TFF's architecture is organized into two main layers:

  1. Federated Learning (FL) API: This high-level interface allows developers to apply federated training and evaluation to existing TensorFlow models. It simplifies the process by providing ready-to-use implementations of federated algorithms.
  2. Federated Core (FC) API: The FC API offers lower-level interfaces for expressing custom federated algorithms. It combines TensorFlow with distributed communication operators in a strongly-typed functional programming environment. This layer serves as the foundation for building both learning and non-learning federated computations.

How TensorFlow Federated Works

TFF allows you to write federated learning algorithms by creating two types of computations:

  • Federated Learning Process: This defines the overall training procedure, including how models are updated and how device and server computations interact.
  • Client Update Functions: These functions define how each device performs local training, such as computing gradients or loss based on its local data.

Here’s a simplified example of how TFF works:

  • The server initializes a global model and sends it to all participating devices.
  • Each device trains the model locally on its data and sends the updated weights back to the server.
  • The server aggregates the updates and refines the global model.
  • The updated model is sent back to the devices for the next round of training.
  • This process continues iteratively until the model reaches satisfactory performance.

Building a Simple Federated Learning Model

To demonstrate federated learning with TFF, we’ll use a simple example of training a model on a decentralized dataset. Let’s assume we have a collection of datasets across different devices.

To start using TensorFlow Federated, install it via pip:

!pip install tensorflow==2.10.0
!pip install tensorflow_federated

1. Loading and Preprocessing the Data

Here’s an example of a federated learning setup using the MNIST dataset, The MNIST dataset will be loaded, normalized, and reshaped for CNN input:

Python 
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
x_train = np.expand_dims(x_train, axis=-1)
x_test = np.expand_dims(x_test, axis=-1)

At this point, the training and test data will be normalized and reshaped to the correct format (28x28x1).

2. Simulating Clients and Creating TensorFlow Datasets

This part of the code simulates the federated environment by splitting the dataset into different clients:

Python 
NUM_CLIENTS = 10
client_data = []
data_per_client = len(x_train) // NUM_CLIENTS
for i in range(NUM_CLIENTS):
    start = i * data_per_client
    end = (i + 1) * data_per_client
    client_data.append((x_train[start:end], y_train[start:end]))

Each client will have a portion of the dataset for local training.

3. Training the Federated Learning Model

The Federated Averaging algorithm is used to perform federated learning for multiple rounds:

Python 
NUM_ROUNDS = 10
for round_num in range(NUM_ROUNDS):
    state, metrics = iterative_process.next(state, federated_train_data)
    print(f'Round {round_num+1}, metrics={metrics}')

Output:

Round 1, metrics={'sparse_categorical_accuracy': 0.912, 'loss': 0.302}
Round 2, metrics={'sparse_categorical_accuracy': 0.926, 'loss': 0.256}
Round 3, metrics={'sparse_categorical_accuracy': 0.935, 'loss': 0.224}
...
Round 10, metrics={'sparse_categorical_accuracy': 0.972, 'loss': 0.112}

4. Evaluating the Model

After the federated learning rounds, the global model will be evaluated on the centralized test dataset:

Python 
test_loss, test_acc = global_model.evaluate(x_test, y_test)
print(f'Test accuracy: {test_acc}')

Output:

Test accuracy: 0.977

Advantages and Challenges of Using TFF

Advantages

  • Flexibility: TFF supports a wide range of machine learning models and can be easily integrated with existing TensorFlow workflows.
  • Open Source: As an open-source framework, TFF allows researchers and developers to customize and extend its functionalities.
  • Simulation Capabilities: TFF provides tools for simulating federated learning environments, making it easier to test models before deploying them on real devices.

Challenges

  • Complexity: Implementing federated learning requires understanding both machine learning concepts and distributed computing principles.
  • Resource Intensive: Running simulations or real-world deployments can be resource-intensive due to the need for multiple client devices.

Use Cases of Federated Learning

Federated Learning has numerous applications across various industries:

  • Healthcare: Training models on sensitive patient data without compromising privacy.
  • Finance: Developing fraud detection systems using decentralized transaction data.
  • Mobile Applications: Improving predictive text models on smartphones without sending user typing data to servers.

Conclusion

TensorFlow Federated offers a robust platform for implementing federated learning algorithms, enabling privacy-preserving machine learning on decentralized data sources. By leveraging TFF, developers can build intelligent applications that respect user privacy while benefiting from collective insights across distributed datasets. As the field of federated learning continues to evolve, frameworks like TFF will play a crucial role in advancing research and real-world applications.

Comment

Explore