End-to-End Learning of Probabilistic Constellation Shaping through Importance Sampling

Laurent Schmalen; Shrinivas Chimmalgi; Vahid Aref

REVIEW 2 major objections 4 minor 20 references

End-to-End Learning of Probabilistic Constellation Shaping through Importance Sampling

T0 review · 2 major / 4 minor · reviewed 2026-08-06 · deepseek-v4-flash

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

Pith's one-line read The paper proposes weighted importance-sampling loss functions that let an autoencoder learn constellation point probabilities for probabilistic shaping, with automatic differentiation alone producing exact gradients of mutual information…

desk verdict Useful importance-sampling loss for autoencoder-based constellation shaping, but the exact-gradient proof conflates the proposal distribution with the empirical batch frequency; the claim needs a fix. read the letter →

arxiv 2506.16098 v1 pith:AZECDTJU submitted 2025-06-19 cs.IT eess.SPmath.IT

Shrinivas Chimmalgi , Laurent Schmalen , Vahid Aref This is my paper · ORCID

classification cs.ITeess.SPmath.IT

keywords probabilisticconstellationshapingimportancesamplingautomaticdifferentiationmutualinformationbit-wiseautoencoderAWGNchannelIM/DD

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

Probabilistic constellation shaping adjusts how often each constellation point is sent, and the previous autoencoder approach required manually adding correction terms to the gradients computed by automatic differentiation. This paper removes that step by writing the objective as an importance-sampling estimate: symbols are drawn from an auxiliary distribution $Q_M$ that approximates the target probabilities $P_M$, and every sample is re-weighted by $w_m = p_m/q_m$. The authors prove that gradient ascent on the two proposed losses, one for mutual information and one for bit-wise mutual information, yields exactly the same gradient in the constellation probabilities as the earlier gradient-correction method. They demonstrate on a 256-QAM AWGN channel and an 8-PAM intensity-modulation/direct-detection channel that the learned shaping gains match the baseline within numerical accuracy. If the claim holds, engineers can shape constellations for new channels with standard machine-learning tooling and no hand-derived gradient terms.

What carries the argument

The load-bearing object is the weighted importance-sampling loss, in which a batch of symbols is drawn from a sampling distribution $Q_M$ (a variational-distance approximation of the target $P_M$), each symbol's contribution is weighted by $w_m = p_m/q_m$, and the constellation normalization is evaluated with the induced distribution $\tilde P_M = W_M \cdot Q_M$. The identity that carries the exactness proof is $\partial f/\partial w_m = q_m\,\partial f/\partial p_m$ combined with the fact that the statistics of the drawn batches do not depend on $w_m$, so automatic differentiation through the loss directly returns the true gradient in the constellation probabilities. The outer loop of Algorithm 1 refreshes $Q_M$ once per epoch from the current $P_M$, keeping the two distributions close and ensuring every point has $q_m N \ge 1$; this is what lets the method use relatively short batches without rare-symbol dropout.

What would settle it

Run Algorithm 1 on a strongly shaped 256-QAM distribution with a small batch size (for example $N = 2^8$) so that the rarest points violate $q_m N \ge 1$, and compare the automatic-differentiation update of $w_m$ against the exact gradient of the finite-sample objective, or against the gradient-addition method's update. If the two move in different directions or diverge, the claimed exactness holds only in the regime where $|B_m|/N$ is a faithful surrogate for $q_m$.

Watch

Extended reading notes

Core claim

The paper's central claim is that Eqs. (2) and (5) are valid autoencoder loss functions for learning constellation point probabilities, and that the gradients of these losses with respect to the weights $w_m = p_m/q_m$, computed by automatic differentiation, coincide with the exact gradients of the mutual information and bit-wise mutual information with respect to $p_m$. The analytical step is Eq. (4), which follows from $\partial f/\partial w_m = q_m\,\partial f/\partial p_m$ and reproduces the gradient expression of the previous method. The paper further claims that this removes the need for the manual correction terms in the earlier approach, that the per-epoch refresh of the sampler $Q_M$ via variational-distance approximation keeps $q_m N \ge 1$ and thereby avoids dropout of rare constellation points, and that simulations on AWGN and IM/DD channels match the earlier method to within about 0.01 bits/symbol.

