Algorithm unrolling has emerged as a learning-based optimization paradigm that unfolds truncated iterative algorithms in trainable neural-network optimizers. We introduce Stochastic UnRolled Federated learning (SURF), a method that expands algorithm unrolling to a federated learning scenario. Our proposed method tackles two challenges of this expansion, namely the need to feed whole datasets to the unrolled optimizers to find a descent direction and the decentralized nature of federated learning. We circumvent the former challenge by feeding stochastic mini-batches to each unrolled layer and imposing descent constraints to mitigate the randomness induced by using mini-batches. We address the latter challenge by unfolding the distributed gradient descent (DGD) algorithm in a graph neural network (GNN)-based unrolled architecture, which preserves the decentralized nature of training in federated learning. We theoretically prove that our proposed unrolled optimizer converges to a near-optimal region infinitely often. Through extensive numerical experiments, we also demonstrate the effectiveness of the proposed framework in collaborative training of image classifiers.
翻译:算法展开作为一种基于学习的优化范式,将截断的迭代算法转化为可训练的神经网络优化器。我们提出随机展开式联邦学习(SURF)方法,将算法展开扩展至联邦学习场景。该方法解决了此类扩展面临的两大挑战:一是需要将完整数据集输入展开优化器以寻找下降方向,二是联邦学习的去中心化特性。针对前者,我们通过向每个展开层输入随机小批量数据并施加下降约束来抑制其随机性;针对后者,我们设计基于图神经网络(GNN)的展开架构来解构分布式梯度下降(DGD)算法,从而保持联邦学习训练的去中心化本质。理论上证明了所提出的展开优化器能以无限频繁的次数收敛至近最优区域。通过大量数值实验,验证了该框架在协同训练图像分类器中的有效性。