Deep Unrolled WMMSE for Wireless Resource Allocation

• Deep unrolled WMMSE is a framework that reformulates iterative WMMSE updates into deep neural network layers, embedding domain knowledge for efficient optimization.
• It leverages graph neural network operations and learnable corrections to accelerate convergence and reduce sample complexity in wireless settings.
• Empirical results demonstrate near state-of-the-art sum-rate performance with significantly reduced inference times and robust scalability across diverse network topologies.

Deep unrolled WMMSE describes a class of model-based deep learning architectures that integrate the weighted minimum mean square error (WMMSE) algorithm within deep neural networks through algorithm unrolling or unfolding. By mapping classical WMMSE iterations to neural network layers, these frameworks embed domain knowledge—specifically, the algebraic structure and reasoning of WMMSE—into the learning architecture, reducing both sample complexity and computational delay. Such approaches have been applied to resource allocation in device-to-device (D2D) networks, power control for interference networks, and beamforming in multi-user multi-antenna systems, with state-of-the-art robustness and scalability across network topologies, sizes, and channel statistics (Yang et al., 2023, Chowdhury et al., 2020, Chowdhury et al., 2023).

1. Classical WMMSE Algorithm and its Reformulation

The original WMMSE algorithm is a block coordinate descent procedure for nonconvex resource allocation problems, notably weighted sum-rate (WSR) maximization in multiuser interference settings. For $K$ users and transmit powers $p_k$, WSR maximization takes the form

$$\max_{\{p_k\}} \sum_{k=1}^K \lambda_k \log_2\left(1 + \frac{|h_{kk}|^2 p_k}{\sum_{j \neq k} |h_{kj}|^2 p_j + \sigma^2}\right)$$

which is transformed into the equivalent sum-weighted-MSE minimization: $$\min_{\{u_k, v_k, w_k\}} \sum_{k=1}^K \lambda_k \left(w_k e_k - \log w_k\right), \quad 0 \leq v_k^2 \leq p_{\max}$$ with weights $w_k$, receiver coefficients $u_k$, and $v_k = \sqrt{p_k}$. The closed-form updates (per iteration) are:

• $$u_k^{(t)} = \frac{h_{kk} v_k^{(t-1)}}{\sum_j |h_{kj}|^2 (v_j^{(t-1)})^2 + \sigma^2}$$
• $$w_k^{(t)} = \frac{1}{e_k^{(t)}} = \frac{1}{1-u_k^{(t)} h_{kk} v_k^{(t-1)}}$$
• $$v_k^{(t)} = \Pi_{[0, \sqrt{p_{\max}}]} \left( \frac{\lambda_k u_k^{(t)} h_{kk} w_k^{(t)}}{\sum_j |h_{jk}|^2 (u_j^{(t)})^2 w_j^{(t)}} \right)$$

This approach, extended to MIMO and beamforming scenarios, is robust but incurs high runtime and is not well suited to dynamic or large-scale problems due to iterative convergence (Yang et al., 2023, Chowdhury et al., 2020).

2. Deep Unrolling of WMMSE: Graph-based and Layer-wise Architectures

Deep unrolled WMMSE replaces the iterative loop with a feedforward network, treating $T$ iterations as $T$ stacked layers. Two principal approaches have emerged:

• GNN Unrolling: Classical WMMSE steps are mapped to GNN operations. Nodes represent user pairs; edges encode channel coefficients. Aggregation and update functions utilize summation, ratio, and projection operations analogous to WMMSE updates. For instance, a MAX-aggregation and MLP approximate $\sum_j |h_{ij} v_j|^2$. By parameterizing these aggregations through small MLPs and sharing parameters across layers, the architecture retains WMMSE’s structure while enhancing data efficiency and generalization (Yang et al., 2023, Chowdhury et al., 2020).
• Learnable Correction/Parameterization: Some architectures inject trainable weights or corrections (realized by MLPs or GCNs) into specific update steps—most often the MSE weight update—while preserving the closed-form for others, e.g., for the receive and transmit filter steps. These weights are often generated by small GNNs operating over the interference graph (Chowdhury et al., 2023).

The network is trained end-to-end, usually with an unsupervised negative sum-rate loss over mini-batches of random channels, eliminating the need for explicit WMMSE or expert supervision.

3. Theoretical Underpinnings and Optimization Advantages

The main theoretical advantages of deep-unrolled WMMSE stem from:

