Auto-Unrolled Proximal Gradient Descent: An AutoML Approach to Interpretable Waveform Optimization

Ahmet Kaplan

cs.LGMar 18, 2026

Auto-Unrolled Proximal Gradient Descent: An AutoML Approach to Interpretable Waveform Optimization

Authors: Ahmet Kaplan

Abstract

This study explores the combination of automated machine learning (AutoML) with model-based deep unfolding (DU) for optimizing wireless beamforming and waveforms. We convert the iterative proximal gradient descent (PGD) algorithm into a deep neural network, wherein the parameters of each layer are learned instead of being predetermined. Additionally, we enhance the architecture by incorporating a hybrid layer that performs a learnable linear gradient transformation prior to the proximal projection. By utilizing AutoGluon with a tree-structured parzen estimator (TPE) for hyperparameter optimization (HPO) across an expanded search space, which includes network depth, step-size initialization, optimizer, learning rate scheduler, layer type, and post-gradient activation, the proposed auto-unrolled PGD (Auto-PGD) achieves 98.8% of the spectral efficiency of a traditional 200-iteration PGD solver using only five unrolled layers, while requiring only 100 training samples. We also address a gradient normalization issue to ensure consistent performance during training and evaluation, and we illustrate per-layer sum-rate logging as a tool for transparency. These contributions highlight a notable reduction in the amount of training data and inference cost required, while maintaining high interpretability compared to conventional black-box architectures.

Figures & tables

Fig. 1: Abstract workflow of Auto-Unrolled Proximal Gradient Descent: classical PGD is transformed into a deep-unfolded network, AutoGluon automates architecture search and hyperparameter optimization, and the system outputs optimized beamforming vectors for wireless waveform optimization.

Parameter	Symbol	Value
Number of BS antennas	$M$	8
Number of users	$K$	4
Maximum transmit power	$P_{\max}$	1
AWGN variance	$\sigma^{2}$	0.1
Channel model	—	Rayleigh, i.i.d. $\mathcal{CN}(0,\mathbf{I})$
Carrier frequency	—	28 GHz (3GPP UMi)

TABLE I: Simulation and Training Parameters

$N$	ZF	PGD (200-it)	WMMSE (100-it)	MLP	LISTA	PGD-Net	Auto-PGD
$10^{2}$	13.33	14.81	14.82	$1.58{\pm}0.03$	$2.26{\pm}0.37$	$12.70{\pm}0.89$	$\mathbf{14.52}{\pm}0.12$
$10^{3}$	13.33	14.81	14.82	$2.38{\pm}0.14$	$6.52{\pm}0.02$	$12.39{\pm}0.66$	$\mathbf{14.64}{\pm}0.00$
$10^{4}$	13.33	14.81	14.82	$4.47{\pm}0.04$	$6.15{\pm}0.01$	$11.80{\pm}0.09$	$\mathbf{13.58}{\pm}0.15$
$5\times 10^{4}$	13.33	14.81	14.82	$6.45{\pm}0.29$	$6.98{\pm}0.01$	$11.65{\pm}0.09$	$\mathbf{11.89}{\pm}0.06$

TABLE II: Test Sum-Rate (bits/s/Hz) vs. Training Set Size.

Fig. 2: Sum-rate vs. training set size for all methods (5-seed mean

\pm

std). Auto-PGD achieves

14.52{\pm}0.12

N=10^{2}

(

98.0\%

) and

14.64{\pm}0.00

N=10^{3}

(

98.8\%

) of Classical PGD. The performance drop at

N{\geq}10^{4}

coincides with the truncated HPO budget (20 and 5 trials, respectively). WMMSE (

14.82

bits/s/Hz) is omitted as it is visually indistinguishable from Classical PGD.

Feature	Traditional PGD	Black-Box DL	Proposed Framework
Model Type	Model-Based	Data-Driven	Gray-Box (Hybrid)
Parameters	Hand-tuned	$10^{5}-10^{7}$	$10^{1}-10^{2}$
Inference Latency	High (Iterative)	Low	Low
Data Efficiency	N/A	Low	High
Interpretability	Mathematical	Post-hoc (e.g., SHAP/LIME)	Algorithmic
Optimization	Manual	Stochastic	AutoML-Optimized

