Sequential Neural Probabilistic Amplitude Shaping: Learning the Channel's Language

Abstract: We present the first neural probabilistic amplitude shaping that outperforms existing methods while accounting for all implementation losses, using a block-less, easily implementable sequential autoregressive encoder compatible with arithmetic distribution matching, yielding reduced rate loss and higher achievable information rates.

• The paper introduces a rate-loss-aware training formulation that jointly optimizes bit-metric decoding, rate loss, and MB target regularization to improve PAS efficiency.
• It presents a sequential autoregressive encoder (Seq-NPAS) that eliminates block boundaries, capturing long-range temporal dependencies for enhanced ADM compatibility.
• Empirical results in dual-polarization WDM systems show that Seq-NPAS++ significantly reduces ADM rate loss and boosts AIR, outperforming traditional blockwise methods.

Sequential Neural Probabilistic Amplitude Shaping: A Rate-Loss-Aware Joint-Distribution Learning Paradigm

Introduction and Context

The paper "Sequential Neural Probabilistic Amplitude Shaping: Learning the Channel’s Language" (2605.28143) presents a significant advancement in the domain of Probabilistic Amplitude Shaping (PAS) for coherent optical transmission systems operating over nonlinear channels. Traditional PAS schemes achieve shaping gains by tailoring amplitude distributions and leveraging Forward Error Correction (FEC) for sign assignment, enabling flexible rate adaptation. However, finite-length distribution matchers such as Enumerative Sphere Shaping (ESS) introduce a performance tradeoff: while shorter shaping blocklengths improve nonlinear tolerance, they exacerbate rate loss due to suboptimal mapping efficiency.

State-of-the-art approaches, including sequence selection and neural shaping via recurrent architectures, reveal that nonlinear distortion depends not only on the symbol marginals but also on the joint distribution—specifically, on the temporal ordering of transmitted symbols. Sequence-selection strategies exploit joint structure heuristically via candidate selection and distortion prediction, but are encumbered by computational complexity and excessive rate loss due to candidate-index signaling. Prior neural shaping formulations using blockwise encoders (NPS/NPAS) have shown potential by learning optimized joint distributions within blocks, but inherently neglect dependencies across blocks and fail to explicitly control rate-loss, which can negate the information rate (AIR) gains.

Rate-Loss-Aware Formulation

The authors introduce a novel, rate-loss-aware training objective for neural probabilistic amplitude shaping. Unlike conventional approaches which optimize for information-theoretic quantities without explicit regularization, this formulation integrates:

• The bit-metric decoding rate, explicitly decomposed into contributions from marginal entropies and demapper mismatch
• An implementation rate loss term, representing the entropy gap attributable to symbol interdependency and the traversal of a distribution matcher
• A regularization term enforcing closeness of the learned symbol marginal to a Maxwell–Boltzmann (MB) target, controlled by a tunable $\lambda$ parameter

This composite objective ensures that the neural encoder's expressivity does not come at the expense of impractically high rate loss, thereby securing realizable gains in practical AIR.

Sequential Autoregressive Architecture

A key architectural innovation is the block-less, sequential autoregressive encoder (Seq-NPAS). The model predicts the next amplitude symbol conditioned on a fixed-length context of past symbols, akin to next-token prediction in large language modeling but with the constellation as the discrete alphabet. This formulation is distinguished by several operational advantages:

• Stationarity: The learned symbol process exhibits stationary statistics, as the same time-invariant prediction mechanism applies regardless of symbol position.
• Removal of Artificial Block Boundaries: Contrary to blockwise NPS/NPAS, Seq-NPAS captures temporal dependencies across arbitrarily long sequences, eliminating the non-stationarities and inefficiencies induced by block boundaries.
• Implementation Compatibility: The sequential design integrates natively with Arithmetic Distribution Matching (ADM), streamlining bit-to-symbol mapping and facilitating practical deployment.

Empirical Results

The practicality and efficacy of Seq-NPAS++ (Seq-NPAS trained with the rate-loss-aware objective) are demonstrated in a dual-polarization WDM optical system. The architecture is compared against NPAS, NPAS++, ESS, and ESS with sequence selection, with key observations including:

• Significant reduction in empirical ADM rate loss, especially under rate-loss-aware training, narrowing the gap to the theoretical lower bound as ADM input length increases (Figure 1).

  Figure 1: Empirical ADM rate loss versus ADM input length and average output length, showing the benefit of rate-aware optimization in minimizing implementation penalties.

• AIR performance: Seq-NPAS++, in both low and high launch power regimes, matches or surpasses prior approaches. In low-power regimes, classical ESS is competitive due to weak nonlinearity and lower rate loss, but as the launch power increases and nonlinear effects become prominent, NPAS++ and Seq-NPAS++ outperform ESS and sequence-selection, achieving up to 0.05 bits/2D higher AIR by leveraging learned joint distribution without incurring prohibitive rate loss (Figure 2).

  Figure 2: AIR versus launch power, illustrating the superiority of neural joint-distribution learning in nonlinear regimes.

• The gains are achieved with concise models (e.g., single-layer Transformer with rotary embeddings and GLU activations) and tractable ADM coding lengths, validating the feasibility of practical, end-to-end deployment.

Theoretical and Practical Implications

This work substantiates several important theoretical and practical insights:

• Joint-distribution learning in shaping must account for rate loss: Neglecting rate loss when introducing temporal dependencies can eliminate any net gain, as the cost of encoding dependencies can counteract improvements in nonlinear tolerance.
• Sequential architectures generalize blockwise approaches: By eliminating boundaries and exploiting temporal correlations over the true channel memory, Seq-NPAS is both more expressive and more efficient.
• Practical shaping methods must close the gap between theory and realization: Incorporating rate loss explicitly in the loss, and implementing with efficient matchers like ADM, bridges the chasm between idealized shaping gains and deployed system performance.

The presented methodology is not limited to fiber nonlinearity, but is broadly applicable to any communication system where the channel is nonlinear and exhibits memory, making it a versatile tool for modern optical and wireless systems employing high-order QAM constellations.

Future Directions

Potential avenues for expansion include:

• Extension to multi-span long-haul links, where channel memory is more pronounced and gains from cross-block modeling are greater
• Integration of more advanced channel models with full polarization and span-wise impairment accumulation
• Exploration of larger, more expressive sequential models, and adaptive context lengths matched to varying channel memory
• Application to other physical-layer systems impacted by nonlinearities, such as satellite, mmWave, or molecular communications

Conclusion

By introducing a rate-loss-aware training regime and a sequential autoregressive neural encoder, this paper establishes a robust paradigm for neural probabilistic amplitude shaping that achieves superior practical AIR and shaping efficiency for optical communications in nonlinear channels. The methodological innovations reconcile the theoretical potential of joint-distribution shaping with the implementation constraints of real-world communication systems, and set a foundation for further advancements in AI-driven physical layer optimization.