Load-bearing premise

The exact-gradient claim assumes that within each batch the empirical frequency $|B_m|/N$ of a constellation point is a faithful stand-in for its true sampling probability $q_m$, so the whole argument depends on batches being large enough and the variational-distance sampler giving every point $q_m N \ge 1$.

Editorial extensions

If this is right

The same optimization trajectory as the gradient-addition method is available from a simpler implementation, because the loss itself, rather than extra manual terms, carries the exact gradient.
Engineers can optimize for either symbol-metric decoding (MI) or bit-metric decoding (BMI) with the same weight-learning scheme, switching only the receiver output in the loss.
Training with relatively short symbol sequences remains viable for strongly shaped constellations, since the per-epoch sampler refresh keeps low-probability points represented ($q_m N \ge 1$).
The reported MI gap matches the baseline within about 0.01 bits/symbol on 256-QAM AWGN, and the BMI gains match on 8-PAM IM/DD, so the method's practical behavior tracks the theoretically corrected baseline.
Since the demapper can be a neural network and the channel model is just a pass-through, the loss functions transfer to other channels without re-deriving gradients.

Reading between the lines

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

The exact-gradient claim is best read as asymptotic in batch size: the proof identifies $|B_m|/N$ with $q_m$, so at small batches the estimated loss is biased and the guarantee should soften gradually rather than fail at a sharp threshold.
A natural extension the paper does not test is joint geometric and probabilistic shaping: with the constellation coordinates also trainable, the automatic-differentiation path through normalization should expose gradients with respect to coordinates as well, and the same importance-sampling loss could learn both.
On channels where noise depends on the transmitted symbol sequence, such as nonlinear optical fibers, the per-epoch refresh of $Q_M$ should beat fixed-distribution sampling; the paper motivates this scenario but leaves the experiment for future work.

Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

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

Colleague,

The short version: this paper offers a genuinely simpler way to train autoencoder-based constellation shaping — replace the manual gradient correction of Aref–Chagnon with importance-weighted losses and refresh the proposal each epoch — and the simulations show it matches the baselines on AWGN and a simplified IM/DD channel. The code is public. That part is real and useful.

The soft spot is in the proof. The paper claims Eq. (4) is the exact gradient of the MI lower bound with respect to the symbol probabilities. The derivation hinges on the line "Then q_m = |B_m|/N" in Sec. II. That equates the per-epoch proposal distribution, which is fixed before sampling, with the batch's empirical frequency. These are different objects. Under the algorithm's actual definition w_m = p_m/q_m, differentiating (2) gives an extra bias term of order (1-q_m)/(2N q_m). For constellation points at the variational lower bound q_m N = 1 that bias is O(1), not O(1/N). So the exact-gradient claim is not established for finite batch sizes. The method may still be consistent as N→∞, and the experiments suggest it works, but the paper should say so honestly.

Two more minor issues: the simulation curves have no error bars, so we cannot tell whether the sub-0.01 bit/symbol differences are noise; and the advertised advantage for short symbol sequences is not experimentally tested — the paper only mentions the peak-power AWGN agreement, and those results are omitted.

My overall impression: the idea is worth taking seriously, and the empirical match to [1] suggests the biased gradient does no harm in these cases. But the central theoretical claim needs repair: either present (2) with the true proposal q_m in the first logarithm, or provide a finite-sample bias analysis. As written, the paper overstates its analytic guarantee.

I'd send it to review — a good referee would catch the same issue and the authors can fix it. I'd bring it up in a reading group if you want a discussion about importance sampling in autoencoder training. I would not cite the exact-gradient claim in my own work until the bias is addressed.

Referee Report

2 major / 4 minor

