Neural Probabilistic Amplitude Shaping

• Neural probabilistic amplitude shaping is a deep learning-driven extension of PAS that jointly optimizes symbol distributions, constellation geometry, and FEC compatibility.
• It leverages autoregressive models and rate-loss control to adapt to both memoryless and nonlinear channels, improving bit-wise mutual information and effective SNR.
• Integrating sequential modeling with arithmetic distribution matching, the method mitigates implementation losses and enhances modulation performance in practical systems.

Neural probabilistic amplitude shaping denotes a family of learned coded-modulation methods that retain the central PAS idea—nonuniform shaping of constellation amplitudes combined with binary FEC—while replacing fixed hand-designed shaping rules with trainable models. In the literature, the term covers at least two related directions. One direction generalizes PAS for memoryless or mismatched channels by jointly learning the symbol distribution, constellation geometry, bit labeling, and demapping, and by conditioning the shaping law on SNR (Aoudia et al., 2020). A second direction is specific to nonlinear coherent fiber systems and learns the joint distribution of unsigned symbol sequences, rather than only a marginal amplitude law, because nonlinear interference depends on temporal symbol patterns (Askari et al., 2 Feb 2026). More recent work makes this sequence modeling rate-loss-aware and block-less through sequential autoregressive encoders compatible with arithmetic distribution matching (Askari et al., 27 May 2026).

1. Classical PAS as the substrate for neural extensions

Probabilistic amplitude shaping is a layered coded-modulation architecture in which the transmitted symbol distribution is factorized into a shaped amplitude component and a typically uniform sign component. In the standard PAS construction, a distribution matcher generates amplitudes according to a target nonuniform law, and a systematic FEC encoder provides parity bits that act as signs. For finite constellations and bit-metric decoding, PAS is attractive because it preserves compatibility with binary coding while approximating capacity-achieving nonuniform inputs (Böcherer, 2017).

The information-theoretic backdrop is that layered probabilistic shaping separates the shaping-induced transmission rate from the code rate of a larger random codebook. For a general decoding metric $q$, the achievable transmission rate can be written as

$$t_x = \left[ H(X)-\mathbb E\!\left[-\log_2\frac{q(X,Y)}{\sum_a q(a,Y)}\right] \right]^+,$$

and with the optimal metric it becomes $I(X;Y)$ (Böcherer, 2017). PAS is a practical instance of this framework, and later theoretical treatments recast it as a JSCC architecture with a distribution matcher, a systematic FEC encoder, a mismatched FEC decoder, and a dematcher, clarifying when practical decoder mismatch is asymptotically harmless (Amjad, 2018). A complementary weak-typicality treatment shows that PAS achievability can be expressed through random sign-coding arguments, with amplitudes produced constructively and only some or all signs randomized (Gültekin et al., 2020).

These classical formulations also identify structural constraints that motivate neural generalizations. Standard PAS assumes a symmetric target distribution and, in the usual 2D setting, is restricted to code rates $\frac{m-2}{m}$ or higher. Finite-length PAS additionally incurs rate loss because practical distribution matchers operate on blocks rather than on ideal i.i.d. amplitude sources. Neural PAS methods are best understood as attempts to relax one or more of these constraints while preserving PAS-compatible implementation structure (Aoudia et al., 2020).

2. Neural generalization beyond fixed Maxwell–Boltzmann shaping

A foundational neural extension appears in a trainable coded-modulation scheme that jointly optimizes probabilistic shaping, geometric shaping, bit labeling, and demapping for a specific channel model and over a wide range of SNRs (Aoudia et al., 2020). In this formulation, the learned system is not restricted to symmetric distributions, can be optimized for channels beyond AWGN, and is compatible with any code rate $r=\frac{k}{m}$ with $k\in\{1,\dots,m-1\}$. This directly targets two limitations of classical PAS: the symmetry constraint imposed by Maxwell–Boltzmann shaping on QAM, and the rate constraint inherited from the standard sign/parity construction.

The architecture partitions the constellation into $2^{m-k}$ sub-constellations of size $2^k$, shapes the information-bearing part according to a trainable law $p_{\boldsymbol{\theta}}$, and keeps the parity bits uniform and i.i.d. because channel coding does not preserve shaping. Its optimization objective is the bit-wise mutual information

$$R(\boldsymbol{\chi}) \coloneqq \left[ H(X) - \sum_{i=1}^m H(B_i|Y) \right]^+,$$

with a BCE-based surrogate that simultaneously maximizes BMI and trains the differentiable demapper to approximate the true posterior. An important implementation device is that gradients with respect to $$t_x = \left[ H(X)-\mathbb E\!\left[-\log_2\frac{q(X,Y)}{\sum_a q(a,Y)}\right] \right]^+,$$0 are obtained by sampling the channel state and summing explicitly over constellation points, instead of relying on Gumbel-Softmax or other stochastic-neuron relaxations (Aoudia et al., 2020).

