Stochastic Unrolled Neural Networks

Samar Hadou, Navid NaderiAlizadeh, Alejandro Ribeiro

Keywords: unrolled optimization, learning to learn, deep unfolding, interpretable deep architecture, constrained learning

TL;DR: This paper develops stochastic unrolled neural networks as learned optimizers for empirical risk minimization (ERM) problems.

Abstract: This paper develops stochastic unrolled neural networks as learned optimizers for empirical risk minimization (ERM) problems. We view a fixed-depth unrolled architecture as a parameterized optimizer whose layers define a trajectory from an initial random model to a task-specific solution. To handle full datasets, we let each layer interact with randomly drawn mini-batches from the downstream dataset, so that the optimizer incrementally absorbs the entire task. We then train the unrolled optimizer under descent constraints that encourage reductions in loss gradient norms along this trajectory, shaping its dynamics to mimic a convergent stochastic descent method. We prove that such stochastic unrolled networks converge to near-stationary downstream models and quantify performance changes under shifts in the task distribution. As a case study, we instantiate this framework in federated learning by designing an unrolled graph neural network (GNN) architecture derived from decentralized gradient descent, and show that it maintains strong performance under data heterogeneity and asynchronous communication on collaborative image classification tasks.

Submission Number: 89