Federated Learning

In 2017, Researchers at Google published a paper^[1] titled “Communication-Efficient Learning of Deep Networks from Decentralized Data”, which introduced and popularised the idea of a Privacy Oriented, Distributed AI Training System in mainstream Computer Science research. This project later turned into what we today call Federated Learning.

Why do we need Privacy Oriented AI Training System?

AI Systems are Data hungry. Most models require Gigabytes to Terrabytes of Data, to learn patterns and understand the data. This poses a huge problem. For models that are specifically built for commercial usage, we need to extract the data generated by the commercial users. One consumer may not provide enough data to train a whole model. To solve this, many companies would silently record usage statistics and other details as consumers would use the application, and the application would send this collected data to large data storage facilities. When thousands to millions consumers send such statistics, the company would have enough data to finally train the AI model.

This however breaches privacy for a lot of users. Sharing personal data is strictly prohibited unless users “Opt-in” in many countries and regions. A lot of people, including governments are turning towards Privacy-First approach^[2]. For example, according to a few sources^[3], even as recent as 2021, only 11% of 5.3 million opted in for ad tracking for Apple’s IDFA technology.

Google initially faced a similar problem while building the Autocorrect/Next Word Prediction system in Google Keyboards. To predict words better, google required more data. This was one of the original inspiration to build the first mass scale Federated Learning System.^[4]

How does it work?

We learnt that Federated Learning Systems are distributed, and preserve privacy. How does that work?

In a Centralized solution, we send data from consumer machines to the center Server (we will refer to these machines that send information to the Server as “Clients” from now onwards), which aggregates the data and trains the AI. This architecture sacrifice network bandwidth, and saves on client’s sompute power.

However, in a Federated Learning architecture, instead of the data being transferred, the AI model parameters are sent to the Sever instead. Here, the clients host the AI models, which are then updated to the latest model parameters recieved from the Server. The Clients then train this model parameters recieved from the Server, based on the data it has collected. This new trained model’s parameters are then uploaded back to the Server. The Server aggregates all the clients’ parameters, and updates the global model. The cycle then repeats when the global model parameters are then sent back to the clients for a new epoch.

[Main Server]<-[<a id=a>Server]
[Main Server]<-[<a id=b>Client]
[Main Server]<-[<a id=c>Client]
[<a id=g> Server]<-[<a id=e>Client]
[<a id=g> Server]<-[<a id=f>Client]
[Main Server]<-[<a id=g> Server]
[<a id=zz> Server]<-[<a id=i>Client]
[<a id=zz> Server]<-[<a id=j>Client]
[<a id=a>]<-[<a id=zz> Server]
[<a id=a> Server]<-[<a id=h>Client]

In a Federated Learning (FL) architecture, we aren’t sharing the personal data, but instead we’re sharing the trained model’s details/parameters over the network. The Server aggregates all the model parameters it recieves from all the clients into a few (or one) models.

This aggregated model combines the learnt behaviour of all the clients. We then repeat this loop of sending the model, and recieving updated models from the clients for a few times. We then stop once the model has reached a certain threshold (such as accuracy of model > 90%).

Sending and recieving model updates’ bandwidth is a small fraction of sending original data’s bandwidth. For a small CNN Network, such as AlexNet as implemented in pytorch’s torchvision.models.alexnet, the model contains around 61 Million parameters. Each parameter, if taken to be float32, are around 4 bytes. Hence the total model size is [4 bytes * 61 Million] around 250MB. Assuming, we run a 100 rounds of sending the model to and fro the Server and 50 Clients, we get a bandwith of around [250MB * 2 trips to and from server * 100 rounds * 50 clients] 2,500GB spread around the entire train time for FL. This is way better than the centralized training solution where we send the training data of size [1 Million Images * 300KB per image * 50 clients] 15,000GB before training even begins, where we have to wait for clients to upload their data, breaching their privacy, and other issues.

We are essentially sacrificing the compute on the client’s side, while lowering bandwidth, lowering wasted time, and improving privacy of Clients. Sacrificing compute on smaller devices maybe worth the tradeoff, as the computing capabilities increases exponentially. Smaller Devices like Smart Phone nowadays are equivalent to computers back in 2000s+ in terms of Computing strength. A small AI model can easily be trained on your smart phone, given enough time.

What are the Cons of FL Architecture?

In a Centralised architecture, all the training data is transferred to a central data pool. The central model learns the data as provided, along with the respective biases in the training data. In a theoretical FL scenario, the training data is divided among the clients equally (independent and identically distributed data, i.e IID Data)

For an IID Dataset, Image recognition FL example:

• Client A has 110 cat images, 120 dog images, and 125 horse images.
• Client B has 90 cat images, 100 dog images, and 105 horse images.
• Client C has 100 cat images, 105 dog images, and 89 horse images.
• …