Summary. The paper proposes an autoencoder-based method for probabilistic constellation shaping in which the mutual information (MI) or bit-wise mutual information (BMI) is optimized with loss functions that use importance-sampling weights w_m = p_m/q_m and a proposal distribution Q_M refreshed once per epoch. The authors claim that, unlike prior work [1], their loss functions allow the constellation-point probabilities to be optimized using automatic differentiation alone, and that the resulting gradients are exact. The method is validated on a 256-QAM AWGN channel and an 8-PAM intensity-modulation/direct-detection channel, with results reported to match the gradient-addition baseline [1] within 0.01 bits/symbol. The paper also states that code and results are publicly available.

Significance. If the exact-gradient claim were correct, the paper would make a useful practical contribution: it would remove the error-prone manual gradient corrections of [1], allow shorter symbol sequences through importance sampling, and apply to arbitrary channel models where the demapper is learned. The paper is clearly written, the simulations cover two relevant channel models, and the release of code is a strength. However, the central theoretical claim is not established by the derivation as written, because the proof identifies the epoch-level proposal q_m with the random empirical frequency |B_m|/N of a finite batch. That identification is load-bearing for Eq. (4) and for the abstract's claim of exact gradients. The empirical agreement with baselines is encouraging but is reported without error bars, so it does not by itself resolve the finite-sample bias concern.

major comments (2)

[Section III, Figs. 2-4] The derivation of Eq. (4) from Eq. (3) uses the identification q_m = |B_m|/N, but in Algorithm 1 q_m is the variational-distance approximation fixed for the whole epoch, whereas |B_m|/N is the random empirical frequency of a single batch. These two quantities are not equal. With q_m fixed and w_m = p_m/q_m, the expected value of the first term of Eq. (2) is -(1/N) \sum_m w_m E[K_m log(K_m/N)] with K_m = |B_m|, which is not equal to -\sum_m p_m log q_m. A Taylor expansion of the multinomial expectation gives E[K_m log(K_m/N)] = N q_m log q_m + (1-q_m)/2 + O(1/N), so the loss itself is a biased estimate of the bound (1). The gradient with respect to p_m divides by q_m, producing a bias of order (1-q_m)/(2N q_m); for a constellation point at the variational lower bound q_m = 1/N, this bias is O(1), not O(1/N). Therefore Eq. (4) is the gradient of the realized batch loss, not the exact gradient of the MI bound, and the abstract's claim that the method 'uses exact gradients' is unsupported. The authors should either prove unbiasedness of the gradient estimator under the actual sampling scheme or explicitly characterize the bias and reframe the claim as a consistent large-N approximation.
[Section III, Figs. 2-4] The numerical comparisons report that differences between the proposed method and the baselines remain below 0.01 bits/symbol, but no confidence intervals, number of independent runs, or error bars are provided. Since the theoretical exactness of the gradients is the main claim of the paper, the simulations should include a direct finite-batch bias check, for example by comparing the batch gradient used in training with a high-accuracy estimate of the gradient of (1) at the same q_m, or by reporting the variability of the optimized probabilities across random restarts. Without such a check, the close agreement with [1] is suggestive but does not validate the exact-gradient claim at the batch size used in the paper.

minor comments (4)

[Section II, Eq. (4)] Equation (4) contains the factor 1/|B_m|, which is undefined when a constellation point is not sampled in a batch. The corresponding automatic-differentiation derivative is zero because the sum over the empty set is empty. Please clarify how empty batches are handled and whether the displayed formula is meant to apply only when |B_m| > 0.
[Algorithm 1] The algorithm updates W_M in the inner loop and then sets P_M = W_M Q_M at the end of the epoch, but the text does not state whether P_M is renormalized to sum to one. Since P_M is used for constellation normalization and as the target distribution, please clarify the normalization step explicitly.
[Section II, paragraph after Eq. (2)] The sentence 'Then q_m = |B_m|/N' is the crux of the derivation and is currently ambiguous. It should be replaced with an explicit statement of whether q_m denotes the epoch-level proposal or the empirical frequency, because the two choices lead to different and incompatible gradient derivations.
[Section III-A] The text mentions that the peak-power AWGN study agrees with Blahut-Arimoto but omits the results due to space constraints. Since the code is available, consider moving this result to an appendix or providing a brief quantitative summary in the text, as it is a strong validation of the method.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: the importance-sampling gradient derivation is self-contained and checked against external benchmarks.

