Joint Beamforming and Integer User Association using a GNN with Gumbel-Softmax Reparameterizations

Mai Vu; Qing Lyu

REVIEW 3 major objections 5 minor 11 references

Joint Beamforming and Integer User Association using a GNN with Gumbel-Softmax Reparameterizations

T0 review · 3 major / 5 minor · reviewed 2026-08-07 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read This paper proposes a graph neural network that jointly optimizes beamforming and user association, using Gumbel-Softmax reparameterization to produce integer association decisions and reporting higher sum-rates than fractional…

desk verdict A useful but not-yet-trustworthy empirical claim: the STGS integer-association GNN is a real contribution, but Eq. (9)'s power projection treats fractional associations unfairly, likely inflating the reported gains. read the letter →

arxiv 2506.05241 v1 pith:GROT3AVW submitted 2025-06-05 eess.SP

Qing Lyu , Mai Vu This is my paper

classification eess.SP

keywords graphneuralnetworkbeamforminguserassociationGumbel-Softmaxstraight-throughestimatorsum-ratemaximizationmmWavenetworksunsupervisedlearning

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper is trying to establish that a single GNN, trained unsupervised on random channel samples, can jointly design beamforming vectors and user association decisions for a multicell mmWave downlink while directly outputting integer associations, each user assigned to exactly one base station. Existing machine-learning designs relax the association variable to fractional probabilities and then round, which the paper argues introduces bias and suboptimality. The proposed architecture appends a Gumbel-Softmax reparameterization, with or without Straight-Through estimation, to the association head so the output is one-hot in the forward pass yet still carries gradients in training. Simulation results with two base stations and up to 32 users show higher sum-rate than softmax-GNN, maximum-ratio transmission with max-SINR association, and the fractional association baseline, especially when the trained model is applied to larger networks and different transmit powers without retraining. If correct, this provides a real-time, scalable way to approach an NP-hard integer resource-allocation problem in dense wireless networks.

What carries the argument

The load-bearing mechanism is the Straight-Through Gumbel-Softmax reparameterization in Eq. (13): $d_{\mathrm{STGS}} = d_{\mathrm{1hot}} - [d_{\mathrm{GS}}]_{\mathrm{const}} + d_{\mathrm{GS}}$. In the forward direction the constant-shift term cancels the continuous sample, leaving the one-hot vector $d_{\mathrm{1hot}}$, so the association output is integer and satisfies the unique-association constraint; in the backward direction the gradient flows only through the differentiable Gumbel-Softmax sample $d_{\mathrm{GS}}$, avoiding the zero gradient of the discrete output. This is mounted on an edge-update GNN where UE node representations are converted by Eqs. (14) and (12) into association probabilities over base stations, edge representations are projected by Eq. (9) into beamforming vectors under each base station's power budget, and training minimizes the negative expected sum-rate in Eq. (10). The temperature parameter $\tau = 1$ controls how close the continuous relaxation is to one-hot.

What would settle it

Run the trained 2-BS/8-UE STGS model on several thousand fresh 32-UE channel realizations and compare its mean sum-rate with the fractional association baseline [2] and MRT maxSINR; if the GNN's sum-rate is not consistently above both, or if any association row fails to sum to one, the central claim is contradicted.

Watch

Extended reading notes

Core claim

On its own terms, the paper claims that the joint beamforming and user association problem in (4) can be solved by an edge-updating bipartite GNN whose association factors are produced by Gumbel-Softmax reparameterization. With Straight-Through estimation, the forward computation emits a strict one-hot vector satisfying constraint (4a), while the backward pass treats the continuous Gumbel-Softmax sample as the differentiable proxy; this gives integer decisions at no extra inference cost. The paper also claims that this reparameterized GNN, trained only to maximize expected sum-rate, generalizes to networks with more users than seen in training and to varied per-base-station powers, and that it achieves a higher sum-rate than all fractional association methods tested, with the pure GS variant matching or slightly exceeding STGS when the number of users grows large.