The model in each client learns the exact same patterns across various clients. The aggregated model perfectly advances its understanding of the data, as all client models train in the same direction. Think of it this way, image each model’s parameter-space to look like a landscape with hills and valleys, where a particular model is represented by a position in this landscape (Vector representation of Model Parameters). Currently all models move/learn in the same direction, as the data distribution looks similar. Therefore, the aggregate model on the Server highly matches all the trained clients.

Models initialised to a random location.

Each client trains its model. Most models move in a similar direction as data is IID.

All client models are all aggregated together. Hence a Round is complete, and the global model has advanced.

However, in practical applications of FL, the data distribution across all the clients is rarely Independent and Identically distributed (It is, Non-IID Data). Here, the training data in some clients are biased to some degree.

For a Non-IID Dataset, Image recognition FL example:

• Client A has 1 cat images, 170 dog images, and 0 horse images. [Client is located in a dog shelter]
• Client B has 2 cat images, 60 dog images, and 30 horse images. [Client is located in a big stable, with plenty dogs]
• Client C has 300 cat images, 0 dog images, and 35 horse images. [Client is once again located in a stable, that is infested with stray cats roaming about]
• …

Here, the models in each client learns their own data, with respect to the distribution of the data given to it. Imagine each client as a marker on the parameter-space landscape. Each client marker is moved in random directions, based on what it learns (each direction is a slightly different topic). The server now struggles to aggregate all the models together for the next round, as all the clients parameters are spaced widely apart.

This is one of the main problems with FL architecture, Non-IID/Heterogeneous Datasets ruins the training processes^[5].

Models initialised to a random location.

All client models are all aggregated together. Hence a Round is complete, but the global model has NOT advanced since the initial state. Models haven't converged, and no progress has been achieved.

Of course, there are other issues with FL, but this is the main issue FL faces.

Some of the other issues FL faces is, Client Unavailability (When some clients don’t participate, maybe due to low resources, Check out ADaRES either here or here), Privacy leakage^[6], Distributed Attacks^[7], etc.

So is FL the system, or the process?

FL architecture is just the distributed training enviroment for AI models. The exact algorithm which dictates how a round of training occurs can be different. There are various protocols available which facilitate training of AI models for various specific usecases in an FL system.

The simplest FL Training protocol is called FedAvg^[1]. I previously explained the process of FL to be very similar to FedAvg. Here’s how FedAvg actually works, written as an algorithm:

1. Server initialises an untrained model. It braodcasts the model parameters to all the clients in the training pool.
2. Each client updates it’s model with the parameters given by the server. All of them parallely train their model based on the training data they have.
3. After a few rounds of training, each model sends their model parameters back to the server.
4. The server aggregates all the parameters into a SINGLE model by taking the averaging of the parameters across the clients.
   Let 2 model’s parameters look like this:
   • Client 1: {"param1": 0.5, "param2": 0.3, ...}
   • Client 2: {"param1": 0.6, "param2": 0.2, ...}
     Then, the aggregated model on the server would look like this:
     {"param1": (0.6+0.5)/2, "param2": (0.2+0.3)/2, ...} => {"param1": 0.55, "param2": 0.25, ...}
5. We go back to STEP 1, by broadcasting this new model to all the clients.
6. We keep doing this, until our accuracy of the model stabilises.

The FedAvg protocol is very easy to implement, and works really well for IID Data^[1], as we’re just averaging the model parameters, where all models learn the same data.

However, it’s terrible under Non-IID settings^[1], where each model learns a different bias, and the averaging of all the biased learnt parameters doesn’t represent a good global model.

To fix this, other researchers have come up with various algorithms/protocols to help models generalize better. Some algorithms include SCAFFOLD^[8], FedProx^[9], etc.

When to use, and when NOT to use FL?

Federated Learning has a lot of pros and cons. Learning how the tradeoffs work will help you learn when to use FL architecture.

Use FL When:

• The model is to be used on Edge Devices, which also collects the training data.
• Data cannot be transferred either due to security policies or excess bandwidth usage.
• Most of the data is IID/Homogenous, this helps the model converge better.
• We can safely Orchestrate FL using the right algorithm, to block certain distributed attacks.

Do not use FL When:

• We can transfer the data to a central server safely, without any attacks or restrictions.
• We need to run fast experiments, or when orchastration is not feasible.
• All Data NEEDS to be processed together before any training occurs (For example, when all data need to go through a centralized pipeline, before training)
• When privacy is not an issue.

You can read more about Federated Learning on scholarly articles like this^[10].

TL;DR Code example

import torch
from torch import nn
import copy
...

type StateDict = dict

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        ...

    def forward(self, x):
        ...
        return x

class Client:
    def __init__(
        self, 
        client_id: int, 
        model: nn.Module, 
        dataset, 
        metadata
    ):
        self.client_id = client_id
        self.model = model
    
    def set_model_params(self, param: StateDict):
        self.model.load_state_dict(param)

    def train(self, ...):
        ...

    def get_model_param(self) -> StateDict:
        return self.model.state_dict()

