About Federated Learning

Update

I haven’t written in a long time, and in the past 2 years, a lot has happened. I have been part of a new startup and contributed sizely to its software, worked on a few research projects and papers, and am currently in an internship. As I have worked on various things over the past few years, I have learnt and have started to care about things I initially thought were unnecessary. These made me realize more about “what not to do” than “what to do”, a theme common to most of my experiences. In this post, I will be sharing my understanding of federated learning (the topic of my bachelor’s thesis), and try to motivate why it’s a good field.

Introduction to FL

Federated learning(FL), simply put, is a decentralized way of training ML models. It facilitates to train these models across several devices in a collaborative manner that allows each to share their information and benefit from others in the network. There are several implications of training model this way, the most important one is that each model can keep their data private and share information regarding the features rather than exact data. What this means is that each model (referred Interchangeably as node or client) can benefit from the collective information of the network without explicitly sharing its own data. This application is crucial, as in recent times there has been a lot of caution over LLMs and safety. This also reduces the need to collect and curate massive datasets in a centralized location, as clients collaborate, learn, and share, everyone collectively improves.

Formally, how do we define the sharing of information? In FL, the notion of ‘information’ is represented mainly in 2 ways: through model parameters and gradients. As these models learn and understand their data distribution, they form patterns in their internal representation that can be expressed through their parameters. The parameters understand the data and are used to process a new sample. In FL, these parameters represent the information related to the data and are shared with other clients. Different algorithms define different approaches to sharing these parameters, and they solve different challenges associated with FL. An important observation would be that if these models’ parameters represent the data, then would sharing them not violate the client’s data privacy? I will talk a bit more about this later on. Other than model parameters, sometimes the gradients are used; these represent the direction of learning for the client, and this information would be used by other clients. Again, similar to model parameters , different algorithms define different approaches to share this information. When clients receive information, they need to decide on an ‘aggregation’ technique, which is used to decide which client’s information is more important than others. As I have been saying, there are several algorithms that solve different problems, but I would broadly classify them into 3 categories: Global FL, Clustered FL, and Personalized FL (These are listed in chronological order). These are categorized on the basis of how the aggregation happens. In global FL approaches, there is a single global server and all clients share and receive their models from this server. The server defines the aggregation strategy, and in turn, this would mean that all clients share the same ‘global model’ which they each update incremetally. Clustered FL deviates from this line of thinking as it is possible that not all clients can benefit from each other. Clearly, clients would benefit from significantly more from collaborating with others that have similar data distrbution. Cluster FL divides clients into groups, and within each group, we would follow an approach similar to global FL. Personalized FL takes this further by not using any global or cluster models; instead, each client independently decides to aggregate or not aggregate based on the inforamtion it receives. I will explain a few important algorithms next.

Domain defining work (In my opinion)

The most seminal algorithm was defined in the first paper in this domain: FedAVG. It follows the global FL strategy where the server aggregates the client parameters weighted by the proportion of data samples each contains. What this means is that the “global model” would give more importance to the clients that have more data samples and inturn dominated by it. In a typical setting, this ensures that the global model follows the data distribution of the clients, and since each client will receive a copy of the global model each round, they would have inforamtion from others in the network. So what benefits does this have? An important thing to understand is that through FL, we intend to achieve different things; the obvious one is data privacy. Other than this, we also want the performance of these client models to be good, and they should benefit from the information provided by others; there would be no point in sharing at all. We hope that by gaining knowledge of other clients’ data distribution, a client would have a superior ability to extract and reason about its own data (Though there are several caveats that would reduce the client’s performance).

If in case the data distribution between clients is not homogeneous, then using a single global model and aggregating all information might be a bad idea. Instead, what if the clients can be grouped based on their data, and the information is only shared among the grouped clients? This is Clustered FL (CFL), and an important work in this field is the Iterative Federated Clustering Algorithm (IFCA). In IFCA, the server (yes, this still exists) which aggregates the client information, maintains multiple “cluster models” instead of a single global model. In this, each client decides its own cluster by comparing its data against all the cluster models. The cluster model that results in the least loss is the one most aligned to the client’s data. This approach addresses the issue of heterogeneity (where client data distributions are not similar) and shows considerable improvement in performance. Here, as we only aggregate within each cluster, it effectively reduces the problem of heterogeneity into homogeneity, which FedAVG solves well.

Even if this approach solves heterogeneity, an important problem still exists. FL systems are designed in an ‘online’ manner, since there is no need for centralized data collection, clients keep learning as new data arrives and so if the data distribution changes over time, then the cluster assignment for the client would also change; this is called data drift. The exact nature of drift can be in any way, it can add new labels to the dataset, or it can change the features. Several algorithms address this issue (mostly by constraining the global or cluster FL strategies with regularization terms), but I feel a newer direction is more interesting and useful. This is personalized FL (PFL), in this, there are no global or cluster models. Instead, each client will receive information from all other clients and selectively decide what to aggregate based on its own data. This behaves in a more “decentralized” manner and is not constrained as that of global FL or CFL (I personally think of this to be similar to message passing in graphs). PFL offers several benefits, as the aggregation logic can be distinct to each client, which would help with more selective gathering of information. FedAMP is an interesting work in this (it might not be the original PFL work, but it is intuitively one of the best works) where each client aggregates information based on similarity between its own model and other clients’ model.

The goal is always to improve the client’s performance to justify information sharing. In cases of heterogeneity or drift, this sharing can be “poisonous” and lead to reduced performance even below the complete local training(without any information sharing). In these cases, the clients could only benefit from information regarding the broad features of the data rather than specific classes/labels. FedRep and FedLayer are 2 algorithms that initially introduce these concepts where only specific parts of the model can be aggregated. By aggregating, say, only the backbone parameters of a feature extractor, the clients can benefit from understanding how different feature representations look rather than sharing the parameters related to classes/labels. This remains an important work that shows that always aggregating everything might not be a good idea.

This was not an exhaustive list of important algorithm, but these are what I feel are some of few interesting algorithms that have greatly steered the direction of the domain.

Challenges and Next Steps

As I have stated, if the model parameters represent the data distribution, would sharing them now violate privacy? It is important to understand what the model parameters represent and in what cases this might actually be a problem. In general, we believe that the parameters understand and store information regarding the data that they can later use for a new sample. In cases of generative models, the parameters mimic the data distributions, so sharing these could be a problem as one could reconstruct the original data (atleast to some extent) using speciailized techniques. In discriminative models, this information does not directly represent the data; instead, it is about differences in the data that allow the model to distinguish between different labels. Sharing of this information might be as problematic as the generative model case, but it is still concerning. There has been active research in this direction that introduces secure aggregation or differential privacy schemes, which noise/mask the parameters/gradients to reduce the chances of reconstructing the data (Though it has its own drawbacks, mainly reduced performance).

There are several things that I have not explicitly explained, things like concept, feature, and client drift, extent of heterogeneity, client participation and drawbacks of differential privacy. Most of these are still open problems, and active research is going on to help improve the overall performance. Exploring these problems would lead to a new FL paradigm and broaden the scope of the field. On a side note, programmatically, there are several interesting implementations of these algorithms showcasing intuitive implementation for distributed training, Flower is an excellent framework that internally uses Ray, they have many examples and usecases for FL.