Load-bearing premise

The load-bearing premise is that the shortcut gradient used in training, treating the smooth Gumbel-Softmax value as a stand-in for the gradient of the one-hot choice, is accurate enough to guide the network to good associations, because if that proxy points the wrong way in some operating regime, the claimed sum-rate advantage over fractional association methods could fail.

Editorial extensions

If this is right

A trained model can be applied without retraining to networks with more users than in training, preserving integer associations and reporting higher sum-rate than fractional baselines.
Integer association decisions are produced in a single feedforward pass, avoiding the extra rounding step and the bias of relaxing integer constraints.
Training requires no labelled optimal association, only random channel samples, base-station powers, and noise powers, because the loss is the negative average sum-rate.
Pure Gumbel-Softmax without Straight-Through estimation shows a slight advantage over STGS at the largest tested user count (32 UEs), suggesting that fractional outputs during training help exploration and scalability.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

Beyond the paper, the same constant-shift trick could be reused for other discrete one-hot wireless allocation variables, such as beam selection, subcarrier assignment, or RIS element activation, giving integer outputs without changing inference cost.
Beyond the paper, the crossover where fractional GS beats discrete STGS at 32 UEs suggests a trade-off between exploration and exploitation that a bias analysis of the straight-through gradient could quantify; one testable prediction is that annealing the temperature over training would move the crossover.
Beyond the paper, the generalization claims assume the base-station count stays fixed because the association matrix dimension is tied to it; a set-based association head whose output size adapts to the number of base stations would be a natural extension for fully variable network sizes.

Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

Short version: this is a solid engineering paper with a genuine novelty—using Gumbel-Softmax with straight-through estimation inside a GNN to output provably integer user associations—and the integer mechanism itself is sound. But the central performance claim is weakened by a normalization issue in Eq. (9) that I checked carefully. For fractional association values a_{k,m} in (0,1), the projection transmits strictly less than P_m, because the denominator uses sum a ||u||^2 while the numerator uses a^2 ||u||^2 in the summed power. That means any fractional-association baseline, including the [2] baseline in Fig. 4, is evaluated below its nominal power budget. The proposed GS/STGS methods, by contrast, evaluate sum-rate with one-hot associations, so they satisfy (4b) exactly. This asymmetry mechanically inflates the claimed margin. It is load-bearing for the abstract's "higher sum-rate than all other fractional association methods." The fix is straightforward—normalize with sqrt(sum a^2 ||u||^2)—or state explicitly that comparisons are restricted to integer associations.

What the paper does well: the STGS forward pass in Eq. (13) is correct and guarantees integer outputs satisfying (4a); the edge-update GNN is clearly described; and the generalization tests across UE counts and power levels are a real step beyond fixed-size ML designs. The training time table and the complexity analysis, despite a small typo (T^2 vs T), give useful practical information. The authors also honestly discuss the trade-off between GS exploration and STGS exploitation in Section IV-D.

Soft spots beyond the power issue: no error bars, confidence intervals, or significance tests across the 3000 test samples; the baseline set omits WMMSE with integer rounding, which would provide a stronger upper bound; and no code or data is released, making replication harder. The STGS gradient bias is a theoretical gap, but it's secondary—the empirical comparison is the real problem.

Who is this for? Wireless communication researchers working on ML-based resource allocation, especially those interested in discrete-output GNNs. It deserves peer review—the core idea is publishable and the architecture is reusable—but acceptance should be conditional on fixing the power normalization and either adding error bars or tempering the dominance claims. My verdict on this version: reject for now, but with a clear path to acceptance.

Recommendation: send to review, but tell the authors the Eq. (9) issue is fatal to the current empirical claims and must be addressed before the paper can be accepted.

Referee Report

3 major / 5 minor