TABLE III: Comparison of Optimization Paradigms

Fig. 3: Training loss convergence for all learned methods at

N{=}100

and

N{=}1000

. At

N{=}1000

, Auto-PGD and LISTA both converge to the lowest training loss while MLP and PGD-Net plateau higher; PGD-Net’s high loss at both sizes is consistent with the high test variance in Table II . At

N{=}100

, LISTA quickly reaches low training loss yet achieves only

2.26{\pm}0.37

bits/s/Hz on the test set (Table II ), indicating overfitting; Auto-PGD’s oscillatory trajectory reflects its mixed-sign step-size schedule.

Method	HPO	Auto depth	Power feas.	Params
IAIDNN [ 4 ]	✗	✗	✗	$O(M^{2}KL)$
Deep W-MMSE [ 5 ]	✗	✗	✗	$O(M^{2}KL)$
DeepFP [ 8 ]	✗	✗	✗	$O(MKL)$
PGD-Net (fixed)	✗	✗	✓	$O(L)$
Auto-PGD (ours)	✓	✓	✓	$O(L)$

TABLE IV: Feature Comparison with State-of-the-Art DU Beamforming Methods

Deep Unfolding

Deep unfolding turns iterative algorithms into trainable networks, combining model-based structure with learned components for efficient inference and optimization.

See all 37 papers

Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization
Jun 29, 2026Haoming Meng, Anton Sugolov, Vardan PapyanDeep Learning Optimization Neural Network Optimization
Department of Mathematics, University of Toronto, Toronto, Canada · Vector Institute, Toronto, Canada
Alternating Gradient Flow Utility: A Unified Metric for Structural Pruning and Dynamic Routing in Deep Networks
Mar 12, 2026Tianhao Qian, Zhuoxuan Li, Jinde Cao +2Neural Network Pruning Structured Sparsity
the School of Mathematics, Southeast University, Nanjing 210096, China · the School of Mathematics, Southeast University, Nanjing 210096, China, and also with the Systems Research Institute of the Polish Academy of Sciences, Warsaw 01-447, Poland · the School of Cyber Science & Engineering, Southeast University, Nanjing 210096, China +1
When Losses Align: Gradient-Based Composite Loss Weighting for Efficient Pretraining
May 8, 2026Ivan Karpukhin, Andrey SavchenkoSelf-Supervised Pre-Training Adaptive Loss Weighting
Sber AI Lab · HSE University ISP RAS Research Center for Trusted Artificial Intelligence

Jun 29, 2026J/K move · Enter open · S save

Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization

Haoming Meng, Anton Sugolov, Vardan Papyan

Department of Mathematics, University of Toronto, Toronto, Canada · Vector Institute, Toronto, Canada

Deep Learning Optimization Neural Network Optimization

Open paper AskPDF arXiv

Deep neural networks with repeated architectural blocks, such as transformers, often exhibit structured relationships across layers that emerge during training. Motivated by this observation, we introduce Depth-wise Gradient Augmentation, a general optimization paradigm in which the update applied to each layer is obtained by transforming the collection of block-wise optimizer updates along the depth dimension. Within this framework, we study Gradient Smoothing, a family of depth-wise smoothing methods, and instantiate it with a simple local Window Smoothing operator. The resulting method operates directly on block-wise updates produced by arbitrary base optimizers (e.g., SGD, Adam, Muon), incurs minimal computational overhead, and is compatible with existing optimization pipelines. We evaluate Gradient Smoothing across a diverse set of architectures and training regimes, including language model pretraining, RL post-training of LLMs for reasoning, diffusion modeling, and image classification with Vision Transformers. Across these settings, Gradient Smoothing consistently improves optimization and generalization performance without modifying model architectures or training objectives. We further show that it promotes more structured representation evolution across depth, consistent with its interpretation as a structured depth-wise preconditioning method. Together, these results establish Depth-wise Gradient Augmentation as a promising framework for exploiting cross-depth structure in optimization and demonstrate Gradient Smoothing as a simple and broadly applicable instantiation.

Abstract

Deep Unfolding

Explore similar work

Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization

Alternating Gradient Flow Utility: A Unified Metric for Structural Pruning and Dynamic Routing in Deep Networks

When Losses Align: Gradient-Based Composite Loss Weighting for Efficient Pretraining