The same work also defines a deep-learning-extended PAS baseline in which a neural network maps SNR to the MB shaping parameter $$t_x = \left[ H(X)-\mathbb E\!\left[-\log_2\frac{q(X,Y)}{\sum_a q(a,Y)}\right] \right]^+,$$1, thereby learning a continuum of MB distributions for QAM without changing the classical PAS structure. This distinction is consequential. In the reported AWGN experiments with $$t_x = \left[ H(X)-\mathbb E\!\left[-\log_2\frac{q(X,Y)}{\sum_a q(a,Y)}\right] \right]^+,$$2, PS-GS at $$t_x = \left[ H(X)-\mathbb E\!\left[-\log_2\frac{q(X,Y)}{\sum_a q(a,Y)}\right] \right]^+,$$3 and MB-QAM achieve essentially the same BMI, whereas under mismatched Rayleigh block fading the jointly learned PS-GS variants outperform MB-QAM and achieve the highest BMI lower bound at $$t_x = \left[ H(X)-\mathbb E\!\left[-\log_2\frac{q(X,Y)}{\sum_a q(a,Y)}\right] \right]^+,$$4 (Aoudia et al., 2020). This suggests that the main value of neural PAS is not merely tuning a Maxwell–Boltzmann family parameter, but replacing the assumed family itself when the channel departs from the conditions under which MB is near-optimal.

3. Joint-distribution learning for nonlinear fiber channels

In nonlinear coherent fiber links, the salient design variable is no longer only the marginal amplitude distribution. The central observation behind neural probabilistic amplitude shaping for fiber is that Kerr nonlinearity makes the distortion depend on neighboring symbols in time and across WDM channels, so the joint distribution of symbol sequences matters (Askari et al., 2 Feb 2026). Standard PAS shapes amplitudes largely through a prescribed marginal law, and sequence-selection methods search among candidate blocks generated under a fixed marginal law, but neither directly learns the conditional probabilities defining the full block distribution.

The proposed PAS-compatible NPAS framework therefore models a block of unsigned amplitudes $$t_x = \left[ H(X)-\mathbb E\!\left[-\log_2\frac{q(X,Y)}{\sum_a q(a,Y)}\right] \right]^+,$$5 through the autoregressive factorization

$$t_x = \left[ H(X)-\mathbb E\!\left[-\log_2\frac{q(X,Y)}{\sum_a q(a,Y)}\right] \right]^+,$$6

implemented with a recurrent neural model. Training uses Gumbel-Softmax sampling with a straight-through estimator, while runtime deployment uses an arithmetic distribution matcher that deterministically maps information bits to unsigned symbols according to the learned probabilities. The sign bits remain uniform and are supplied by the FEC parity stream, preserving compatibility with PAS architectures. The work explicitly contrasts NPAS with NPS, which learns joint distributions over signed symbols and is therefore not directly PAS-compatible (Askari et al., 2 Feb 2026).

The reported fiber setup is a single-span 205 km single-mode-fiber link with 5 WDM channels, dual-polarization 64-QAM, 50 GBd, 55 GHz channel spacing, dispersion $$t_x = \left[ H(X)-\mathbb E\!\left[-\log_2\frac{q(X,Y)}{\sum_a q(a,Y)}\right] \right]^+,$$7, nonlinear coefficient $$t_x = \left[ H(X)-\mathbb E\!\left[-\log_2\frac{q(X,Y)}{\sum_a q(a,Y)}\right] \right]^+,$$8, attenuation $$t_x = \left[ H(X)-\mathbb E\!\left[-\log_2\frac{q(X,Y)}{\sum_a q(a,Y)}\right] \right]^+,$$9, an EDFA with 5 dB noise figure, RRC roll-off $I(X;Y)$0, electronic CD compensation, and pilot-aided linear CPR with 2.5% pilot rate. Training uses an additive-multiplicative perturbation model as a differentiable nonlinear-fiber surrogate, whereas evaluation is performed in SSFM. At optimal launch power, NPAS and NPS outperform ESS and ESS plus sequence selection by more than $I(X;Y)$1 dB in effective SNR and about $I(X;Y)$2 bits/2D in AIR; with $I(X;Y)$3, the study reports stable improvement for NPAS while NPS plateaus and then degrades as block length grows (Askari et al., 2 Feb 2026).

