Check this blog entry for more information
Federated Learning
This project is an implementation of a Federated Learning system consisting on a Parameter Server and an Android application that can be used as a client
Both components are implemented in Kotlin using DL4J as the Machine Learning framework
Enjoy!
The problem
Training a machine learning model requires data. The more we have, the better (well... not always, but let's allow some simplifications). However, data is not cheap and more importantly, it can contain sensitive and personal information.
Recent developments in privacy in the form of new laws as GDPR and the increase of awareness of users and citizens in the value of their data is generating a need for techniques to enforce more privacy
Though techniques as anonymisation can greatly help with the privacy issue the fact that all the data is being sent to a central location to train the machine learning models is always a motive to be worried about
Federated Learning as a solution to privacy
Federated Learning turns the update of Machine Learning models upside-down by allowing the devices on the edge to participate in the training.
Instead of sending the data in the client to a centralised location, Federated Learning sends the model to the devices participating in the federation. The model is then re-trained (using Transfer Learning) with the local data
And the data, your data, never leaves the device, let that be your phone, your laptop or your IoT gadget
High-level Architecture
Very briefly the process of training a model goes as follows:
- The server opens a new round of training
- The clients that are going to participate in the training round download the latest version of the model from the server
- Using their local data, each client updates the model
- Those updates are sent to the server
- The server gathers all updates and applies Federated Averaging to improve the shared model
- The shared model is now ready for all clients to use
Use Case. Classifying and training images
To demonstrate how Federated Learning works, I have implemented a system based on Cifar-10, a well-known image classification dataset
Android client
| 
| 
|
Client Architecture
The architecture allows to remove the UI bit in Android and apply the rest with little effort to another type of client that supports Kotlin
Parameter Server
Server Architecture
Federated Averaging
The averaging is done by the server once it has received a minimum number of updates. It applies Federated Averaging as defined in 1
See FederatedAveragingStrategy.kt
override fun processUpdates(): ByteArrayOutputStream {
val totalSamples = repository.getTotalSamples()
val model = ModelSerializer.restoreMultiLayerNetwork(repository.retrieveModel())
val shape = model.getLayer(layerIndex).params().shape()
val sumUpdates = repository.listClientUpdates().fold(
Nd4j.zeros(shape[0], shape[1]),
{ sumUpdates, next -> processSingleUpdate(next, totalSamples, sumUpdates) }
)
model.getLayer(layerIndex).setParams(sumUpdates)
val outputStream = ByteArrayOutputStream()
ModelSerializer.writeModel(model, outputStream, true)
repository.storeModel(outputStream.toByteArray())
return outputStream
}