FedPM

[yc7z]

PR Type

[Feature]

Short Description

This PR implements the "Federated Probabilistic Mask Training" algorithm described in http://arxiv.org/pdf/2209.15328. An example in the `examples/fedpm_example' folder runs this algorithm on the MNIST dataset, and it reproduces part of the experimental results in the original paper.

On a high level, the implementation consists mainly of the following components:

1. masked_layers.py: this module uses the linear and convolutional layers provided by PyTorch as a basis to implement masked versions of those layers, which ensures that only probability scores receive gradients.
2. fedpm.py: this module handles the server-side aggregation. The Bayesian aggregation approach described in the paper is implemented here.
3. The parameter exchanger part: this part is responsible for selecting the probability scores and then use them to sample binary masks to be sent to the server.

Tests Added

Tests have been added to ensure that the newly added masked layers, parameter exchanger, strategy, and functions work correctly.

[emersodb]

PR and Design look really awesome. Fantastic job on all the tests! I learned some new stuff about PyTorch too. Overall, everything looks pretty close to ready. Just some fairly minimal comments and discussion points.

[jewelltaylor]

Thanks for contributing Yuchong ❤️

[emersodb]

I think the only minor question left before this is read to go is whether the lock file really needs to change (since no toml updates have happened). Outside of that, this is ready to go!

[emersodb]

Ready to go!