Learning to accelerate distributed ADMM using graph neural networks

2025-09-05

AI Summary

Distributed ADMM suffers from slow convergence and high sensitivity to hyperparameters, limiting its practicality in large-scale decentralized settings. To address this, we propose GNN-ADMM—a framework that unrolls ADMM iterations within a graph neural network (GNN) message-passing architecture, enabling differentiable modeling. The GNN dynamically predicts iteration-specific step sizes and inter-node communication weights, allowing end-to-end optimization of convergence behavior for targeted problem classes. We provide theoretical guarantees that the learned adaptive strategy preserves the original ADMM’s convergence properties. Empirical evaluation across diverse distributed optimization tasks—including consensus optimization, distributed logistic regression, and sparse signal recovery—demonstrates substantial acceleration in convergence, improved solution accuracy, and strong generalization to unseen topologies and problem instances. The implementation is publicly available.

Technology Category

• Search and Optimization: Distributed Search
• Machine Learning: Distributed Machine Learning & Federated Learning
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Distributed optimization is fundamental in large-scale machine learning and control applications. Among existing methods, the Alternating Direction Method of Multipliers (ADMM) has gained popularity due to its strong convergence guarantees and suitability for decentralized computation. However, ADMM often suffers from slow convergence and sensitivity to hyperparameter choices. In this work, we show that distributed ADMM iterations can be naturally represented within the message-passing framework of graph neural networks (GNNs). Building on this connection, we propose to learn adaptive step sizes and communication weights by a graph neural network that predicts the hyperparameters based on the iterates. By unrolling ADMM for a fixed number of iterations, we train the network parameters end-to-end to minimize the final iterates error for a given problem class, while preserving the algorithm's convergence properties. Numerical experiments demonstrate that our learned variant consistently improves convergence speed and solution quality compared to standard ADMM. The code is available at https://github.com/paulhausner/learning-distributed-admm.

Problem

• Accelerating distributed ADMM convergence using GNNs
• Learning adaptive hyperparameters to improve ADMM performance
• Enhancing solution quality and speed via learned step sizes

Innovation

• Uses graph neural networks for ADMM
• Learns adaptive hyperparameters via GNN
• Unrolls ADMM iterations for end-to-end training