Training models under resource constraints for cross-device federated learning

A computer-implemented method of training a global student model, comprising:

• receiving, on a client device comprising a data set, from a server, a local student model based on the global student model, the local student model comprising an embedding layer and a first layer;
• inputting, on the client device, the data set to the embedding layer to generate an embedding layer output;
• transmitting, from the client device and to the server, the embedding layer output;
• performing, on the client device, a forward pass on the first layer, with the embedding layer output as an input, to generate a student model first layer output;
• receiving, on the client device and from the server, a teacher model first layer output;
• calculating, on the client device, a loss based on the student model first layer output and the teacher model first layer output, wherein the student model first layer output and the teacher model first layer output are generated using the embedding layer output;
• training, on the client device, the first layer of the local student model until the student model first layer output converges with the teacher model first layer output;
• and
• transmitting, from the client device and to the server, first layer weights of the first layer of the local student model.

A system of training a global student model stored on a server, the server comprising a processing device and a memory comprising instructions that are executed by the processing device to perform a method comprising:

• storing, on the server, a global student model comprising a first layer and a teacher model comprising a first layer;
• transmitting, from the server and to one or more client devices, local student models based on the global student model, each client device comprising a corresponding data set, the local student models each comprising an embedding layer and a first layer;
• inputting, on each client device of the one or more client devices, the corresponding data set to the embedding layers of the local student models to generate a plurality of embedding layer outputs;
• receiving, at the server, an embedding layer output of the plurality of embedding layer outputs of one of the local student models;
• performing, on the server, a forward pass on the first layer of the teacher model, with the embedding layer output as an input, to generate a teacher model first layer output;
• transmitting, from the server and to the one or more client devices, the teacher model first layer output;
• updating, on each client of the one or more client devices, first layer weights of the local student models based on a loss between the teacher model first layer output and a student model first layer output, each generated using the embedding layer output;
• receiving, at the server, the first layer weights of the local student models;
• and
• calculating, on the server, first layer weights of the global student model using the received first layer weights of the local student models.