Summary. The paper addresses joint beamforming and user association in a multi-cell mmWave downlink. It proposes an edge-update GNN in which UE node representations are converted to association factors via Gumbel-Softmax (GS) or Straight-Through Gumbel-Softmax (STGS), and edge representations are projected to beamformers satisfying per-BS power constraints. Training is unsupervised, minimizing the negative sum-rate. The claimed contributions are (i) an architectural guarantee of integer association outputs satisfying constraint (4a), and (ii) simulation evidence that the GS/STGS GNN achieves higher sum-rate than fractional association baselines, including when generalized to larger networks and different transmit powers.

Significance. The architectural mechanism is clean and correct: STGS in Eq. (13) provides one-hot outputs in the forward pass while preserving a gradient path, and this is a useful, low-complexity way to enforce integer association in a GNN. The paper is also honest about the exploration/exploitation trade-off between GS and STGS. If the equal-power comparison and statistical support are fixed, the result would be a solid engineering contribution. The current empirical section, however, does not yet establish the dominance claims as stated.

major comments (3)

[Section III-B.3, Eq. (9)] The projection in Eq. (9) only enforces the per-BS power constraint (4b) when all association coefficients a_{k,m} are 0 or 1. For fractional associations, summing squared norms gives P_m times (sum_k a_{k,m}^2 ||u_{m,k}||^2) / (sum_k a_{k,m} ||u_{m,k}||^2), which is strictly smaller than P_m whenever any a_{k,m} lies in (0,1). Since the GS training loss in Section IV-C.1 and the fractional baseline of [2] use fractional associations, those configurations transmit strictly below their power budget. The comparison against fractional baselines is therefore not made at equal transmit power, which mechanically lowers the fractional sum-rate and inflates the reported margin. Please normalize with sqrt(sum_k a_{k,m}^2 ||u_{m,k}||^2) in the denominator, or state explicitly that Eq. (9) is applied only to one-hot associations and describe how fractional baselines are power-normalized.
[Section V-A and V-B] All empirical claims rest on single-curve averages over 3000 test samples; no error bars, confidence intervals, or multiple-seed variances are reported for Figs. 2, 4, 5, and 6. The claim that STGS or GS 'consistently' achieves higher sum-rate requires an estimate of variability. In particular, the crossing at 32 UEs in Fig. 4 could be within noise. Please report mean plus/minus standard deviation over at least several training seeds and independent test channel batches, and ideally a paired significance test for the GS/STGS variants against each baseline.
[Section V-B.3 and Section VI] The dominance claim that the proposed method achieves 'higher sum-rate when generalized to larger networks and various transmit power levels than all other fractional association solutions' is supported only by two conventional baselines (MRT maxSINR and the fractional method of [2]) plus softmax/GS variants of the same GNN. There is no WMMSE-with-integer-association baseline, no exhaustive search for small networks, and no rounding-based fractional-to-integer baseline. Since WMMSE [1] is cited as the standard beamforming approach, its absence from the comparison makes the stated dominance claim disproportionate to the experiments. Adding at least one strong integer-association baseline would materially strengthen the paper.

minor comments (5)

[Section III-B.3] The text 'the resulting UE representation matrix B(L)' appears to be a typo: the UE representations are denoted C, not B, throughout the architecture description.
[Section IV-E] The complexity expression 'L×2M KT2' is ambiguous; if T denotes the processing time of one MLP, the term should be written as 2 M K T^2 or 2 M K T depending on the intended unit, and it should be defined clearly.
[Section V-B.3] The sentence 'The results can be divided into 3 groups: GS with and without ST, softmax with and without ST, and traditional methods' is confusing because 'GS with ST' is the STGS method; please rephrase using the names introduced earlier.
[Section V-B.3] There is a typo in 'can limit the generation ability'; this should read 'generalization ability'.
[References] Reference [7] spells 'Gumbel-Softmax' as 'gumble-softmax'; please correct it.

