Learning to accelerate distributed ADMM using graph neural networks

Henri Doerks, Paul Häusner, Daniel Hernández Escobar, Jens Sjölund

Abstract:

Distributed optimization is fundamental to large-scale machine learning and control applications.
Among existing methods, the alternating direction method of multipliers (ADMM) has gained pop-
ularity due to its strong convergence guarantees and suitability for decentralized computation. How-
ever, ADMM can suffer from slow convergence and high sensitivity to hyperparameter choices.
In this work, we show that distributed ADMM iterations can be naturally expressed within the
message-passing framework of graph neural networks (GNNs). Building on this connection, we
propose learning adaptive step sizes and communication weights through a GNN that predicts these
hyperparameters based on the current iterates. By unrolling ADMM for a fixed number of itera-
tions, we train the network end-to-end to minimize the solution distance after these iterations for a
given problem class, while preserving the algorithm’s convergence properties. Numerical experi-
ments demonstrate that our learned variant consistently improves convergence speed and solution
quality compared to standard ADMM, both within the trained computational budget and beyond.