Skip to content

Neural Architecture Search

FedRLNAS

FedRLNAS is an algorithm designed to conduct Federated Neural Architecture Search without sending the entire supernet to the clients. Instead, clients still perform conventional model training as in Federated Averaging, and the server will search for the best model architecture. In this example, the server overrides aggregate_weights() to aggregate updates from subnets of different architectures into the supernet, and implements architecture parameter updates in weights_aggregated(). In its implementation, only only DARTS search space is supported.

cd examples/model_search/fedrlnas
uv run fedrlnas.py -c FedRLNAS_MNIST_DARTS.toml

Reference: Yao et al., "Federated Model Search via Reinforcement Learning," in Proc. International Conference on Distributed Computing Systems (ICDCS), 2021.

PerFedRLNAS

PerFedRLNAS is an algorithm designed to personalize different models on each client considering data and system heterogeneity, via Federated Neural Architecture Search. Different from FedRLNAS, where the server searches a uniform architecture for all clients. In this algorithm, each client will be given a different model structure and learn personalized architecture and model weights. In this example, the update rules and sample rules are redesigned to support this feature. In current implementation, examples of NASVIT MobileNetV3, and DARTS search space are provided.

NASVIT search space:

cd examples/model_search/pfedrlnas/VIT
uv run fednas.py -c ../configs/PerFedRLNAS_CIFAR10_NASVIT_NonIID01.toml

MobileNetV3 search space (synchronous mode):

cd examples/model_search/pfedrlnas/MobileNetV3
uv run fednas.py -c ../configs/PerFedRLNAS_CIFAR10_Mobilenet_NonIID03.toml

MobileNetV3 search space (asynchronous mode):

cd examples/model_search/pfedrlnas/MobileNetV3
uv run fednas.py -c ../configs/MobileNetV3_CIFAR10_03_async.toml

DARTS search space:

cd examples/model_search/pfedrlnas/DARTS
uv run fednas.py -c ../configs/PerFedRLNAS_CIFAR10_DARTS_NonIID_03.toml

Reference: Yao et al., "PerFedRLNAS: One-for-all Personalized Federated Neural Architecture Search," in Proc. 38th Annual AAAI Conference on Artificial Intelligence (AAAI), 2024.

FedTP

FedTP is proposed to improve personalized federated learning with transformer structured models. For each client, the attention maps in the transformer block are generated and updated by a hypernet working on the server, instead of being updated by average aggregation. The core part is in fedtp_server: customize_server_payload reloads the weights of attention maps with attention generated by hypernet before sending the models to clients and aggregate_weights updates the hypernet besides doing averaging aggregation of other parts of the model.

cd examples/model_search/fedtp
uv run fedtp.py -c FedTP_CIFAR10_ViT_NonIID03_scratch.toml

Reference: Li et al., "FedTP: Federated Learning by Transformer Personalization." Arxiv, 2022.

HeteroFL

HeteroFL is an algorithm aimed at solving heterogeneous computing resources requirements on different federated learning clients. They use five different complexities to compress the channel width of the model. In the implementation, we need to modify the model to implement those five complexities and scale modules. We provide examples of ResNet family and MobileNetV3 family here. The core operations of assigning different complexities to the clients and aggregate models of complexities are in function get_local_parameters and aggregation respectively, in heterofl_algorithm.py.

cd examples/model_search/heterofl
uv run heterofl.py -c heterofl_resnet18_dynamic.toml

Reference: Diao et al., "HeteroFL: Computation and Communication Efficient Federated Learning for Heterogeneous Clients," in Proc. International Conference on Learning Representations (ICLR), 2021.

FedRolex

FedRolex argues that the statistical method of pruning channels in HeteroFL will cause unbalanced updates of the model parameters. In this algorithm, they introduce a rolling mechanism to evenly update the parameters of each channel in the system-heterogeneous federated learning. In this implementation, models of ResNet and ViT are supported.

cd examples/model_search/fedrolex
uv run fedrolex.py -c example_ViT.toml

Reference: Alam et al., "FedRolex: Model-Heterogeneous Federated Learning with Rolling Sub-Model Extraction," in Proc. NeurIPS, 2022.

AnyCostFL

AnyCostFL is an on-demand system-heterogeneous federated learning method to assign models of different architectures to meet the resource budgets of devices in federated learning. In this algorithm, it adopts the similar policy to assign models of different channel pruning rates as the HeteroFL. But they prune the channel on the basis of the magnitude of the \(l_2\) norms of the channels. In this implementation, models of ResNet and ViT are supported.

cd examples/model_search/anycostfl
uv run anycostfl.py -c example_ResNet.toml

Reference: Li et al., "AnycostFL: Efficient On-Demand Federated Learning over Heterogeneous Edge Device," in Proc. INFOCOM, 2022.

SysHeteroFL

SysHeteroFL is a system-heterogeneous federated learning algorithm that assigns models of different architectures to the clients to achieve better performance when there are resource budgets on the clients. In this implementation, subnets of ResNet model with different architectures are sampled.

cd examples/model_search/sysheterofl
uv run sysheterofl.py -c config_ResNet152.toml

Reference: D. Yao, "Exploring System-Heterogeneous Federated Learning with Dynamic Model Selection," [https://arxiv.org/abs/2409.08858](https://arxiv.org/abs/2409.08858).
--