Plankton-FL: Exploration of Federated Learning for Privacy-Preserving Training of Deep Neural Networks for Phytoplankton Classification

Federated learning has been shown to be very effective for training deep neural networks on decentralized data while ensuring data privacy [5]. Specifically, when there is sensitive data from various sources, federated learning can be leveraged. A typical federated learning framework consists of 2 components: a global model and $K$ clients. Each client contains its own local model, and they are trained iteratively and independently on their respective data. It is assumed that all of the data available is partitioned into $K$ clients, $P_{k}$. The local models are then used to update the global model [5]. This process is repeated for $N$ rounds in order for the global model to generalize. The objective for federated learning is calculated using equation 1.

In equation 1, $f_{i}(w)$ denotes the loss function $\ell(x_{i},y_{i};w)$ for an observation $(x_{i},y_{i})$ and model parameters $w$. Also, $n_{k}$ and $n$ denote the $|P_{k}|$ and the total number of observations, respectively.

The centralized federated learning algorithm which governs the communication between the global and local models is known as FederatedAveraging (FedAvg) and was introduced by McMahan et al. [5]. FedAvg is the iterative process of training all the local models, taking the average of all the updated weights from the local models, and then using it to update the global model. As described by McMahan et al., pseudo-code is provided in Algorithm 1.

In this paper, we used the FedAvg method, as described in Algorithm 1, when training our two instances of Plankton-FL. To test the feasibility and potential of federated learning, three different experiments were simulated. Specifically, a centralized learning baseline (CL), a (centralized) federated learning framework (FL), and a mutually exclusive, federated learning framework (ME-FL). Figure 1 provides a visual representation of each of the three experiments.

For the CL experiment, we have two data sources that we consolidate into a single server. From there we train a centralized model, which can then be deployed back to the edge devices for assessment and monitoring. The model was trained for a maximum of 75 epochs, with an early stopping criteria: A minimum of 50 epochs and a $\delta$ between test accuracies of 0.000001.

In the FL experiment, all of the training data was combined, randomly shuffled, and distributed to clients. Each client trained their own local model, on-prem, and only communicated their parameters back to the global server. FL was run for 10 iterations, where each iteration number corresponded to the number of clients. Namely, for the first iteration, there was only one client containing all of the training data, identical to CL, and with each increasing iteration another client was added (i.e. second iteration utilized two clients, etc.). Although, in reality, no single client will contain all of the data, for the purpose of comparing to CL, it made sense to start with a single client.

Unlike FL, in ME-FL, instead of combining all of the data together, shuffling, and distributing them, the clients only contained data from a single source, making them mutually exclusive. Again, each of the clients trained their own model, on-prem, and only communicated their parameters back to the global server. ME-FL was run for 9 iterations, starting from 2 clients up to 10. Iterations start from 2 clients due to the nature of the experiment; since each client only has data from a single source, it would not make sense to only have a single client. This modified setup ensures that we always have data from both sources.

The dataset was provided by Blue Lion Labs and was collected from two mutually exclusive sources, Halifax and Waterloo. It contained 301 distinct microscope specimen photos, each at a resolution of 3208 x 2200 pixels. The phytoplankton contained in these images are from eleven different species: Entomoneis paludosa, Alexandra catenella, Pymnesium parvum, Navicula sp, Heterosigma akashiwo, Prorocentrum lima, Alexandrium ostenfeldii, Porphyridium purpureum, Dolichospermum, Phaeodactylum tricornutum M1, and Phaeodactylum tricornutum M2.

For the purpose of this exploration, we took the majority class present in each image as the label to do image classification. This ensures that all models receive the same amount of information. Given the task, we built a custom convolutional neural network with four convolutional layers, three max-pooling layers, two dense layers, and an output layer. Across all intermediate layers, the ReLU activation function was used, and at the output, a softmax activation was used to predict the probabilities of each class. For all the convolutional layers, a kernel of size 3x3 and a stride of 1x1 was used and for all of the pooling layers, a pool and stride size of 2x2 was used. Additionally, dropout was used with a rate of 0.25 after the convolutional layers and a rate of 0.5 after the first dense layer.

When training, the images were resized to a resolution of 128 x 128 pixels and further augmented, using a horizontal flip, vertical flip, rotation, and color jitter, to create a larger data set of 2107 images. Across all experiments, the model architectures were held the same, a batch size of 8 was used, and the data was split into 80% training and 20% test. We also tune the learning rate across all experiments via a grid search over three different learning rates (LR) of 0.001, 0.0001, and 0.0005. Both of the federated learning experiments were run for 75 rounds and each local model was trained for 1 epoch. Furthermore, given it is a multi-label image classification task, the metric considered is prediction accuracy and the loss function is categorical cross-entropy.

Table 1 provides a numerical comparison across all experiments utilizing a learning rate of 0.0001. Note that, CL, FL, and ME-FL were trained for 51, 75, and 75 epochs, respectively. Firstly, when comparing CL and FL, we observe that for a single client FL outperforms CL. However, this is expected because FL with a single client is the exact same setup as CL; we expect the test accuracies to be very close in magnitude, and it is entirely possible that FL can outperform CL in this scenario. For all other number of clients, FL has a progressively worse test accuracy and is continuously outperformed by CL. In addition, across all number of clients, we observe that ME-FL consistently gets outperformed by CL and FL, which further demonstrates the impracticality of this method.

Figure 2 provides a visual comparison of CL, FL, and ME-FL across all epochs. We specifically look at the results for 10 clients of both FL and ME-FL, as in reality there are often large numbers of clients. From the figure, CL and FL both appear to learn, whereas ME-FL does not appear to learn at all. Comparing CL and FL, we observe that CL converges much faster than FL, which tells us that CL learns faster than FL. Overall, across all experiments, generally, we observed that CL performed the best, FL performed the second best, and ME-FL had the worst performance and this same trend is observed across each learning rate.

From table 1 and figure 2, we observe that FL performs relatively well, which prompted an investigation into its properties. Figure 3 displays the global model test accuracies for FL, for all numbers of clients, across the three learning rates. We observe a downward trend in the test accuracies as the number of clients increases. With an increasing number of clients, the global model needs to process and learn more information (i.e. the global model has to aggregate weights from more sources). With more information to process, learning is slowed down, yielding a worse generalization.

The largest contributor as to why ME-FL had a subpar performance relative to FL was because ME-FL was trained on individual, mutually exclusive clients. In our FL experiments we utilize a homogeneous model architecture, that is, all clients and the global server have the same model architecture. The nature of FL yields an independent and identically distributed (IID) distribution of labels across clients. However, in ME-FL, the distribution across clients is non-IID. As discussed in other literature, homogeneous federated learning performs poorly on non-IID data distributions [6]. Given this limitation, heterogeneous federated learning  [8, 9] is an alternative approach that should be explored. In this method, clients are allowed to differ in network architecture, allowing for more flexibility. Research has been done to explore applications of heterogeneous federated learning to mutually exclusive data [10] and it has typically been the preferred approach over homogeneous federated learning.