• Decomposed Mapping: Unrolling decomposes the complex end-to-end mapping $G: v^{(t-1)} \mapsto v^{(t)}$ into simpler sub-mappings $g_1$, $g_2$, $g_3$ for $u$, $w$, $v$, learning each with lower effective complexity (Yang et al., 2023).
• Reduced Exploration Space: By restricting learning to WMMSE-inspired transformations, feature correlations rise while the effective hypothesis space contracts, improving generalization and accelerating convergence (Yang et al., 2023).
• Convexity and Monotonicity: For unfolded schemes with PGD or gradient-based substeps, under standard step-size bounds, every update is monotonic in the objective, and the limiting points are stationary for the original WSR problem (Pellaco et al., 2022).

These properties yield lower sample complexity and faster convergence relative to both black-box message-passing GNNs and purely supervised MLP baselines.

4. Empirical Performance, Complexity, and Robustness

Across diverse topologies and problem scales, deep unrolled WMMSE methods consistently achieve:

• Higher Rates with Fewer Layers: For $K=10,30,50$ users, unrolled GNN architectures (“UWGNN”—Editor's term) achieve $102$–$103\%$ of the converged WMMSE sum-rate with only 3 layers, using $1.18$G MACs and 1.9K parameters (Yang et al., 2023).
• Complexity Reduction: Compared to 100-iteration WMMSE ($1.8$G MACs, 2.5K params, $16$ ms inference), UWGNN reduces inference to $\approx$2 ms. Similar trends are observed in UWMMSE for power allocation (Chowdhury et al., 2020).
• Robust Generalization: Trained at one network size or topology, models generalize without retraining to other sizes (up to $K=100$), channel distributions (Rayleigh→Rician), spatial densities, and under topology and channel mobility (Yang et al., 2023, Chowdhury et al., 2020, Chowdhury et al., 2023).
• Sample and Data Efficiency: Fewer training samples are needed to achieve performance parity.
• Resilience to Distribution Shift: UWGNN maintains $>85$–$90\%$ of baseline performance under abrupt changes in channel statistics, topology, or user mobility (Yang et al., 2023).

5. Architectural and Practical Design Choices

Key design elements common to state-of-the-art unrolled WMMSE architectures include:

• Parameter Sharing: Aggregation and update functions (MLPs or GNN modules) are shared across layers, promoting efficient parameter use and facilitating transfer to varying network sizes (Yang et al., 2023, Chowdhury et al., 2020, Chowdhury et al., 2023).
• Graph Construction: Channel coefficients ($h_{ij}$) define the edge weights; node features encapsulate local power, weight, and filter state.
• Unsupervised End-to-End Loss: The loss is typically the negative expected sum-rate, evaluated using the current architecture’s outputs over random channel batches.
• Efficient Aggregation: Mechanisms (e.g., GNN MAX-pooling plus MLP) approximate sums and ratios central to WMMSE algebra.
• Scalability: Both per-layer and total parameter footprints are independent of network size, promoting scalability to large wireless networks (Chowdhury et al., 2023).

6. Open Challenges and Future Research Directions

Recognized open problems and future avenues include:

• Extension to MIMO Beamforming and Multi-cell Scenarios: While current D2D unrolled solutions are for scalar transmitters, extending to full MIMO, multi-antenna, and multi-cell systems (with corresponding matrix-valued updates) is ongoing (Chowdhury et al., 2023, Schynol et al., 2022).
• Online and Nonstationary Adaptation: Real-time, data-efficient adaptation under extreme channel nonstationarity and dynamic topologies remains a challenge (Yang et al., 2023).
• Integration with Scheduling and Joint Optimization: Incorporating edge scheduling, user admission, and cross-layer resource allocation within the unrolled framework merits further attention.
• Algorithm-aligned Network Design: The paradigm of algorithmically aligned, interpretable, and scalable architectures is poised to become foundational for wireless resource management under strict latency and complexity requirements (Yang et al., 2023).

7. Comparative Summary Table

Approach	Core Method	Key Gain Over Classic WMMSE	Robustness/Scaling
UWGNN (Yang et al., 2023)	GNN-unrolled WMMSE	$1.18$G MACs, 1.9K params, $2$ ms	Generalizes to $K=10$–$100$, channel & topology
UWMMSE (Chowdhury et al., 2020)	GNN-param. unfold.	$\sim2$ ms inference, few layers	Cross-size/density generalization
GCN-WMMSE (Chowdhury et al., 2023)	GCN-corrected WMMSE	$20$–$30\times$ faster inference	MIMO, arbitrary network size, topology shift

The approaches cited highlight that unrolling and graph neural message-passing, informed by WMMSE’s algebraic structure, yield architectures with rigorous performance guarantees, low complexity, transparent interpretability, and exceptional practical robustness. These characteristics position deep unrolled WMMSE as a preferred methodology for real-time, large-scale, and dynamically evolving wireless networks (Yang et al., 2023, Chowdhury et al., 2023, Chowdhury et al., 2020).