# Fast State-Augmented Learning for Wireless Resource Allocation with Dual Variable Regression

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 44 references that are in the map.
- First posted on arXiv 23 Jun 2025.

## Builds on (most cited first, 27 of 27)

- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Edge Artificial Intelligence for 6G: Vision, Enabling Technologies, and Applications](https://mlsanity.com/paper/2111.12444/) (2021)
- [Optimal Wireless Resource Allocation with Random Edge Graph Neural Networks](https://mlsanity.com/paper/1909.01865/) (2019)
- [Unfolding WMMSE using Graph Neural Networks for Efficient Power Allocation](https://mlsanity.com/paper/2009.10812/) (2020)
- [Resource Management in Wireless Networks via Multi-Agent Deep Reinforcement Learning](https://mlsanity.com/paper/2002.06215/) (2020)
- [Learning Optimal Resource Allocations in Wireless Systems](https://mlsanity.com/paper/1807.08088/) (2018)
- [State Augmented Constrained Reinforcement Learning: Overcoming the Limitations of Learning with Rewards](https://mlsanity.com/paper/2102.11941/) (2021)
- [Graph Embedding based Wireless Link Scheduling with Few Training Samples](https://mlsanity.com/paper/1906.02871/) (2019)
- [A Graph Neural Network Approach for Scalable Wireless Power Control](https://mlsanity.com/paper/1907.08487/) (2019)
- [Link Scheduling using Graph Neural Networks](https://mlsanity.com/paper/2109.05536/) (2021)
- [Learning Resilient Radio Resource Management Policies with Graph Neural Networks](https://mlsanity.com/paper/2203.11012/) (2022)
- [Intelligent O-RAN for Beyond 5G and 6G Wireless Networks](https://mlsanity.com/paper/2005.08374/) (2020)
- [State-Augmented Learnable Algorithms for Resource Management in Wireless Networks](https://mlsanity.com/paper/2207.02242/) (2022)
- [Modular Meta-Learning for Power Control via Random Edge Graph Neural Networks](https://mlsanity.com/paper/2108.13178/) (2021)
- [Unsupervised Learning for Asynchronous Resource Allocation in Ad-hoc Wireless Networks](https://mlsanity.com/paper/2011.02644/) (2020)
- [Near-Optimal Solutions of Constrained Learning Problems](https://mlsanity.com/paper/2403.11844/) (2024)
- [Diffusion Model Based Resource Allocation Strategy in Ultra-Reliable Wireless Networked Control Systems](https://mlsanity.com/paper/2407.15784/) (2024)
- [Power Allocation for Wireless Federated Learning using Graph Neural Networks](https://mlsanity.com/paper/2111.07480/) (2021)
- [Generative Diffusion Models for Resource Allocation in Wireless Networks](https://mlsanity.com/paper/2504.20277/) (2025)
- [Inexact Stochastic Mirror Descent for two-stage nonlinear stochastic programs](https://mlsanity.com/paper/1805.11732/) (2018)
- [Accelerated Distributed Allocation](https://mlsanity.com/paper/2401.15598/) (2024)
- [Survey of Distributed Algorithms for Resource Allocation over Multi-Agent Systems](https://mlsanity.com/paper/2401.15607/) (2024)
- [Opportunistic Routing in Wireless Communications via Learnable State-Augmented Policies](https://mlsanity.com/paper/2503.03736/) (2025)
- [Learning to Slice Wi-Fi Networks: A State-Augmented Primal-Dual Approach](https://mlsanity.com/paper/2405.05748/) (2024)
- [Joint Task Offloading and Routing in Wireless Multi-hop Networks Using Biased Backpressure Algorithm](https://mlsanity.com/paper/2412.15385/) (2024)
- [Momentum-based Distributed Resource Scheduling Optimization Subject to Sector-Bound Nonlinearity and Latency](https://mlsanity.com/paper/2503.06167/) (2025)
- [Wireless Link Scheduling with State-Augmented Graph Neural Networks](https://mlsanity.com/paper/2505.07598/) (2025)

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

- [Graph Signal Diffusion Models for Wireless Resource Allocation](https://mlsanity.com/paper/2604.05175/) (2026)
- [Constrained Diffusion Models with Primal-Dual Inference](https://mlsanity.com/paper/2606.17192/) (2026)
- [Generative Diffusion Models of Stochastic Graph Signals](https://mlsanity.com/paper/2607.06833/) (2026)

## Abstract (arXiv)

We consider resource allocation problems in multi-user wireless networks, where the goal is to optimize a network-wide utility function subject to constraints on the ergodic average performance of users. We demonstrate how a state-augmented graph neural network (GNN) parametrization for the resource allocation policy circumvents the drawbacks of the ubiquitous dual subgradient methods by representing the network configurations (or states) as graphs and viewing dual variables as dynamic inputs to the model, treated as graph signals supported over the graphs. Lagrangian maximizing state-augmented policies are learned during the offline training phase, and the dual variables evolve through gradient updates while executing the learned state-augmented policies during the inference phase. Our main contributions are to illustrate how near-optimal initialization of dual multipliers for faster inference can be accomplished with dual variable regression, leveraging a secondary GNN parametrization, and how maximization of the Lagrangian over the multipliers sampled from the dual descent dynamics substantially improves the training of state-augmented models. We demonstrate the superior performance of the proposed algorithm with extensive numerical experiments in a case study of transmit power control. Finally, we prove a convergence result and an exponential probability bound on the excursions of the dual function (iterate) optimality gaps.