Circularity Check

0 steps flagged · score 1.0 of 10

No significant circularity: the objective in Eq. (10) is the negative sum-rate used for both training and evaluation, which is normal for an optimization method, and no load-bearing claim reduces to a fitted parameter or to a self-citation.

full rationale

The derivation chain is self-contained. The GNN is trained with the negative sum-rate loss in Eq. (10) and evaluated on the same sum-rate; this is an engineering optimization loop, not a prediction derived from its own input. No per-sample constant is fitted and then reported as a prediction; hyperparameters such as temperature tau=1, L=2 layers, and hidden size 1024 are explicitly described as tuning choices in Section V-A. The Gumbel-Softmax and Straight-Through estimators in Eqs. (11)-(13) follow standard external references [6] and [7] and are not used to assume the claimed integer-association advantage; that advantage is established by simulation against external baselines ([2] and MRT maxSINR). The only self-citation, [9] in Section III-B, is a related-work pointer for edge-update GNNs and is not load-bearing. One non-circular benchmarking concern should be noted: the projection in Eq. (9) satisfies constraint (4b) exactly only for one-hot associations, since for fractional a in (0,1) the transmitted power is below P_m; this could affect fair comparison with fractional baselines, but it is a correctness/fairness issue, not a circularity step, because no result is equivalent to its input by construction.

Assumptions & free parameters 6 free parameters · 5 assumptions · 0 invented entities

The central claim depends on six hand-tuned hyperparameters and five domain/technical assumptions inherited from prior work. None of the assumptions are validated inside the paper, but they are standard in the wireless-ML literature. The tuning of tau, L, and the MLP sizes is acknowledged in Section V-A; these are not presented as predictions. No new entities (particles, forces, dimensions) are introduced.

free parameters (6)

Gumbel-Softmax temperature tau = 1
Chosen via 'extensive testing' in Section V-A; controls the hard/soft trade-off of association outputs and affects training dynamics.
Number of GNN update layers L = 2
Set in Section V-A; determines the model's receptive field and capacity.
Node/edge representation dimension = 512
Set in Section V-A; dimensions of B, C, and E representations in preprocessing.
MLP architecture (hidden layers, width) = 2 hidden layers, 1024 nodes
Set in Section V-A for all f_1..f_7; defines the learned function class.
Mini-batch size and batches per epoch = 5 samples, 400 mini-batches
Set in Section V-A; controls gradient noise during unsupervised training.
Warm restart schedule (SGDR) = initial period 50 epochs, multiplier 2, LR 1e-8 to 5e-5
Set in Section V-A following [11]; affects convergence and final performance.

assumptions (5)

domain assumption mmWave channel model of Eq. (15) from [10] with a uniform linear array
All simulations use this synthetic channel model; the generalization claims are only valid for this channel family.
domain assumption Perfect channel state information at the GNN
Explicitly assumed in Section V-A; imperfect CSI would change beamforming quality and association decisions.
domain assumption Unique association with each UE served by exactly one BS
Constraint (4a); the method enforces it via one-hot outputs, which is necessary for the sum-rate expression to be well-defined.
standard math Gumbel-Softmax is a valid differentiable surrogate for categorical sampling
Borrowed from [6],[7]; the gradient flow of the whole training relies on Eq. (12).
domain assumption GNN combinatorial generalization: shared per-node/per-edge functions transfer across graph sizes
Borrowed from [3] and used to test on larger UE counts without retraining (Section V-B3); no proof is given for this specific objective.

how reviews work

0 comments

Cite this review

Pith. "Pith review of Joint Beamforming and Integer User Association using a GNN with Gumbel-Softmax Reparameterizations." pith.science (2026). https://pith.science/paper/GROT3AVW

@misc{pith2026250605241,
  author       = {Pith},
  title        = {Pith review of: Joint Beamforming and Integer User Association using a GNN with Gumbel-Softmax Reparameterizations},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/GROT3AVW}},
  note         = {Machine review of arXiv:2506.05241}
}