These results establish the defining claim of fiber-oriented neural PAS: the optimized object is a block distribution over amplitudes, not merely a Maxwell–Boltzmann-like histogram. At the same time, the same study states that ADM rate loss is not yet explicitly quantified for NPAS and NPS because the training relies on Gumbel-Softmax sampling, so the reported ESS AIR is an optimistic bound that does not subtract ESS rate losses (Askari et al., 2 Feb 2026). That caveat became the starting point for later rate-loss-aware neural PAS.

4. Sequential neural PAS and explicit rate-loss control

Sequential Neural Probabilistic Amplitude Shaping, or Seq-NPAS, extends neural PAS by making the cost of learned dependencies explicit in both the formulation and the implementation (Askari et al., 27 May 2026). Its starting point is that temporal correlations can improve nonlinear tolerance but also induce rate loss when realized by an actual distribution matcher. The relevant quantity is defined as

$I(X;Y)$4

which measures the gap between the sum of marginal entropies and the entropy rate of the joint process. In this formulation, $I(X;Y)$5 is the minimum rate loss of an ideal distribution matcher for the learned joint distribution and depends only on the learned sequence statistics (Askari et al., 27 May 2026).

Earlier neural shaping objectives optimized a BCE-based surrogate for instantaneous bit-metric rate and demapper mismatch. Seq-NPAS augments that objective to

$I(X;Y)$6

thereby penalizing implementation cost and balancing learned temporal structure against marginal shaping toward a Maxwell–Boltzmann target. This is not a cosmetic modification: the reported experiments show that ignoring $I(X;Y)$7 leads to hidden rate penalties large enough to erase apparent nonlinear gains, whereas including it yields better net AIR (Askari et al., 27 May 2026).

The architectural shift is from blockwise recurrent modeling to a stationary finite-memory autoregressive source,

$I(X;Y)$8

so that the same prediction rule is applied at every position. The method is described as block-less because there are no artificial block boundaries, no per-block reinitialization, and no deliberate neglect of cross-block correlations. It is also naturally compatible with arithmetic distribution matching: ADM can generate the shaped sequence symbol by symbol from the learned conditional probabilities $I(X;Y)$9. The reported practical ADM rate loss,

$\frac{m-2}{m}$0

approaches the theoretical lower bound $\frac{m-2}{m}$1 as the ADM input length grows (Askari et al., 27 May 2026).

In the nonlinear dual-polarization WDM fiber-link experiments, NPAS++ and Seq-NPAS++ achieve up to $\frac{m-2}{m}$2 bits/2D higher AIR than uniform signaling, ESS, ESS plus sequence selection, and earlier NPAS baselines by learning symbol dependencies while controlling rate loss. Among the tested methods, Seq-NPAS++ provides the best overall performance, and the work presents it as the first practical neural shaping method to outperform existing schemes once all implementation losses are included (Askari et al., 27 May 2026). A plausible implication is that, in nonlinear channels, the central question is not whether neural models can learn dependencies, but whether those dependencies survive the entropy-rate budget of the distribution-matching layer.

5. Finite blocklength, sphere shaping, and sequence-level baselines

Neural PAS is embedded in a longer line of PAS research on finite-length shaping and sequence-level nonlinear mitigation. For short blocklengths on AWGN, the dominant issue is often not exact PMF matching but the rate loss created by the finite shaping set. A systematic comparison of CCDM, MPDM, ESS, and shell mapping argues that, at practical blocklengths, shaping should be viewed as designing the most energy-efficient finite signal space for a given rate rather than matching a target PMF (Gültekin et al., 2019). In the specific $\frac{m-2}{m}$3 example, the reported rate losses are $\frac{m-2}{m}$4 bit/1-D for CCDM, $\frac{m-2}{m}$5 bit/1-D for MPDM, and $\frac{m-2}{m}$6 bit/1-D for ESS; ESS and MPDM are about $\frac{m-2}{m}$7 dB more power-efficient than uniform 8-ASK at $\frac{m-2}{m}$8 bit/1-D, and end-to-end PAS simulations show up to 1 dB power-efficiency improvement over uniform signaling around blocklength 200 (Gültekin et al., 2019).

A related complexity-performance compromise is partial enumerative sphere shaping, which shapes only a subset of amplitude bits while leaving the remainder uniform and independent. For 16-ASK at 3 bit/symbol, shaping 2 amplitude bits rather than all 3 yields almost the same performance as full ESS, with about $\frac{m-2}{m}$9 dB improvement over uniform signaling and only about $r=\frac{k}{m}$0 dB gap to 3-bit ESS, while reducing storage and computational complexity by factors of 6 and 3, respectively (Gültekin et al., 2019). These results are not neural, but they identify the same design axes that later neural PAS methods confront: entropy-rate efficiency, complexity, and the degree of structure imposed on amplitude sequences.