class Server:
    def __init__(
        self, 
        model: nn.Module, 
        clients: list[Client]
    ):
        self.clients = clients
        self.global_model = copy.deepcopy(model) #untrained model

    def get_global_model_param(self) -> StateDict:
        return self.global_model.state_dict()

    def set_global_model_param(self, param: SateDict):
        self.global_model.load_state_dict(param)

    def send_model_to_clients(self):
        for client in self.clients:
            client.set_model_params(self.get_global_model_param)

    def get_client_models(self) -> list[StateDict]:
        return [client.get_model_param() for client in self.clients]

    def aggregate_fedavg(self, params: list[StateDict]) -> StateDict:
        # total for avg
        total = len(params)
        
        avg_weights = copy.deepcopy(params[0]) # copy of StateDict
        
        for key in avg_weights.keys(): # for each key in avg_param
            val = sum(c_param[key] for c_param in params)
            val /= total # take avg

            avg_weights[key] = val

        return avg_weights

    def single_train_round(self):
        self.send_model_to_clients()

        for client in self.clients:
            client.train() 
            # It happens parallely IRL
            # we are just simulating the train sequentially
        # wait until training for all clients is complete

        params = self.get_client_moels()
        global_param = self.aggregate_fedavg(params)
        
        self.set_global_model_param(global_param)

    def train(self, epochs: int):
        for epoch in range(1, epochs+1):
            self.single_train_round()
            print(f"Epoch {epoch}/{epochs} Done!")


def create_clients(num_clients) -> list[Client]:
    import datasets_for_clients as data
    clients = [
        Client(
            i, 
            Net(), 
            data.dataset(i), 
            data.metadata(i)
        ) for i in range(num_clients)
    ]
    return clients

def main():
    model = Net()

    clients = create_clients(50)
    # Creating 50+ copies of the Pytorch model 
    # would eat up your memory.
    # This is just a simple Tl;dr simulation of 
    # how clients might look like.

    server = Server(model, clients)

    server.train(epochs=100)

    model.load_state_dict(server.get_global_model_param())
    # this model now contains the final trained model

if __name__ == "__main__":
    main()

References

1. H. Brendan McMahan and Eider Moore and Daniel Ramage and Seth Hampson and Blaise Aguera y Arcas. “Communication-Efficient Learning of Deep Networks from Decentralized Data” in Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS) 2017. http://arxiv.org/abs/1602.05629
2. GDPR- Europe
   LGPD - Brazil
   APPI - Japan
   PIPEDA - Canada
   And many more
3. AppleInsider Flurry
4. Andrew Hard and Chloé M Kiddon and Daniel Ramage and Francoise Beaufays and Hubert Eichner and Kanishka Rao and Rajiv Mathews and Sean Augenstein. “Federated Learning for Mobile Keyboard Prediction”, in arxiv 2018. https://arxiv.org/abs/1811.03604
5. Hsu, Tzu-Ming Harry and Qi, Hang and Brown, Matthew. “Measuring the effects of non-identical data distribution for federated visual classification”, in arXiv preprint arXiv:1909.06335 2019. https://arxiv.org/pdf/1909.06335 Li, Xiang and Huang, Kaixuan and Yang, Wenhao and Wang, Shusen and Zhang, Zhihua. “On the convergence of fedavg on non-iid data”, in arXiv preprint arXiv:1907.02189 2019. https://arxiv.org/pdf/1907.02189
6. K. Wei et al., “Federated Learning With Differential Privacy: Algorithms and Performance Analysis,” in IEEE Transactions on Information Forensics and Security, vol. 15, pp. 3454-3469, 2020, doi: 10.1109/TIFS.2020.2988575
7. Xianghua Xie, Chen Hu, Hanchi Ren, Jingjing Deng, A survey on vulnerability of federated learning: A learning algorithm perspective, Neurocomputing, Volume 573, 2024, 127225, ISSN 0925-2312, https://doi.org/10.1016/j.neucom.2023.127225.
8. Karimireddy, Sai Praneeth and Kale, Satyen and Mohri, Mehryar and Reddi, Sashank and Stich, Sebastian and Suresh, Ananda Theertha. “Scaffold: Stochastic controlled averaging for federated learning”, in International conference on machine learning 2020. https://proceedings.mlr.press/v119/…
9. Li, Tian and Sahu, Anit Kumar and Zaheer, Manzil and Sanjabi, Maziar and Talwalkar, Ameet and Smith, Virginia. “Federated optimization in heterogeneous networks”, in Proceedings of Machine learning and systems, Volume 2, 2020. https://proceedings.mlsys.org/paper/2020/file/…
10. Jung, Alexander. “Federated learning: From theory to practice”, in Springer Nature 2026. https://arxiv.org/pdf/2505.19183