Federated Learning for TinyML Devices

Microcontrollers can now run machine learning models, but those models are almost always trained elsewhere, on powerful computers, and only flashed to the device afterwards. During my Bachelor’s thesis at UPC I explored the opposite: training the model directly on the microcontroller, and then letting several boards learn together with federated learning (FL) while their raw data never leaves the device.

The work led to two peer-reviewed publications and an open-source prototype.

Papers:
On-Device Training of Machine Learning Models on Microcontrollers with Federated Learning, Electronics 11(4): 573, 2022 [Open Access]
On-Device Training of Machine Learning Models on Microcontrollers With a Look at Federated Learning, ACM GoodIT ’21, pp. 198–203 [DOI]
TinyML: From Basic to Advanced Applications, Bachelor’s thesis, UPC, 2021 [Open Access]

Motivation

Typical TinyML workflows train models offline, optimise them (quantisation, pruning) and deploy a static model that only does inference. That means a device can’t keep learning from the data it sees in the field. We asked: can training run on the device itself, and can several devices improve a shared model without sending their recordings anywhere?

Prototype

• Task: keyword spotting. Press a button, say a word, and the board records one second of audio.
• Features: the board computes MFCCs (13 coefficients × 50 frames = 650 values) from the recording.
• Model: a small feed-forward neural network with one hidden layer (25 nodes by default) and 3 outputs, one per keyword. Stored as floats, the model takes about 65 kB, which fits in the board’s 256 kB of RAM so it can be updated during training.
• Training: backpropagation runs on the board itself, with a learning rate of 0.6 and momentum.
• Federated learning: three boards act as FL clients and talk to a Python server over USB serial. The server averages the boards’ weights (FedAvg) and sends the new global model back.

What we found

• On-device training works. The training loss decreased steadily as the boards trained on their own recordings.
• FL frequency is a trade-off. Running an FL round every 2 epochs reduced the loss much earlier than every 30 epochs, but each round costs bandwidth and time: sending the model took about 1.8 s from board to server and more than three times longer in the other direction, because of the board’s small USB buffer.
• Model size matters. We compared four hidden-layer sizes. Bigger layers reduced the loss sooner, while smaller ones needed more training samples, a direct trade-off against the microcontroller’s limited memory.
• Non-IID data. When the boards’ training data was not identically distributed, very frequent FL rounds (down to every epoch) helped reduce the loss of the local models.

The ACM GoodIT ’21 paper presented the first version of this work: on-device training for keyword spotting and a first look at federated learning between two boards. The Electronics paper extended it to three boards and a full analysis of FL rounds, model size, non-IID data and resource usage.

Code

Repository: https://github.com/marcmonfort/TinyML-FederatedLearning
The code from my Bachelor’s thesis:

• Arduino firmware (PlatformIO) for the Nano 33 BLE Sense clients: audio recording, MFCC features and on-device training of the neural network
• fl_server.py, the Python federated learning server that talks to the boards over serial
• Instructions to wire the buttons and run the demo: three buttons train three keywords, a fourth tests the model, and the boards merge their models every 10 seconds

Closing Thoughts

These experiments showed that even small microcontrollers can learn in the field while keeping their raw data local. The natural next steps are the ones we pointed to in the papers:

1. Studying non-IID data with more devices and less local data per device
2. Replacing USB serial with real wireless links (e.g. LoRa) and measuring their bandwidth and energy costs