full rationale

The paper's central claim is that losses (2) and (5), with weights w_m = p_m/q_m, supply exact gradients of MI/BMI with respect to the constellation probabilities using automatic differentiation. The derivation in Sec. II starts from a known lower bound [15, Eq. 34], defines the loss, and computes partial derivatives directly; Eq. (4) follows by algebra and is then compared with [1, Eq. (3)] as a consistency check, not derived from it. Reference [1] is co-authored by one of the present authors, but it is used as a baseline and as a point of comparison, not as the source of the gradient formula, so the self-citation is not load-bearing. The variational-distance approximation and sampling procedure are taken from an external textbook [14], and the simulations are benchmarked against optimized Maxwell–Boltzmann distributions [9], [16], [17] and, in the omitted peak-power study, against Blahut–Arimoto (noted in Sec. III-A as omitted for space, with code made available). The one questionable identification is the in-text statement “Then q_m = |B_m|/N” in Sec. II: this equates the fixed proposal distribution with the empirical batch frequency. That is a finite-sample statistical approximation (and a possible source of estimator bias at small batch size) rather than a circular self-definition, because q_m is originally fixed per epoch and the target result (gradient of the loss) is not defined in terms of q_m. The omitted Blahut–Arimoto comparison and any finite-batch bias are correctness/completeness concerns, not circular reductions. The derivation is therefore self-contained against external benchmarks, and no step reduces to its own input by construction.

Assumptions & free parameters 0 free parameters · 4 assumptions · 0 invented entities

The method has no hand-fitted constants: the target probabilities are learned. The assumptions are the validity of the variational distance approximation for the proposal distribution, the unbiasedness (up to finite-sample bias) of the importance-sampling estimator, the exactness of the matched demappers, and the correctness of autodiff through the normalization block.

assumptions (4)

domain assumption The variational distance minimization (from [14]) yields a proposal distribution QM with q_m N >= 1 for all m.
Invoked in Sec. II to construct QM and to bound the importance weights.
standard math The empirical loss (2) is an importance-sampling estimate of the MI bound (1) when symbols are drawn from QM.
Standard importance sampling property used in Sec. II, up to the finite-sample bias of replacing q_m with |B_m|/N in the log terms.
domain assumption The demappers used in simulations are perfectly matched to the channel laws, isolating the learning method.
Stated in Sec. III-A and III-B.
standard math Automatic differentiation computes the derivative of (2) with respect to the weights w_m, including the term from the normalization of constellation points via p_tilde_M.
The algorithm depends on autodiff's ability to propagate gradients through the normalization block (Alg. 1, Sec. II).

how reviews work

0 comments

Cite this review

Pith. "Pith review of End-to-End Learning of Probabilistic Constellation Shaping through Importance Sampling." pith.science (2026). https://pith.science/paper/AZECDTJU

@misc{pith2026250616098,
  author       = {Pith},
  title        = {Pith review of: End-to-End Learning of Probabilistic Constellation Shaping through Importance Sampling},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/AZECDTJU}},
  note         = {Machine review of arXiv:2506.16098}
}

read the original abstract

Probabilistic constellation shaping enables easy rate adaption and has been proven to reduce the gap to Shannon capacity. Constellation point probabilities are optimized to maximize either the mutual information or the bit-wise mutual information. The optimization problem is however challenging even for simple channel models. While autoencoder-based machine learning has been applied successfully to solve this problem [1], it requires manual computation of additional terms for the gradient which is an error-prone task. In this work, we present novel loss functions for autoencoder-based learning of probabilistic constellation shaping for coded modulation systems using automatic differentiation and importance sampling. We show analytically that our proposed approach also uses exact gradients of the constellation point probabilities for the optimization. In simulations, our results closely match the results from [1] for the additive white Gaussian noise channel and a simplified model of the intensity-modulation direct-detection channel.

Figures

Figures reproduced from arXiv: 2506.16098 by the authors.

