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.

Background

Recently, I’ve become fascinated with graph networks and understanding their inner workings. Following this interest, I’m attempting to create something similar to Graph Network Dojo, which will contain implementations of various graph layers and models. The purpose of this project is for me to deepen my understanding of GNNs.

Graph Neural Networks (GNNs) represent a novel type of neural network architecture that operates on a graph-in, graph-out principle. The input to the model is a graph, which is processed and transformed into an output graph. This output maintains the same structure as the input, but with modified node and edge embeddings. These transformed graphs can then be utilized for various downstream tasks such as link prediction, node classification, or edge classification.

While there are many variants, modern GNNs generally fall into three categories:

1. Message Passing: Nodes exchange information (messages) with neighbors iteratively.
2. Convolution: Generalizing the convolution operator from image processing to graph structures (e.g., GCNs).
3. Attention: Using mechanisms like Graph Attention Networks (GAT) to weigh the importance of different neighbors dynamically.

For better understanding, articles on Distill are highly recommended.

Introduction

In my quest to learn more about GNNs, I’ve been exploring various applications to implement. Forst implementation was to solve (or atleast try!) the HAMD, the Hamiltonian cycle decision problem. The results were pretty satisfacory atleast for graphs with N <= 30, for N > 30, creating a dataset was the difficult part. After HAMD i began to search for papers to implment, I encountered the concept of molecule generation. After looking at different approaches, I was decided to test out Xavier Bresson’s two-step decoder method and chose to implement it.

Throughout this process, I found Chaitanya Joshi’s blogs and projects to be incredibly helpful resources. However, despite this guidance, I still encountered numerous challenges while attempting to implement the model. This project has been be an excellent learning experience, allowing me deepen my understanding of both GNNs.

GVAEs: Model Architecture

The outline of the approach is to map the input graph into a latent space using an encoder, and then sample from the space to reconstruct the molecules using a decoder, this is the basic approach in VAEs.

1. The Encoder: A generic encoder’s task is to compress high dimensional input data into low dimensional representation called latent space (often represented by z). The variational part of VAE, comes from the fact that instead of generating a low dimensional feature vector z, we generate 2 quantitites mu and sigma, which represnet the mean and variance of the low dimensional distribution we want to mimic.

   The goal is to make the low dimensional distribution, a standard gaussian distribution, so the generated mu and sigma for an ideal encoder should be 0, 1. This is achieved through reparametrization. For this GVAE, the encoder is similar to a generic encoder, a GNN is used as encoder in GVAE.

2. 2 Step Decoder:

   In a generic decoder the low dimensional input is upsampled and the output follows the same structure as the input to the encoder, so in this case it should be a graph. But for graphs fixing the number of nodes is important, so the forst step of the decoder is node predictin, and the sencond step is edge prediction.

The Sparsity Problem and Edge Loss

A major challenge in molecular generation is the “Edge Loss” function. In a molecular graph, the number of actual bonds is very less compared to the number of possible bonds (a fully connected graph). If you train a naive model, it achieves high accuracy simply by predicting “No Bond” for every edge.

After testing reading about different approaches, weighted cross-entropy loss seemed appropriate. In this, weights are applied to each class based on the inverse of the frequency of that class in the dataset, so class that appeared less would have high weight, so a lower weight to no bond class and the model is forced to learn the structure of the graph rather than predict “no bond” all the time. Additionally, I would it challenging to handle proper indexing of nodes within the mini-batch for libraries like pytorch geometric, as in PyG, a minibatch is treated as a single graph o disconnected components.

Balancing the KL Divergence During model training, the KL divergence term dropped to low values very early, this made the model unstable and it didn’t improve that much afterwards.

To prevent this, I tried KL Annealing. In this a weight is given to the KL term that starts at 0 and gradually increases to 1 during training. This allows the model to first learn how to reconstruct the molecule (reducing reconstruction loss) before being trying to fit the latent space into a gaussian distribution.

Depth and Residual Connections

During experiments, I observed that deeper GNNs often struggle to propagate gradients effectively. Residual connections were very imporatant to reduce this effect. This allows gradients to flow backwards more efficiently, which stabilizes the training and results in better feature extraction.

Future Directions in Generative Graphs Implementing a GVAE for molecular generation highlights both the power and the limitations of current methods.

• Adjacency Matrix Complexity: Reconstructing a dense adjacency matrix is computationally expensive, For larger molecules, this becomes a problem.

• Output molecule strucure: The model might predict a carbon atom with 6 bonds, which is not possible. Although the paper suggests a method for this, I didn’t implement it.

• Looking forward, the field is moving toward autoregressive models (building the graph node-by-node) and reinforcement learning approaches. These methods can explicitly enforce chemical strcutre at each step of the generation process, potentially solving solving the strcture issues.

This is Kireeti’s personal website to know more About