# Stochastic Unrolled Federated Learning

arXiv: https://arxiv.org/abs/2305.15371 · HTML: https://mlsanity.com/paper/2305.15371/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 56 of its 88 references that are in the map.
- First posted on arXiv 24 May 2023.

## Builds on (most cited first, 30 of 56)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Deep Learning with Differential Privacy](https://mlsanity.com/paper/1607.00133/) (2016)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [Model-Contrastive Federated Learning](https://mlsanity.com/paper/2103.16257/) (2021)
- [Learning to learn by gradient descent by gradient descent](https://mlsanity.com/paper/1606.04474/) (2016)
- [Federated Learning Based on Dynamic Regularization](https://mlsanity.com/paper/2111.04263/) (2021)
- [Algorithm Unrolling: Interpretable, Efficient Deep Learning for Signal and Image Processing](https://mlsanity.com/paper/1912.10557/) (2019)
- [Meta-SGD: Learning to Learn Quickly for Few-Shot Learning](https://mlsanity.com/paper/1707.09835/) (2017)
- [A Unified Theory of Decentralized SGD with Changing Topology and Local Updates](https://mlsanity.com/paper/2003.10422/) (2020)
- [Meta-Dataset: A Dataset of Datasets for Learning to Learn from Few Examples](https://mlsanity.com/paper/1903.03096/) (2019)
- [Learning to Optimize: A Primer and A Benchmark](https://mlsanity.com/paper/2103.12828/) (2021)
- [Asynchronous Decentralized Parallel Stochastic Gradient Descent](https://mlsanity.com/paper/1710.06952/) (2017)
- [Stability Properties of Graph Neural Networks](https://mlsanity.com/paper/1905.04497/) (2019)
- [Graphon Neural Networks and the Transferability of Graph Neural Networks](https://mlsanity.com/paper/2006.03548/) (2020)
- [Asynchronous Parallel Stochastic Gradient for Nonconvex Optimization](https://mlsanity.com/paper/1506.08272/) (2015)
- [Decentralized Federated Averaging](https://mlsanity.com/paper/2104.11375/) (2021)
- [Learned Optimizers that Scale and Generalize](https://mlsanity.com/paper/1703.04813/) (2017)
- [Theoretical Linear Convergence of Unfolded ISTA and its Practical Weights and Thresholds](https://mlsanity.com/paper/1808.10038/) (2018)
- [Learning to Optimize Multigrid PDE Solvers](https://mlsanity.com/paper/1902.10248/) (2019)
- [Decentralized Collaborative Learning of Personalized Models over Networks](https://mlsanity.com/paper/1610.05202/) (2016)
- [Graphs, Convolutions, and Neural Networks: From Graph Filters to Graph Neural Networks](https://mlsanity.com/paper/2003.03777/) (2020)
- [Unfolding WMMSE using Graph Neural Networks for Efficient Power Allocation](https://mlsanity.com/paper/2009.10812/) (2020)
- [On the Linear Convergence of the ADMM in Decentralized Consensus Optimization](https://mlsanity.com/paper/1307.5561/) (2013)
- [Decentralized Federated Learning through Proxy Model Sharing](https://mlsanity.com/paper/2111.11343/) (2021)
- [Decentralized Federated Learning: Balancing Communication and Computing Costs](https://mlsanity.com/paper/2107.12048/) (2021)
- [Learning Gradient Descent: Better Generalization and Longer Horizons](https://mlsanity.com/paper/1703.03633/) (2017)
- [Local Differential Privacy for Deep Learning](https://mlsanity.com/paper/1908.02997/) (2019)

## Built on by (most cited first, 6 of 6)

- [Robust Stochastically-Descending Unrolled Networks](https://mlsanity.com/paper/2312.15788/) (2023)
- [Limited Communications Distributed Optimization via Deep Unfolded Distributed ADMM](https://mlsanity.com/paper/2309.14353/) (2023)
- [Unrolled Graph Neural Networks for Constrained Optimization](https://mlsanity.com/paper/2509.17156/) (2025)
- [Improving Federated Aggregation with Deep Unfolding Networks](https://mlsanity.com/paper/2306.17362/) (2023)
- [A Constrained Optimization Perspective of Unrolled Transformers](https://mlsanity.com/paper/2601.17257/) (2026)
- [Unrolled Neural Networks for Constrained Optimization](https://mlsanity.com/paper/2601.17274/) (2026)

## Abstract (arXiv)

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 federated learning in order to expedite its convergence. 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 guarantee its convergence. 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.