**Figure 1.** Block diagram of the proposed end-to-end learning system. The arrow directions mark the order of operations. Operations marked with [PITH_FULL_IMAGE:figures/full_fig_p003_1.png] view at source ↗

**Figure 2.** Gap of MI to Shannon capacity for uniform and PCS 256-QAM. The [PITH_FULL_IMAGE:figures/full_fig_p003_2.png] view at source ↗

**Figure 3.** Gap of BMI to Shannon capacity for uniform and PCS 256-QAM. The [PITH_FULL_IMAGE:figures/full_fig_p003_3.png] view at source ↗

Figures from the paper (1 more)

**Figure 4.** Figure 4: BMI for uniform and PCS 8-PAM. The symbol probabilities have [PITH_FULL_IMAGE:figures/full_fig_p004_4.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

20 extracted references · 20 canonical work pages

[1]

End-to-end learning of joint geometric and probabilistic constellation shaping,

V . Aref and M. Chagnon, “End-to-end learning of joint geometric and probabilistic constellation shaping,” in Proc. OFC, 2022

work page 2022
[2]

On shaping gain in the nonlinear fiber-optic channel,

R. Dar et al., “On shaping gain in the nonlinear fiber-optic channel,” in Proc. ISIT, 2014

work page 2014
[3]

A shaping algorithm for mitigating inter-channel nonlinear phase-noise in nonlinear fiber systems,

O. Geller et al., “A shaping algorithm for mitigating inter-channel nonlinear phase-noise in nonlinear fiber systems,” J. Light. Technol. , vol. 34, no. 16, Aug. 2016

work page 2016
[4]

Introducing enumerative sphere shaping for optical communication systems with short blocklengths,

A. Amari et al., “Introducing enumerative sphere shaping for optical communication systems with short blocklengths,” J. Light. Technol. , vol. 37, no. 23, Dec. 2019

work page 2019
[5]

Probabilistic constellation shaping for optical fiber communications,

J. Cho and P. J. Winzer, “Probabilistic constellation shaping for optical fiber communications,” J. Light. Technol. , vol. 37, no. 6, Mar. 2019

work page 2019
[6]

On probabilistic shaping of quadrature amplitude modulation for the nonlinear fiber channel,

T. Fehenberger et al., “On probabilistic shaping of quadrature amplitude modulation for the nonlinear fiber channel,” J. Light. Technol. , vol. 34, no. 21, Nov. 2016

work page 2016
[7]

Non-linear phase noise mitigation over systems using constellation shaping,

D. Pilori et al., “Non-linear phase noise mitigation over systems using constellation shaping,” J. Light. Technol. , vol. 37, no. 14, Jul. 2019

work page 2019
[8]

Parallel bisection-based distribution matching for nonlinearity-tolerant probabilistic shaping in coherent optical commu- nication systems,

M. Fu et al., “Parallel bisection-based distribution matching for nonlinearity-tolerant probabilistic shaping in coherent optical commu- nication systems,” J. Light. Technol. , vol. 39, no. 20, Oct. 2021

work page 2021

Show all 20 references

[9]

Achievable rates for probabilistic shaping,

G. B ¨ocherer, “Achievable rates for probabilistic shaping,” arXiv preprint arXiv:1707.01134, Jul. 2017

2017 arXiv
[10]

Joint learning of geometric and probabilistic constella- tion shaping,

M. Stark et al., “Joint learning of geometric and probabilistic constella- tion shaping,” in Proc. GLOBECOM, 2019

2019
[11]

Joint learning of probabilistic and geo- metric shaping for coded modulation systems,

F. Ait Aoudia and J. Hoydis, “Joint learning of probabilistic and geo- metric shaping for coded modulation systems,” in Proc. GLOBECOM , 2020

2020
[12]

autodiff, a modern, fast and expressive C++ library for automatic differentiation,

A. M. M. Leal, “autodiff, a modern, fast and expressive C++ library for automatic differentiation,” https://autodiff.github.io, 2018

2018
[13]

Memory-aware end-to-end learning of channel distortions in optical coherent communications,

V . Neskorniuk et al., “Memory-aware end-to-end learning of channel distortions in optical coherent communications,” Optics Express, vol. 31, no. 1, Dec. 2022

2022
[14]

Probabilistic amplitude shaping,

G. B ¨ocherer, “Probabilistic amplitude shaping,” F ound. Trends Commun. Inf. Theory , vol. 20, no. 4, 2023

2023
[15]

Simulation-based computation of information rates for channels with memory,

D. Arnold et al., “Simulation-based computation of information rates for channels with memory,” IEEE Trans. Inf. Theory , vol. 52, no. 8, 2006

2006
[16]

Optimal nonuniform signaling for Gaussian channels,

F. R. Kschischang and S. Pasupathy, “Optimal nonuniform signaling for Gaussian channels,” IEEE Trans. Inf. Theory , vol. 39, no. 3, May 1993

1993
[17]

Probabilistic shaping for the AWGN channel,

S. Delsad, “Probabilistic shaping for the AWGN channel,” arXiv preprint arXiv:2302.08388, Feb. 2023

2023 arXiv
[18]

Training deep networks without learning rates through coin betting,

F. Orabona and T. Tommasi, “Training deep networks without learning rates through coin betting,” in Proc. NeurIPS, vol. 30, 2017

2017
[19]

A note on multivariate Gauss-Hermite quadrature,

P. J ¨ackel, “A note on multivariate Gauss-Hermite quadrature,” 2005. [Online]. Available: http://www.jaeckel.org/ ANoteOnMultivariateGaussHermiteQuadrature.pdf

2005
[20]

End-to-end learning of probabilistic constellation shaping through importance sampling

S. Chimmalgi et al., “End-to-end learning of probabilistic constellation shaping through importance sampling.” [Online]. Available: https: //github.com/kit-cel/E2E-PCS-through-importance-sampling

Pith tools

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

[1] [1]

End-to-end learning of joint geometric and probabilistic constellation shaping,

V . Aref and M. Chagnon, “End-to-end learning of joint geometric and probabilistic constellation shaping,” in Proc. OFC, 2022

work page 2022

[2] [2]

On shaping gain in the nonlinear fiber-optic channel,

R. Dar et al., “On shaping gain in the nonlinear fiber-optic channel,” in Proc. ISIT, 2014

work page 2014

[3] [3]

A shaping algorithm for mitigating inter-channel nonlinear phase-noise in nonlinear fiber systems,

O. Geller et al., “A shaping algorithm for mitigating inter-channel nonlinear phase-noise in nonlinear fiber systems,” J. Light. Technol. , vol. 34, no. 16, Aug. 2016

work page 2016

[4] [4]

Introducing enumerative sphere shaping for optical communication systems with short blocklengths,

A. Amari et al., “Introducing enumerative sphere shaping for optical communication systems with short blocklengths,” J. Light. Technol. , vol. 37, no. 23, Dec. 2019

work page 2019

[5] [5]

Probabilistic constellation shaping for optical fiber communications,

J. Cho and P. J. Winzer, “Probabilistic constellation shaping for optical fiber communications,” J. Light. Technol. , vol. 37, no. 6, Mar. 2019

work page 2019

[6] [6]

On probabilistic shaping of quadrature amplitude modulation for the nonlinear fiber channel,

T. Fehenberger et al., “On probabilistic shaping of quadrature amplitude modulation for the nonlinear fiber channel,” J. Light. Technol. , vol. 34, no. 21, Nov. 2016

work page 2016

[7] [7]

Non-linear phase noise mitigation over systems using constellation shaping,

D. Pilori et al., “Non-linear phase noise mitigation over systems using constellation shaping,” J. Light. Technol. , vol. 37, no. 14, Jul. 2019

work page 2019

[8] [8]

Parallel bisection-based distribution matching for nonlinearity-tolerant probabilistic shaping in coherent optical commu- nication systems,

M. Fu et al., “Parallel bisection-based distribution matching for nonlinearity-tolerant probabilistic shaping in coherent optical commu- nication systems,” J. Light. Technol. , vol. 39, no. 20, Oct. 2021

work page 2021