This repository contains the experiments and possible extensions of the paper MARINA: Faster Non-Convex Distributed Learning with Compression by Eduard Gorbunov, Konstantin Burlachenko, Zhize Li, Peter Richtárik. This repo is initially forked from the original repo, mentioned in the paper.
This paper proposes a new algorithm, which performs communication efficient distributed learning. It utilizes the compression strategy based on compression of gradient differences. And this method deals with the bottleneck of the communication in distributed learning, which is the main problem of the distributed learning.
This method aims to optimize smooth, but non-convex finite-sum problems of the following form:
where
The algoithm implements the following steps:
Generalized version of the MARINA is VR-MARINA, which is based on the variance reduction technique. which can handle the situation when the local functions
VR version of the algorithm implements the following steps:
At this part of the project I reproduced the experiments from the paper. It's simple task of the image classification using the CIFAR-10 dataset, using the ResNet-18 model.
The main difficulty of this experiment is the distributed learning setup. I used 5 workers, each of them have a ResNet-18 with it's own set of parameters. The main idea of the distributed learning is that each worker has a part of the dataset. Loss function is calculated on each worker, and then the gradients are aggregated on the parameter server. After that, the parameter server sends the updated parameters to each worker.
As a method for gradient compression I used RandK, which is a random k-sparse compression method. It's a simple method, which randomly selects k elements from the gradient vector and sets the rest to zero. For the experiments I used
As we can see, the VR-MARINA algorithm converges faster than the MARINA algorithm.
The idea of the original paper is to have part of the dataset on each worker. And the part of the dataset here means that samples of the dataset are distributed between the workers.
And in my work I considered the case of the vertical partitioning of the dataset. It means that each worker has all the samples, but each sample has only a part of the features. And the main idea of the vertical partitioning is that each worker has a part of the features of each sample. This approach can be especially useful then we dealing with the multi-modal data, and each client server works with a different modality of the data, for which it have been pre-trained.
In this case I have simple regression task:
$$
Dx = b
$$
where
And each worker has a part of vector
Function we want to minimize is:
It's derivative (gradient) is:
For gradient we need
And we can't calculate it on each worker, because it will be too expensive. So we need to use the compression of the gredient differences as in the original paper.
Quantized gradient difference is:
And the update rule is:
As we can see, the algorithm converges even with the compression of the gradient differences.
In this part I used multi-modal data, which is a combination of the text and image data. As a dataset I used the CUB-200-2011 dataset link, which contains images of the birds and their descriptions. And the task is to predict the class of the bird based on the given information. Sample from the dataset looks like this:
And the model I used for the prediction is the combination of the ResNet-18 for the images and BERT for the text. Architecture of the model is the following:
First I trained the model with full gradient, to check if fine-tuning of BERT and ResNet-18 needed for this task. And the results are the following:
As we can see, the model with fine-tuning of BERT and ResNet-18 doesn't make much difference. So I used the model with the pre-trained BERT and ResNet-18 for the experiments.
In this setup we don't send the gradients from the server to workers, so we can't use gradient compression here. Instead I tried to use the compression on the forward pass. I compress vectors of embeddings from the ResNet-18 and BERT, and then send them to the server. With this I got the following results with different compression rates:
This method of comparison is not what we would want, as we compressing the data on the forward pass, but not the gradients. So, to implement the MARINA-like algorithm for the vertical partitioning of the dataset, we need to modify the model architecture. We need to have trainable parameters on the clients side, so we can send gradients from the server to them. So, the modified architecture looks like this:
It's identical to the previous one, but now we have trainable parameters on the clients side. And with this we can utilize the compression of the gradient differences.
I tried to train the model with different compression rates, and got the following results:
As we can see, the model converges even with the compression of the gradient differences. So, we can use this method for the distributed learning with the vertical partitioning of the dataset.
During this project I reproduced the experiments from the original paper, and implemented the MARINA-like algorithm for the vertically splitted dataset. And the results show that the algorithm converges even with the compression of the gradient differences. So, we can use this method for the distributed learning with the vertical partitioning of the dataset.