In nonlinear fiber systems, classical baselines also include sequence-selection methods that operate on shaped candidate pools instead of learning a joint distribution. One line of work introduces a linear lowpass filter model from symbol-energy sequences to nonlinear distortion and derives the LSAS metric for PAS with sequence selection, showing that nonlinear shaping gain depends on shaping blocklength, mapping strategy, and the energy-sequence spectrum relative to the channel’s lowpass NLI filter (Askari et al., 2022). Another line introduces a sign-dependent additive-multiplicative perturbative metric and reports more than $r=\frac{k}{m}$1 dB SNR gain and a $r=\frac{k}{m}$2 bits/2D AIR improvement for sequence selection over PAS without selection in single-polarized 256-QAM transmission over a 20-span, 1600 km long-haul link (Askari et al., 2024). Seq-NPAS explicitly differentiates itself from such methods by learning the distribution directly and thereby avoiding repeated metric evaluations, candidate-pool dependence, side information for candidate signaling, extra rate loss, and high computational complexity (Askari et al., 27 May 2026).

6. Receiver interaction, carrier phase recovery, and present design constraints

A recurrent theme in the PAS literature is that shaping gains in nonlinear optical channels depend strongly on the receiver, especially on carrier phase recovery. In a 12-channel, 900 GHz WDM system with a 4-channel superchannel of interest, 2 polarizations, 256-QAM, 41.67 GBd, 80 km SSMF spans, and either EDC or ideal DBP, short-block PAS implemented with ESS can provide a nonlinear shaping gain without blind phase search: in the 15-span case, the reported gain is about $r=\frac{k}{m}$3 bits/symbol with EDC and $r=\frac{k}{m}$4 bits/symbol with ideal DBP (Civelli et al., 2020). However, when BPS is properly optimized, the nonlinear shaping gain essentially disappears at sufficiently high SNR, and long-block PAS performs as well as short-block PAS because PAS and BPS mitigate the same inter-channel nonlinear phase noise from opposite ends of the link (Civelli et al., 2020).

A broader follow-up comparing ESS, shell mapping, and CCDM across several links reaches the same conclusion in stronger form. Without CPR, PAS with optimal finite block length provides a nonlinear shaping gain relative to linearly optimized PAS with infinite block length, and sphere shaping yields the largest such gain; with CPR included, the gain becomes smaller or vanishes in many realistic cases, so the reduction of rate loss obtained by sphere shaping and longer DM blocks becomes more important than nonlinear phase-noise mitigation obtained by constant-energy DMs and short blocks (Civelli et al., 2022). The same work introduces a nonlinear phase-noise metric based on the frequency-resolved logarithmic perturbation models and reports correlation $r=\frac{k}{m}$5 between SNR and $r=\frac{k}{m}$6NPN, versus $r=\frac{k}{m}$7 for EEDI, when CPR is present (Civelli et al., 2022).

PAS also changes the symbol statistics seen by the phase tracker. A Bayesian alternative to BPS, called BaPS, explicitly incorporates both the PAS symbol prior and a phase prior from pilots or Wiener phase tracking through a MAP phase search. In AWGN with Wiener phase noise and 64-QAM PAS, BaPS is consistently superior to ordinary BPS and a PAS-aware BPS variant near the optimal shaping parameter, especially at low SNR and with short windows. In an 11-channel, 32 GBd, single-polarization WDM fiber simulation with 256-QAM PAS at 2.4 bits/amplitude, 15 spans of 80 km SSMF, and EDFA noise figure 6 dB, BaPS outperforms BPS in Q-factor for all launch powers (Askari et al., 2023).

For neural PAS, these results impose a methodological constraint. The existing NPAS fiber study assumes pilot-aided linear CPR with 2.5% pilot rate (Askari et al., 2 Feb 2026), but the wider PAS literature suggests that apparent nonlinear shaping gains can be overstated if CPR is omitted or mismatched. This suggests that neural PAS, especially when evaluated in nonlinear fiber channels, should be assessed jointly with carrier phase recovery and with explicit implementation losses, rather than against transmitter-only baselines.

Neural probabilistic amplitude shaping therefore occupies a technically heterogeneous but conceptually coherent area. It includes neural generalizations of PAS for arbitrary code rates and non-Maxwell–Boltzmann distributions, PAS-compatible joint-distribution learning for nonlinear fiber channels, and rate-loss-aware sequential models that align learned dependencies with arithmetic distribution matching. Across these variants, the central shift is from optimizing a static amplitude PMF to optimizing the full operational interface among source statistics, distribution matching, channel memory, demapping, and receiver-side phase recovery (Aoudia et al., 2020, Askari et al., 2 Feb 2026, Askari et al., 27 May 2026).