read the original abstract

Machine learning (ML) models can effectively optimize a multi-cell wireless network by designing the beamforming vectors and association decisions. Existing ML designs, however, often needs to approximate the integer association variables with a probability distribution output. We propose a novel graph neural network (GNN) structure that jointly optimize beamforming vectors and user association while guaranteeing association output as integers. The integer association constraints are satisfied using the Gumbel-Softmax (GS) reparameterization, without increasing computational complexity. Simulation results demonstrate that our proposed GS-based GNN consistently achieves integer association decisions and yields a higher sum-rate, especially when generalized to larger networks, compared to all other fractional association methods.

Figures

Figures reproduced from arXiv: 2506.05241 by the authors.

**Figure 2.** Convergence behaviors for GNN models with STGS reparameteriza [PITH_FULL_IMAGE:figures/full_fig_p005_2.png] view at source ↗

**Figure 3.** User association factors from 20 random channels. (a) Integer [PITH_FULL_IMAGE:figures/full_fig_p005_3.png] view at source ↗

Figures from the paper (3 more)

**Figure 4.** Figure 4: Generalization behaviors for GNN models trained with 8 UEs and 2 [PITH_FULL_IMAGE:figures/full_fig_p005_4.png]

**Figure 5.** Figure 5: Generalization of GNN models to different power levels. The GNN [PITH_FULL_IMAGE:figures/full_fig_p006_5.png]

**Figure 6.** Figure 6: Performance comparison with different BSs and power settings. The [PITH_FULL_IMAGE:figures/full_fig_p006_6.png]

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

11 extracted references · 8 canonical work pages

[2]

User association for load balancing in heterogeneous cellular networks,

Q. Ye, B. Rong, Y . Chen, M. Al-Shalash, C. Caramanis, and J. G. Andrews, “User association for load balancing in heterogeneous cellular networks,”IEEE Transactions on Wireless Communications, vol. 12, no. 6, pp. 2706–2716, 2013

work page 2013
[1]

An iteratively weighted mmse approach to distributed sum-utility maximization for a mimo interfering broadcast channel,

Q. Shi, M. Razaviyaynet al., “An iteratively weighted mmse approach to distributed sum-utility maximization for a mimo interfering broadcast channel,”IEEE Trans. Signal Process, vol. 59, no. 9, pp. 4331–4340, 2011

work page 2011
[3]

Relational inductive biases, deep learning, and graph networks,

P. W. Battaglia, J. B. Hamricket al., “Relational inductive biases, deep learning, and graph networks,”arXiv preprint arXiv:1806.01261, 2018

arXiv 2018
[4]

Gnn-aided user association and beam selection for mmwave integrated heterogeneous networks,

W. Deng, Y . Liu, M. Li, and M. Lei, “Gnn-aided user association and beam selection for mmwave integrated heterogeneous networks,”IEEE Wireless Communications Letters, 2023

work page 2023
[5]

Gblinks: Gnn-based beam selection and link activation for ultra-dense d2d mmwave networks,

S. He, S. Xionget al., “Gblinks: Gnn-based beam selection and link activation for ultra-dense d2d mmwave networks,”IEEE Transactions on Communications, vol. 70, no. 5, pp. 3451–3466, 2022

work page 2022
[6]

Soft actor-critic with integer actions,

T.-H. Fan and Y . Wang, “Soft actor-critic with integer actions,” in2022 American Control Conference (ACC). IEEE, 2022, pp. 2611–2616

work page 2022
[7]

Categorical reparametrization with gumble-softmax,

E. Jang, S. Gu, and B. Poole, “Categorical reparametrization with gumble-softmax,” inInternational Conference on Learning Represen- tations (ICLR 2017). OpenReview. net, 2017

2017
[8]

Deep learning optimized sparse antenna activation for reconfigurable intelligent sur- face assisted communication,

S. Zhang, S. Zhang, F. Gao, J. Ma, and O. A. Dobre, “Deep learning optimized sparse antenna activation for reconfigurable intelligent sur- face assisted communication,”IEEE Transactions on Communications, vol. 69, no. 10, pp. 6691–6705, 2021

work page 2021

Show all 11 references

[9]

Graph neural network based beamforming and ris reflection design in a multi-ris assisted wireless network,

B. Lim and M. Vu, “Graph neural network based beamforming and ris reflection design in a multi-ris assisted wireless network,” in2023 IEEE SSP, 2023, pp. 120–124

2023
[10]

Millimeter wave channel modeling and cellular capacity evaluation,

M. R. Akdenizet al., “Millimeter wave channel modeling and cellular capacity evaluation,”IEEE journal on selected areas in communications, vol. 32, no. 6, pp. 1164–1179, 2014

2014
[11]

Sgdr: Stochastic gradient descent with warm restarts,

I. Loshchilov and F. Hutter, “Sgdr: Stochastic gradient descent with warm restarts,”arXiv preprint arXiv:1608.03983, 2016

2016 arXiv

Pith tools

Reviewed August 7, 2026 · model on record in the stance chip above.

[1] [2]

User association for load balancing in heterogeneous cellular networks,

Q. Ye, B. Rong, Y . Chen, M. Al-Shalash, C. Caramanis, and J. G. Andrews, “User association for load balancing in heterogeneous cellular networks,”IEEE Transactions on Wireless Communications, vol. 12, no. 6, pp. 2706–2716, 2013

work page 2013

[2] [1]

An iteratively weighted mmse approach to distributed sum-utility maximization for a mimo interfering broadcast channel,

Q. Shi, M. Razaviyaynet al., “An iteratively weighted mmse approach to distributed sum-utility maximization for a mimo interfering broadcast channel,”IEEE Trans. Signal Process, vol. 59, no. 9, pp. 4331–4340, 2011

work page 2011

[3] [3]

Relational inductive biases, deep learning, and graph networks,

P. W. Battaglia, J. B. Hamricket al., “Relational inductive biases, deep learning, and graph networks,”arXiv preprint arXiv:1806.01261, 2018

arXiv 2018

[4] [4]

Gnn-aided user association and beam selection for mmwave integrated heterogeneous networks,

W. Deng, Y . Liu, M. Li, and M. Lei, “Gnn-aided user association and beam selection for mmwave integrated heterogeneous networks,”IEEE Wireless Communications Letters, 2023

work page 2023

[5] [5]

Gblinks: Gnn-based beam selection and link activation for ultra-dense d2d mmwave networks,

S. He, S. Xionget al., “Gblinks: Gnn-based beam selection and link activation for ultra-dense d2d mmwave networks,”IEEE Transactions on Communications, vol. 70, no. 5, pp. 3451–3466, 2022

work page 2022

[6] [6]

Soft actor-critic with integer actions,

T.-H. Fan and Y . Wang, “Soft actor-critic with integer actions,” in2022 American Control Conference (ACC). IEEE, 2022, pp. 2611–2616

work page 2022

[7] [7]

Categorical reparametrization with gumble-softmax,

E. Jang, S. Gu, and B. Poole, “Categorical reparametrization with gumble-softmax,” inInternational Conference on Learning Represen- tations (ICLR 2017). OpenReview. net, 2017

2017

[8] [8]

Deep learning optimized sparse antenna activation for reconfigurable intelligent sur- face assisted communication,

S. Zhang, S. Zhang, F. Gao, J. Ma, and O. A. Dobre, “Deep learning optimized sparse antenna activation for reconfigurable intelligent sur- face assisted communication,”IEEE Transactions on Communications, vol. 69, no. 10, pp. 6691–6705, 2021

work page 2021