Quantization of Neural Network Equalizers in Optical Fiber Transmission
  Experiments
The quantization of neural networks for the mitigation of the nonlinear and
components' distortions in dual-polarization optical fiber transmission is
studied. Two low-complexity neural network equalizers are applied in three
16-QAM 34.4 GBaud transmission experiments with different representative
fibers. A number of post-training quantization and quantization-aware training
algorithms are compared for casting the weights and activations of the neural
network in few bits, combined with the uniform, additive power-of-two, and
companding quantization. For quantization in the large bit-width regime of
$\geq 5$ bits, the quantization-aware training with the straight-through
estimation incurs a Q-factor penalty of less than 0.5 dB compared to the
unquantized neural network. For quantization in the low bit-width regime, an
algorithm dubbed companding successive alpha-blending quantization is
suggested. This method compensates for the quantization error aggressively by
successive grouping and retraining of the parameters, as well as an incremental
transition from the floating-point representations to the quantized values
within each group. The activations can be quantized at 8 bits and the weights
on average at 1.75 bits, with a penalty of $\leq 0.5$~dB. If the activations
are quantized at 6 bits, the weights can be quantized at 3.75 bits with minimal
penalty. The computational complexity and required storage of the neural
networks are drastically reduced, typically by over 90\%. The results indicate
that low-complexity neural networks can mitigate nonlinearities in optical
fiber transmission.
Jamal Darweesh, Nelson Costa, Antonio Napoli, Bernhard Spinnler, Yves Jaouen, Mansoor Yousefi
# Quantization of Neural Network Equalizers in Optical Fiber Transmission Experiments

###### Abstract

The quantization of neural networks for the mitigation of the nonlinear and components' distortions in dual-polarization optical fiber transmission is studied. Two low-complexity neural network equalizers are applied in three 16-QAM 34.4 GBaud transmission experiments with different representative fibers. A number of post-training quantization and quantization-aware training algorithms are compared for casting the weights and activations of the neural network in few bits, combined with the uniform, additive power-of-two, and companding quantization. For quantization in the large bit-width regime of \(>5\) bits, the quantization-aware training with the straight-through estimation incurs a Q-factor penalty of less than 0.5 dB compared to the unquantized neural network. For quantization in the low bit-width regime, an algorithm dubbed companding successive alpha-blending quantization is suggested. This method compensates for the quantization error aggressively by successive grouping and retraining of the parameters, as well as an incremental transition from the floating-point representations to the quantized values within each group. The activations can be quantized at 8 bits and the weights on average at 1.75 bits, with a penalty of \(\leq 0.5\) dB. If the activations are quantized at 6 bits, the weights can be quantized at 3.75 bits with minimal penalty. The computational complexity and required storage of the neural networks are drastically reduced, typically by over 90%. The results indicate that low-complexity neural networks can mitigate nonlinearities in optical fiber transmission.

 Neural network equalization, nonlinearity mitigation, optical fiber communication, quantization.

## I Introduction

The compensation of the channel impairments is essential to the spectrally-efficient optical fiber transmission. The advent of the coherent receivers, combined with the advances in the digital signal processing (DSP) algorithms, has allowed for the mitigation of the fiber transmission effects in the electrical domain [1]. However, real-time energy-efficient DSP is challenging in high-speed communication.

The linear transmission effects, such as the chromatic dispersion (CD) and polarization mode dispersion (PMD), can be compensated using the well-established DSP algorithms [2]. The distortions arising from the fiber Kerr nonlinearity can in principle be partially compensated using the digital back propagation (DBP) based on the split-step Fourier method (SSFM). DBP can be computationally complex in long-haul transmission with large number of steps in distance [3]. The neural networks (NNs) provide an alternative approach to nonlinearity mitigation with flexible performance-complexity trade-off [4, 5, 6, 7, 8]; see Section III-A.

To implement NNs for real-time equalization, the model should be carefully optimized for the hardware. The number of bits required to represent the NN can be minimized by quantization [9] and data compression, using techniques such as pruning, weight sharing and clustering [10]. There is a significant literature showing that these methods often drastically reduce the storage requirement of the NN, and its energy consumption, which is often dominated by the communication cost of fetching words from the memory to the arithmetic units [10, 11, 12]. How the NNs can be quantized with as few bits as possible, while maintaining a given Q-factor, is an important problem. This paper is dedicated to the quantization of the NNs for nonlinearity mitigation, in order to reduce the computational complexity, memory footprint, latency and energy consumption of the DSP.

There are generally two approaches to the NN quantization. In post-training quantization (PTQ), the model is trained in 32- or 16-bit floating-point (FP) precision, and the resulting parameters are then quantized with fewer number of bits [9, 13]. This approach is simple; however, quantization introduces a perturbation to the model parameters incurring a performance penalty. As a consequence, PTQ is usually applied in applications that do not require quantization below 8 bits.

In quantization-aware training (QAT), quantization is integrated into the training algorithm, and the quantization error is partly compensated [11, 14, 15, 16, 12]. However the optimization of the loss function with gradient-based methods is not directly possible, because the quantizer has a derivative that is zero almost everywhere. In the straight-through estimator (STE), the quantizer is assumed to be the identity function, potentially saturated in an input interval, in the backpropagation algorithm used for computing the gradient of the loss function [17, 18]. QAT is used in applications requiring low complexity in inference; however, it can be more complex in training than PTQ, and needs parameter tuning and experimentation. With the exception of a few papers reviewed in Section IV-F, the quantization of the NNs for nonlinearity mitigation has not been much explored.
In this paper, we study the quantization of the weights and activations of a small convolutional fully-connected (Conv-FC) and a bidirectional long short-term memory fully-connected (BiLSTM-FC) equalizer, applied to three 16-QAM 34.4 GBaud dual-polarization fiber transmission experiments. The experiments are based on a 9x50 km true-wave classic (TWC) fiber link, a 9x110 km standard single-mode fiber (SMF) link, and a 17x70 km large effective area fiber (LEAF) link. We compare the Q-factor penalty, computational complexity, and memory requirement of a number of PTQ and QAT-STE algorithms, as a function of the launch power and the quantization rate \(b\). The uniform, additive power-of-two (APoT), companding, fixed- and mixed-precision quantization are compared. It is shown that, these algorithms, if optimized, work well in the large bit-width regime of \(b\geq 5\). However, they do not achieve sufficiently small distortions in our experiments in the low bit-width regime with \(b<5\), where the quantization error needs to be aggressively mitigated. For this case, we propose a companding successive alpha-blending (SAB) quantization algorithm that mitigates the quantization error by successive grouping and retraining of the parameters, combined with an incremental transition from the floating-point representations to the quantized values within each group. The algorithm also accounts for the probability distribution of the parameters. It is shown that the quantization of the activations impacts the Q-factor much more than the weights. The companding SAB algorithm is studied w/wo the quantization of activations.

The results indicate that, for quantization in the large bit-with regime, QAT-STE incurs a Q-factor penalty of less than 0.5 dB relative to the unquantized NN, while reducing the storage and computational complexity of the NN typically by over 90%. This is obtained with the uniform, companding or APoT variant of QAT-STE, depending on the transmission experiment. If the activations are quantized at 8 bits, the weights can be quantized with the companding SAB algorithm at the average rate of 1.75 bits, paving the way to the binary NN equalizers. The quantization of the activations at 6 bits and weights at \(3.75\) bits results in a reduction in the computational complexity by \(95\%\) and memory footprint by \(88\%\), with the Q-factor penalty of 0.2 dB. Overall, the results suggest that nearly-binary NNs mitigate nonlinearities in optical fiber transmission.

This paper is structured as follows. In Section II, we describe the optical fiber transmission experiments. In Section III, we review the use of the NNs for the fiber nonlinearity mitigation, and in Section IV the quantization of the NNs. Finally, we compare the Q-factor penalty and the gains of quantization for several algorithms in Section V, and draw conclusions in Section VI.

## II Dual Polarization Transmission Experiment Setup

Fig. 1 shows the block diagram of the transmission experiments considered in this paper. Three experiments are performed with different representative fibers, described below.

#### Ii-1 Transmitter

At the transmitter (TX), a pseudo-random bit sequence (PRBS) is generated for each polarization \(p\in\{x,y\}\), and mapped to a sequence of symbols \(\mathbf{s}_{p}\) taking values in a 16-QAM constellation according to the Gray mapping. The two complex-valued sequences \(\mathbf{s}_{x}\) and \(\mathbf{s}_{y}\) are converted to four real-valued sequences, and passed to an arbitrary wave generator (AWG) that modulates them to two QAM signals using a root raised cosine pulse shape with the roll-off factor of 0.1 at the rate \(34.4\) GBaud. The AWG includes digital-to-analog converters (DACs) at \(88\) Gsamples/s.

The outputs of AWG are four continuous-time electrical signals \(I_{x}\), \(Q_{x}\), \(I_{y}\) and \(Q_{y}\) corresponding to the in-phase (I) and quadrature (Q) components of the signals of the \(x\) and \(y\) polarization. The electrical signals are converted to optical signals and polarization-multiplexed with a dual-pol IQ Mach-Zehnder modulator (MZM), driven by an external cavity laser (ECL) at wavelength \(1.55~{}\mu m\) with line-width 100 KHz. The output of the IQ-modulator is amplified by an erbium-doped fiber amplifier (EDFA), filtered by an optical band-pass filter (OBPF) and launched into the fiber link. The laser introduces phase noise, modeled by a Wiener process with the Lorentzian power spectral density [19, Chap. 3.5].

#### Ii-2 Fiber-optic Link

The channel is a straight-line optical fiber link in a lab, with \(N_{sp}\) spans of length \(L_{sp}\). An EDFA with 5 dB noise figure (NF) is placed at the end of each span to compensate for the fiber loss. The experiments are performed

\begin{table}
\begin{tabular}{l c c c}  & TWC fiber & SMF & LEAF \\ \cline{2-4} \(L_{sp}~{}\mathrm{km}\) & 50 & 110 & 70 \\ \(N_{sp}\) & 9 & 9 & 17 \\ \(\alpha~{}\mathrm{dB/km}\) & 0.21 & 0.22 & 0.19 \\ \(D~{}\mathrm{ps/(mm.km)}\) & 5.5 & 18 & 4 \\ \(\gamma~{}(\mathrm{W.Km})^{-1}\) & 2.8 & 1.4 & 2.1 \\ PMD \(\tau~{}\mathrm{ps/\sqrt{km}}\) & 0.02 & 0.08 & 0.04 \\ NF \(\mathrm{dB}\) & 5 & 5 & 5 \\ \end{tabular}
\end{table} TABLE I: OPTICAL LINK PARAMETERS

Fig. 1: The block-diagram of the transmission experiments.
with the TWC fiber, SMF and LEAF, and parameters in Table I.

TWC Fiber ExperimentThe first experiment is with a short-haul TWC fiber link with 9 spans of 50 km. The TWC fiber was a brand of nonzero dispersion shifted fiber (NZ-DSF) made by Lucent, with low CD coefficient of \(D=5.5~{}\mathrm{ps}/(\mathrm{nm}\cdot\mathrm{km})\) at 1550 nm wavelength and a high nonlinearity parameter of \(\gamma=2.8~{}(\mathrm{Watt}\cdot\mathrm{km})^{-1}\). Thus, even though the link is short with 450 km length, the channel operates in the nonlinear regime at high powers. The link parameters, including the fiber loss coefficient \(\alpha\) and PMD value \(\tau\), can be found in Table I.

SMF ExperimentThe second experiment is based on a long-haul 9x110 km standard single-mode fiber link, with parameters in Table I.

LEAF ExperimentLEAF is also a brand of NZ-DSF, made by Corning, similar to the TWC fiber but with a smaller nonlinearity coefficient due to the larger cross-section effective area. This experiment uses a 17x70 km link described in Table I.

#### Ii-B3 Receiver

At the receiver, the optical signal is polarization demultiplexed, and converted to four electrical signals using an integrated coherent receiver driven by a local oscillator (LO). Next, the continuous-time electrical signals are converted to the discrete-time signals by an oscilloscope, which includes analog-to-digital converters (ADCs) that sample the signals at the rate of \(50\) Gsamples/s, and quantize them with the effective number of bits of around \(5\). The digital signals are up-sampled at 2 samples/symbol, and equalized in the DSP chain shown in Fig. 1.

The equalization is performed by the conventional dual-polarization linear DSP [1], followed by a NN. The linear DSP consists of a cascade of the frequency-domain CD compensation, multiple-input multiple-output (MIMO) equalization via the radius directed equalizer to compensate for PMD [1, Sec. VII-], [20], polarization separation, carrier frequency offset (CFO) correction, and the carrier-phase estimation (CPE) using the two-stage algorithm of Pfau _et al._ to compensate for the phase offset [21]. The linearly-equalized symbols are denoted by \(\tilde{\mathbf{s}}_{p}\).

Once the linear DSP is applied, the symbols are still subject to the residual CD, dual-polarization nonlinearities, and the distortions introduced by the components at TX and RX. Define the residual channel memory \(M\) to be the maximum effective length of the auto-correlation function of \(\tilde{\mathbf{s}}_{p}\) over \(p\in\{x,y\}\).

The outputs of the CPE block \(\tilde{\mathbf{s}}_{p}\) are passed to a low-complexity NN, which mitigates the remaining distortions, and outputs \(\tilde{\mathbf{s}}_{p}\). The architecture of the NN depends on the experiment, and will be explained in Section III-B.

## III Neural Networks for Nonlinearity Mitigation

### _Prior Work_

The NN equalizers in optical fiber communication can be classified into two categories. In _model-based equalizers_, the architecture is based on the parameterization of the channel model. An example is learned DBP (LDBP) [8], where the NN is a parameterization of the SSFM which is often used to simulate the fiber channel. The dual-polarization LDBP is a cascade of layers, each consisting of two complex-valued symmetric filters to compensate for the CD, two real-valued asymmetric filters for the differential group delays, a unitary matrix for the polarization rotation, and a Kerr activation function for the mitigation of the fiber nonlinearity. It is shown that LDBP outperforms DBP [8].

On the other hand, in _model-agnostic equalizers_, the architecture is independent of the channel model [4, 5, 6, 7]. The model-agnostic schemes do not require the channel state information, such as the fiber parameters. Here, the NNs can be placed at the end of the conventional linear DSP for nonlinearity mitigation [22], or after the ADCs for compensating the linear and nonlinear distortions (thereby replacing the linear DSP) [23, 24].

A number of NN architectures have been proposed for the nonlinearity mitigation. Fully-connected (FC) or dense NNs with 2 or 3 layers, few hundred neurons per layer, and tanh activation were studied in [25, 26]. The overfitting and complexity become problems when the models get bigger. The convolutional NNs can model the linear time-invariant (LTI) systems with a finite impulse response. The application of the convolutional networks for compensating the nonlinear distortions is investigated in [27], showing that one-dimensional convolution can well compensate the CD. The bi-directional recurrent and long-short term memory networks (LSTM) receivers are shown to perform well in fiber-optic equalization [24]. Compared to the convolutional and dense networks, BiLSTM networks better model LTI systems with infinite impulse response, such as the response of the CD. A comparison of the different architectures in optical transmission in [25] shows that, dense and convolutional-LSTM models perform well at low and high complexities, respectively.

An effect that particularly impacts the performance of the NN is PMD. In most papers, random variation of the polarization-dependent effects during the transmission have not been carefully studied. The polarization effects are sometimes neglected [22], or assumed to be static during the transmission [8]. In such simulated systems, the dual-polarization NN receivers are subject to a performance degradation compared to real-life experiments [25].

### _Two NN Models Considered in This Paper_

In this Section, we describe two NN equalizers used in this paper. The NN is placed at the end of the linear DSP shown in Fig. 1. In consequence, since the PMD is compensated by the MIMO equalizer, the NN is static and trained offline. Due to the constrains of the practical systems, low-complexity architectures are considered. A Conv-FC network is applied in the TWC fiber and SMF links, and a BiLSTM-FC network in the LEAF link. The BiLSTM-FC model has more parameters, and performs better; however, the smaller Conv-FC model is sufficient in short-haul links.
#### Ii-B1 Conv-FC Model

The four sequences of linearly-equalized symbols \(\Re(\tilde{\mathbf{s}}_{x})\), \(\Im(\tilde{\mathbf{s}}_{x})\), \(\Re(\tilde{\mathbf{s}}_{y})\) and \(\Im(\tilde{\mathbf{s}}_{y})\) are passed to the NN. We consider a many-to-one architecture, where the NN equalizes one complex symbol per polarization given \(n_{i}\) input symbols. The inputs of the network are four vectors, each containing a window of \(n_{i}=M+1\) consecutive elements from each of the four input sequences, where \(M\) is the residual channel memory defined in Section II-3. The network outputs a vector of \(n_{o}=4\) real numbers, corresponding to the real and imaginary parts of the symbols of the two polarizations after full equalization. The size of the concatenated input of the NN is thus \(\bar{n}_{i}=4(M+1)\). The NN operates in a sliding-window fashion: as each of its input vectors are shifted forward one element, \(4\) real numbers are produced.

The Conv-FC model is a cascade of a complex-valued convolutional layer, a FC hidden layer, and a FC output layer. The first layer implements the discrete convolution of \(\tilde{\mathbf{s}}_{p}\), \(p\in\{x,y\}\), with a kernel \(\mathbf{h}\in\mathbb{C}^{K}\), to compensate primarily the residual CD, where \(\mathbb{C}\) denotes the complex numbers and \(K\) is the number of kernel taps. The two complex convolutions \(\tilde{\mathbf{s}}_{p}\ast\mathbf{h}\) are implemented using eight real convolutions in terms of two filters \(\Re(\mathbf{h})\) and \(\Im(\mathbf{h})\), according to

\[\tilde{\mathbf{s}}_{p}\ast\mathbf{h} =\Re(\tilde{\mathbf{s}}_{p})\ast\Re(\mathbf{h})-\Im(\tilde{ \mathbf{s}}_{p})\ast\Im(\mathbf{h})\] \[+j\Big{\{}\Re(\tilde{\mathbf{s}}_{p})\ast\Im(\mathbf{h})+\Im( \tilde{\mathbf{s}}_{p})\ast\Re(\mathbf{h})\Big{\}}. \tag{1}\]

The first layer thus contains eight parallel real-valued one-dimensional convolutions, with the stride one and "same padding," and no activation. There are total \(2K\) trainable real filter taps, typically far fewer than in generic convolutional layers used in the literature with large feature maps. The eight real convolutions are combined according to (1) or Fig. 2(a), obtaining \(\Re(\tilde{\mathbf{s}}_{x}\ast\mathbf{h})\), \(\Im(\tilde{\mathbf{s}}_{x}\ast\mathbf{h})\), \(\Re(\tilde{\mathbf{s}}_{y}\ast\mathbf{h})\) and \(\Im(\tilde{\mathbf{s}}_{y}\ast\mathbf{h})\), which are then concatenated. The resulting vector is fed to a FC hidden layer with \(n_{h}\) neurons, and tangent hyperbolic (tanh) activation. The joint processing of the two polarizations in the dense layer is necessary in order to compensate the nonlinear interactions between the two polarizations during the propagation. Finally, there is an output FC layer with \(2\) neurons for each complex-valued polarization symbol, and no activation.

The computational complexity \(\mathcal{C}\) of the unquantized NNs can be measured by the number of the real multiplications per polarization, considering that the cost of the additions and computation of the activation is comparatively negligible. For the Conv-FC model

\[\mathcal{C}_{\text{Conv-FC}}=4n_{i}K+2n_{i}n_{h}+\frac{n_{h}n_{o}}{2}. \tag{2}\]

#### Ii-B2 BiLSTM-FC Model

The second model is a cascade of a concatenator, a BiLSTM unit and FC output layer, shown in Fig. 2(b). At each time step \(t\) in the recurrent model, \(n_{i}=M+1\) linearly-equalized complex symbols are taken from each polarization. The resulting vectors \(\Re(\tilde{\mathbf{s}}_{x}^{(t)})\), \(\Im(\tilde{\mathbf{s}}_{x}^{(t)})\), \(\Re(\tilde{\mathbf{s}}_{y}^{(t)})\), \(\Im(\tilde{\mathbf{s}}_{y}^{(t)})\), are concatenated in a vector of length \(\bar{n}_{i}=4(M+1)\) and fed to a many-to-many BiLSTM unit. Each LSTM cell in this unit has an input of length \(2(M+1)\) corresponding to the one-sided memory, \(n_{h}\) hidden state neurons, the recurrent activation \(\tanh\), and the gate activation sigmoid. The output of the BiLSTM unit is a vector of length \(2n_{h}\), that is fed to a FC output layer with no activation and \(n_{o}=4\) neurons1. The computational complexity of the BiLSTM-FC model is

Footnote 1: Equivalently, the input output of the BiLSTM unit may be expressed in arrays of shape \((4,M+1)\), without concatenation.

\[C_{\text{BiLSTM-FC}}=n_{h}\Big{(}4n_{h}+16n_{i}+3+n_{o}\Big{)},\]

real multiplications per polarization.

The many-to-many variants of the above models are straightforward. In this case, there are \(n_{o}=4(M+1)\) neurons at the output, so that all \(M+1\) complex symbols are equalized in one shot; thus \(n_{i}=M+L\), \(\bar{n}_{i}=n_{o}=4(M+L)\). The many-to-many versions are less complex per symbol and parallelizable, but also less performant.

The performance of the receiver is measured in terms of

\[\text{Q-factor}=10\log_{10}\!\left(2\operatorname{erfc}^{-2}(2\text{BER}) \right)\quad\text{dB},\]

where the BER is the bit error rate, and \(\operatorname{erfc}(.)\) is the complementary error function. The Q-factor of the NNs is
compared with that of DBP and linear equalization. The DBP replaces the CD compensation unit at the beginning of the DSP chain and is applied with single step per span, and 2 samples per symbol. This comparison is done to evaluate the effectiveness of the NN in jointly mitigating the residual CD and Kerr nonlinearity.

Fig. 3(a) shows the Q-factor gain of the unquantized Conv-FC model over the linear DSP in the TWC fiber experiment (\(K=M=40\)) [28]. The results demonstrates that the NN offers a Q-factor enhancement of \(0.5\) dB at -2 dBm, and \(2.3\) dB at 2 dBm. The raw data before the linear DSP were not available to add the DBP curve to Fig. 3(a). The TWC fiber link is short. On the other hand, the nonlinearities are stronger in the fiber link in the SMF experiment than in the TWC fiber experiment, due to the longer length. For the SMF experiment, Fig. 3(b) shows that the Conv-FC model provides a performance similar to that of DBP with 1 sample/symbol (SpS). The improvement results from the mitigation of the dual-polarization nonlinearities, as well as the equipment's distortions. The BiLSTM based receiver in the LEAF experiment (with \(n_{h}=100\), \(M=40\)) also gives a comparable performance to the DBP as shown in Fig. 3(c).

In general, the implementation of the NN can be computationally expensive. In order to reduce the complexity, in the next section, we quantize the NNs, casting the weights and activations into low precision numbers.

## IV Quantization of the Neural Networks

The parameters (weights and biases) of the NN, activations and input data are initially real numbers represented in FP 32 (FP32) or 64 bit numbers, described, e.g., in the IEEE 754 standards. The implementation of the NNs in memory or computationally restricted environments requires that these numbers to represented by fewer number of bits and in different format, e.g., in INT8.

Define the quantization grid \(\mathcal{W}\) as a finite set of numbers

\[\mathcal{W}=\big{\{}\hat{w}_{0},\hat{w}_{1},\cdots,\hat{w}_{n}\big{\}},\]

where \(\hat{w}_{i}\in\mathbb{R}\) are the quantization symbols. A continuous random variable \(w\in\mathbb{R}\) drawn from a probability distribution \(p(w)\) is quantized to \(\hat{w}=Q(w)\), where \(Q:\mathbb{R}\mapsto\mathcal{W}\) is the quantization rule or quantizer

\[Q(w)=\sum_{i=0}^{N}\hat{w}_{i}\;\mathbbm{1}_{I_{i}}(w).\]

Here, \(I_{i}=[\Delta_{i},\Delta_{i+1})\), where \(\{\Delta_{i}\}_{i=0}^{N+1}\) are the quantization thresholds, and \(\mathbbm{1}\) is the indicator function, i.e., \(\mathbbm{1}_{I_{i}}(w)=1\) if \(w\in I_{i}\), and \(\mathbbm{1}_{I_{i}}(w)=0\) otherwise. The intervals \(\{I_{i}\}_{i=0}^{N}\) are the quantization cells, partitioning the real line. The quantization rate of \(\mathcal{W}\) is \(b=\log_{2}(N+1)\) bits, assuming that \(\hat{w}_{i}\) are equally likely. The hardware support is best when \(b\) is a power of two, commonly \(b=8\).

The quality of reproduction is measured by a distortion which is often the mean-square error (MSE) \(D(b)=\mathbb{E}(w-\hat{w})^{2}\), where the expectation \(\mathbb{E}\) is with respect to the probability distribution of \(w\) and \(Q\) (if it includes random elements). For a fixed rate \(b\), the symbols \(\hat{w}_{i}\) and \(\Delta_{i}\) (or \(Q(.)\)) are found to minimize the distortion \(D(b)\).

### _Quantization Schemes_

There is a significant literature on the quantization algorithms in deep learning. However, most of these algorithms have been developed for over-parameterized NNs with large number of parameters. These networks have many degrees-of-freedom to compensate for the quantization error. It has been experimentally demonstrated that the over-parameterized NNs are rather resilient to the quantization, at least up to 8 bits. In contrast, the NNs used for fiber equalization are small, typically with few hundred or thousands of weights, smaller than the models deployed even in smartphones and Internet of Things applications [29]. Below, we review a number of the quantization algorithms suitable for the NN equalizers.

#### Iv-A1 Uniform Quantization

In uniform quantization, the quantization symbols \(\hat{w}_{i}\) are uniformly placed. Given a step size (or scale factor) \(s\) and a zero point \(z\), the uniform quantization rule is

\[\hat{w}=s(\bar{w}-z),\]

Fig. 3: Q-factor of the linear DSP, DBP with 1 SpS, and unquantized NN equalizers in (a) TWC fiber, (b) SMF, and (c) LEAF experiments.
where \(\bar{w}\in\bar{\mathcal{W}}=\{0,1,\cdots,N\}\). The integer representation of \(w\) is

\[\bar{w}=\text{clip}\left(\left\lfloor\frac{w}{s}\right\rceil+z;0,N\right),\]

where \(\text{clip}(w,a,b)\), \(a\leq b\), is the clipping function

\[\text{clip}(w,a,b)=\begin{cases}a,&w<a,\\ w,&a\leq w<b,\\ b,&w\geq b,\end{cases}\]

in which \(\left\lfloor x\right\rceil\) is the rounding function, mapping \(x\) to an integer in \(\bar{\mathcal{W}}\), e.g., to the nearest symbol. The quantization grid is thus

\[\mathcal{W}_{u}(s,z,b)=\Big{\{}-zs,-sz+s,\cdots,-sz+sN\Big{\}}. \tag{3}\]

The scale factor \(s\) and zero point \(z\) can be determined by considering an interval \([\alpha,\beta]\) that contains most of the weights. Then, \(s(a,c,N){=}(\beta-\alpha)/N\) and \(z=\left\lfloor-\alpha/s\right\rfloor\). The interval \([\alpha,\beta]\) is called the clipping (or clamping or dynamic) range, and is selected by a procedure called calibration, which may require a calibration dataset (a small set of unlabeled examples). The parameters of the uniform quantizer are thus \(\alpha\), \(\beta\), \(b\) and the choice of the rounding function.

For a fixed rate \(b\), the remaining parameters can be obtained by minimizing the MSE. However, it is simpler, and sometimes about equally good (especially when \(b\geq 4\)), to set the clipping range to be an interval centered at the mean \(\mu\) of \(w\), with a duration proportional to the standard deviation \(\sigma\) of \(w\)

\[\alpha=\mu-\kappa\sigma,\quad\beta=\mu+\kappa\sigma,\]

where, e.g., \(\kappa=4\). Even a simpler method of calibration is setting \(\alpha\) and \(\beta\) to be the minimum and maximum value of the weights \(w\), respectively [12]. The min-max choice can be sensitive to the outlier parameter values, increasing unnecessarily the step size and rounding error.

In the symmetric quantization, \(z=0\). Thus, \(w=0\) is mapped to \(\bar{w}=0\) and \(\hat{w}=0\). The grid of the uniform unsigned symmetric quantization is thus \(\mathcal{W}_{\text{uns}}(s)=\big{\{}0,s,\cdots,sN\big{\}}\). If the distribution of \(w\) is symmetric around the origin, symmetric signed quantization is applied, where

\[\mathcal{W}_{\text{uns}}(s,b)=\Big{\{}ks:\quad k=-(N+1)/2,\cdots,(N-1)/2\Big{\}}. \tag{4}\]

The common practice is to cast the weights with the signed symmetric quantization. However, the output of the rectified linear unit and sigmoid activation is not symmetric. Moreover, the empirical distribution of the weights can sometimes be asymmetric. For instance, Fig. 4 shows the weight distribution of a NN used in Section V. It can be seen that the distribution has a negative mean. In these cases, asymmetric, or unsigned symmetric, quantization is used.

The quantization is said to be static if \(\alpha\) and \(\beta\) are known and hard-coded a priori in hardware. The same values are used in training and inference, and for any input. In contrast, in dynamic-range quantization, \(\alpha\) and \(\beta\) are computed in real-time for each batch of the inputs to the NN. Since activations depend on input, their clipping range is best determined dynamically. This approach requires real-time computation of the statistics of the activations, bringing about an overhead in computational and implementation complexity, and memory.

The computation composed of the addition and multiplication of the numbers in \(\mathcal{W}_{u}\) can be performed with integer arithmetic, with the scale factor and zero point applied in FP32 at the end. In what follows, the notation UN-\(b\) is used to indicate uniform quantization of the weights and activations at \(b\) bits (with a similar notation for other quantizers).

#### Iii-B2 Additive Power-of-two Quantization

In non-uniform quantization, the quantization symbols are not uniformly placed. The hardware support for these schemes is generally limited, due to, e.g., the requirements of the iterative clustering (e.g., via \(k\)-means) [30]. Thus, the majority of studies adopt uniform quantization. On the other hand, the empirical probability distribution of the weights is usually near bell shaped [31]; see Fig. 4. Thus, logarithmic quantization [32, 33, 34] could provide lower rate for a given distortion compared to the uniform quantization.

In the power-of-two (PoT) quantization, the quantization symbols are powers of two [32]

\[\mathcal{W}_{\text{par}}(s,r,b)=\pm s\Big{\{}0,2^{0},2^{-r},\cdots,2^{-r(2^{ \text{\tiny{b-1}}}-1)}\Big{\}},\]

where \(r\in\mathbb{N}\) controls the width of the distribution of symbols, and \(s\in\mathbb{R}\) is the scale factor. The scale factor is stored in FP32, but is applied after the multiply-accumulate operations, and can be trainable. The PoT simplifies the computation by performing the multiplications via bit shifts. However, PoT is

Fig. 4: a) Probability density function (PDF) of the weights is bell-shaped with non-zero mean, suggesting that uniform quantization is not optimal. b) APoT-4, illustrating that the quantization symbols are irregularly placed; c) CP-3.
not flexible in the above form, and the symbols are sharply concentrated around zero. Further, increasing the bit-width merely sub-divides the smallest quantization cell around zero, without generating new symbols in other cells.

The APoT introduces additional adjustable parameters, that can be used to control the distribution of the symbols, introducing new symbols generally everywhere [33]. The APoT grid is the sum of \(n\) PoT grids with a base bit-width \(b_{0}\) and different ranges, for a given \(n\in\mathbb{N}\) and \(b_{0}\). The bit-width is thus \(b=nb_{0}\). Choosing \(b_{0}\) such that \(n=b/b_{0}\) is an integer, the quantization grid of APoT is

\[\mathcal{W}_{\text{apot}}(s,r,b,b_{0},\gamma)=\pm\,s\sum_{i=0}^{n-1}2^{-i} \big{|}\mathcal{W}_{\text{pot}}\big{|}(1,n,b_{0}+1)+\gamma,\]

where \(s\) and \(\gamma\) are trainable scale and shift factors in FP32, the absolute value in the set \(|\mathcal{W}|\) is defined per component, and \(\Sigma\) is the Minkowski set sum. It can be verified that \(|\mathcal{W}_{\text{apot}}|=2^{b}\). The shift parameter \(\gamma\) allows restricting the quantized weights to unsigned numbers.

As with the PoT, the main advantage of APoT representation is that it is multiplier-free, thus considerably less complex than the uniform quantization. The PoT and APoT gives rise to more efficient quantizers such as in DeepShift, where the bit-shifts or exponents are learned directly via STE [34]. The use of APoT in fiber-optics equalization is discussed in [28].

### _Companding Quantization_

In companding (CP) quantization, an appropriate nonlinear transformation is applied to the weights so that the distribution of the weights becomes closer to a uniform distribution, and a uniform quantizer can be applied afterwards [35]. A companding quantizer is composed of a compressor, a uniform quantizer, and an expander. The \(\mu\)-law is an example of a compressor

\[w_{c}=F(w)=\operatorname{sign}(w)\frac{\log(1+\mu|w|)}{\log(1+\mu)}, \tag{5}\]

where \(\mu>0\) is the compression factor. Its inverse

\[w=\mu^{-1}\operatorname{sign}(w_{c})\Big{(}1+\mu)^{|w_{c}|}-1\Big{)}, \tag{6}\]

is the expander.

Companding quantization has been widely used in data compression and digital communication. It is shown that the logarithmic companding quantization can cast the weights and biases of the NN image classifiers at 2 bits [36], and outperforms the uniform and APoT quantization in the same task [37]. However, the use of companding quantization in NN equalizers has not been investigated.

### _Mixed-precision Quantization_

The majority of the quantization schemes consider fixed-precision quantization, where a global bit-width is predefined. In the mixed-precision quantization, different groups of weights or activations are quantized generally at different rates [38]. The groups could be defined by layers, channels, feature maps, clusters, etc. One approach to determine the bit-width of each group is based on the sensitivity of the model using the Hessian matrix of the loss function [39]. If the Hessian matrix has a large norm on average over a particular group, a larger bit-width is assigned to that group. The output (and sometimes input) layer is often quantized at high precision, e.g., at 16 bits, as it directly influences the prediction. The biases impart a small overhead and usually not quantized. In our work, the quantization rates are determined from the sensitivity of the loss function. The hardware support for mixed-precision quantization is limited compared to the fixed-precision quantization.

### _PTQ and QAT_

#### Iii-D1 Post-training Quantization

In PTQ, training is performed in full or half precision. The input tensor, activation outputs, and the weights are then quantized at fewer bits and used in inference [40]. In practice, the quantized values are stored in integer or fixed-point representations in field-programmable gate array (FPGA) or application-specific integrated circuit (ASIC), and processed in arithmetic logic units with bit-wise operations. However, the general-purpose processors include the FP processing units as well, where the numbers are stored and processed in FP formats. Thus, to simulate PTQ in general-purpose hardware, the quantizer \(Q(.)\) is introduced in the computational graph of the NN after each weight, bias and activation stored in FP.

The PTQ has little overhead, and is useful in applications where the calibration data are not available. However, quantization below 4-8 bits can cause a significant performance degradation [41]. Several approaches have been proposed to recover the accuracy in the low bit-width regimes. Effort has been dedicated to finding a smaller clipping range from the distribution of the weights, the layer- and channel-wise mixed precision, and the correction of the statistical bias in the quantized parameters. Moreover, rounding a real number to the nearest quantization symbol may not be optimal [42]. In adaptive rounding, a real number is rounded to the left or right symbol based on a Bernoulli distribution, or deterministic optimization. It has been shown that PTQ-4 with adaptive rounding incurs a small loss in accuracy in some applications [43].

#### Iii-D2 Quantization-aware Training

In QAT, quantization is co-developed with the training algorithm. This usually enhances the prediction accuracy of the model by accounting for the quantization error during the training.

QAT is simulated by placing the quantizer function after each weight and activation in the computational graph of the NN. The output of the quantizer is a piece-wise constant function of its input. This function is not differentiable at the points of discontinuity, and has a derivative that is zero everywhere else, i.e., \(Q^{\prime}(w)=\partial\hat{w}/\partial w=0\). Thus, the gradient of the loss function with respect to the weights is zero almost everywhere, and learning with the gradient-based methods is not directly possible. There are a number of approaches to address the zero gradient problem, such as approximating \(Q^{\prime}(w)\) with a non-zero function, as in STE.

QAT usually achieves higher prediction accuracy than PTQ when quantizing at low number of bits, at the cost of the
increased overhead. On the other hand, if the approximation technique is not carefully chosen, QAT may perform even worse than PTQ [44]. Training can be performed from scratch, or from a pre-trained model, followed by QAT fine-tuning the result.

The Straight-thorough EstimationIn STE, the derivative of the quantizer is approximated with the identity function, potentially truncated on the clipping range \([\alpha,\beta]\)

\[Q^{\prime}(w)\approx\begin{cases}0,&w<\alpha,\\ 1,&\alpha\leq w<\beta,\\ 0,&w\geq\beta.\end{cases} \tag{7}\]

During the NN training, in the forward pass \(Q(.)\) is used. In the backward pass, \(Q^{\prime}(.)\) in (7) is applied, which is then used in the chain rule to back-propagate the errors in training [41, 18]. Moreover, the weights remains in FP in the backward pass, to recover the accuracy lost in the forward pass. Even though (7) is not a good approximation to the zero, STE works surprisingly well in some models when \(b\geq 5\)[44]. The gradient is usually sensitive to quantization, even more than activations. It is thus either not quantized, or quantized with at least 6 bits [45].

There are non-STE approaches as well. For instance, an appropriate regularization term can be added to the loss function that penalizes the weights that take on values outside the quantization set. Another approach is the alpha-blending (AB) quantization.

Alpha-blending QuantizationThe AB quantization addresses the problem of the quantizer's zero derivative by replacing each weight with a convex combination of the full precision weight \(w\in\mathbb{R}\) and its quantized version \(\tilde{w}=Q(w)\)[46]:

\[\tilde{w}=(1-\alpha_{j})w+\alpha_{j}\hat{w}, \tag{8}\]

where the coefficient \(\alpha_{j}\) is changed from \(0\) to \(1\) with the epoch index \(j\in\{k_{1},\cdots,k_{2}\}\) according to

\[\alpha_{j}=\begin{cases}0,&j\leq k_{1},\\ \left(\frac{k_{1}-j}{k_{2}-k_{1}}\right)^{3},&k_{1}<j\leq k_{2},\\ 1,&j\geq k_{2},\end{cases} \tag{9}\]

for some \(k_{1}\leq k_{2}\). This approach enables a smooth transition from the unquantized weights corresponding to \(\alpha_{k_{1}}=0\) to the quantized ones corresponding to \(\alpha_{k_{2}}=1\). The AB quantization is integrated into the computational graph of the NN, by placing the sub-graph shown in Fig. 5 at the end of each scalar weight.

Considering \(Q^{\prime}(.)=0\), we have \(\partial\tilde{w}/\partial w=1-\alpha\), and \(\partial L(\tilde{w})/\partial w=L^{\prime}\left(\tilde{w}\right)(1-\alpha)\neq 0\). Thus, even though the quantizer has zero derivative, the derivative of the loss function with respect to \(w\) is non-zero, and the weights are updated in the gradient-based training. The activations can still be quantized with STE.

The AB QAT starts with \(j=k_{1}\), and trains with one or more epochs. Then, \(j\) is incremented to \(k_{1}+1\), and the training continues, initialized with the weights obtained at \(j=k_{1}\). It has been shown that the AB quantization provides an improvement over QAT-STE in different scenarios [46].

Given a base quantizer \(Q(.)\), the AB quantization may be viewed as using the quantizer \(Q_{ab}(w)=(1-\alpha_{j})w+\alpha_{j}Q(w)\). As shown in Fig. 5(b), when \(Q(.)\) is the uniform quantizer, \(Q_{ab}(.)\) is a piece-wise linear approximation to \(Q(.)\), with slope \(1-\alpha_{j}\). As \(\alpha_{j}\to 1\), the approximation error tends to zero, and \(w\) is quantized.

#### Iii-B3 Successive Post-training Quantization

Successive PTQ (SPTQ) may be viewed as a combination of PTQ and QAT [47], and is particularly effective for quantizing small NNs such as those encountered in optical fiber communication as discussed in [48]. The idea is to compensate for the quantization error in the training. The parameters of the NN are partitioned into several sets and sequentially quantized based on a PTQ scheme. This approach is simple and tends to perform well in practice, with a good PTQ scheme and hyper-parameter optimization.

At stage \(i\), the set of weights in the layer \(\ell\) denoted by \(\mathcal{W}_{i}^{(\ell)}\) is partitioned into two subsets \(\mathcal{W}_{i,1}^{(\ell)}\) and \(\mathcal{W}_{i,2}^{(\ell)}\) corresponding to the quantized and unquantized weights, respectively, i.e.,

\[\mathcal{W}_{i}^{(\ell)}=\Big{\{}\mathcal{W}_{i,1}^{(\ell)},\mathcal{W}_{i,2}^ {(\ell)}\Big{\}},\quad\mathcal{W}_{i,1}^{(\ell)}\cap\mathcal{W}_{i,2}^{(\ell )}=\emptyset. \tag{10}\]

The model is first trained over weights in \(\mathcal{W}_{i}^{(\ell)}\) in FP32. Then, the resulting weights in \(\mathcal{W}_{i,1}^{(\ell)}\) are quantized under a suitable PTQ scheme. Next, the weights in \(\mathcal{W}_{i,1}^{(\ell)}\) are fixed, and the model is retrained by minimizing the loss function with respect to the weights in \(\mathcal{W}_{i,2}^{(\ell)}\), starting from the previously trained values. The second group is retrained in order to compensate for the quantization error arising from the first group, and make up for the loss in the accuracy. In stage \(i+1\), the above steps are repeated upon substitution \(\mathcal{W}_{i+1}^{(\ell)}\overset{\Delta}{=}\mathcal{W}_{i,2}^{(\ell)}\). The weight partitioning, group-wise quantization, and retraining is repeated until the network is fully quantized. The total number of partition sets is denoted by \(N_{p}\).

In another version of this algorithm, the partitioning for all stages is set initially. That is to say, the weights of layer \(\ell\) are partitioned into \(N_{p}\) groups \(\{\mathcal{W}_{i}^{(\ell)}\}_{i=1}^{N_{p}}\) and successively quantized, such that at each stage the weights of the previous groups are quantized and fixed, and those of the remaining groups are retrained.

The hyper-parameters of the SPTQ are the choice of the quantizer function in PTQ and the partitioning scheme. There

Fig. 5: (a) Sub-graph introduced after each weight \(w\) in the computational graph of the NN in the AB quantization; (b) the AB quantizer, when the base quantizer is the uniform one.
are several options for the partitioning, such as random grouping, neuron grouping and local clustering. It has been demonstrated that models trained with SPTQ provide classification accuracies comparable to their baseline counterparts trained in 32-bit, with fewer bits [47]. Fig. 9(c) shows that SPTQ improves the Q-factor considerably, around 0.8 dB.

#### Iv-D4 Successive Alpha-blending Quantization

In this section, we propose SAB, a quantization algorithm suitable for the conversion of a small full-precision model to a low-precision one, in the low bit-width regime 1-3 bits, depending on whether or not the activations are quantized.

SAB is an iterative algorithm with several stages, blending SPTQ and AB quantization in a particular manner described below. At stage \(i\), the weights are partitioned into the set \(\mathcal{W}_{i,1}^{(\ell)}\) and \(\mathcal{W}_{i,2}^{(\ell)}\) as in (10). First, each weight \(w\in\mathcal{W}_{i,1}^{(\ell)}\) is updated according to the AB relation (8) as \(\tilde{w}=(1-\alpha_{j})w+\alpha_{j}\tilde{w}\), where \(\alpha_{j}\) is given by (9) at \(j=k_{1}\). Then, the weights \(\tilde{w}\in\mathcal{W}_{i,1}^{(\ell)}\) are fixed, while those in \(\mathcal{W}_{i,2}^{(\ell)}\) are retrained from their previous values. Next, \(\alpha_{j}\) is incremented to the value in the sequence (9) at \(j=k_{1}+1\). The process of partitioning, AB updating, and retraining is repeated until \(\alpha_{j}=1\) is reached at \(j=k_{2}\), where all weights in \(\mathcal{W}_{i,1}^{(\ell)}\) are fully quantized. The algorithm then advances to the next stage \(i+1\), by partitioning \(\mathcal{W}_{i,2}^{(\ell)}\) into two complementary sets. The last partition is trained with the AB algorithm instead of being fixed, to address the problem of the performance drop in the last set that was encountered in SPTQ. The quantization process is summarized in Algorithm 1.

Note that SAB is not directly a combination of SPTQ and AB: the successive retraining strategy is distributed within the AB algorithm with respect to \(\alpha_{j}\). Therefore, SAB quantization improves upon SPTQ and AB quantization, since each partition is not quantized in one shot, rather is incrementally quantized by increasing \(\alpha_{j}\). This allows the trained set \(\mathcal{W}_{i,2}^{(\ell)}\) to adapt to the changes in \(\mathcal{W}_{i,1}^{(\ell)}\). Instead of fixing the last partition as in the SPTQ scheme, the AB algorithm is applied to train the last partition and fix the quantization error. This modification leads to a reduction in the drop in performance occurred in the last partition.

In uniform SAB quantization, the grid is (3). On the other hand, in the companding SAB quantization, first the compressor (5) is applied so that the probability distribution of the weights is approximately uniform on the clipping range. Then, all weights are quantized with the uniform SAB algorithm, and passed through the expander (6).

### _Computational Complexity of the Quantized NNs_

In this Section, we present expressions for the computational complexity of the two NN equalizers described in Section III-B after quantization, in order to quantify the gains of quantization in memory and computation. The complexity is measured in the number of the elementary bit-wise operations (BO) [49]. The reduction in memory is simply \(1-b/32\), where \(b\) is the quantization rate.

#### Iv-E1 FC Layers

Consider a FC layer with \(n_{i}\) inputs each with bit-width \(b_{i}\), \(n_{o}\) neurons at output, and per-weight bit-width of \(b_{w}\). There are \(n_{o}\) inner products, each between vectors of length \(n_{i}\). The main step is the BO to compute an inner product, which is bounded in Appendix A. From (16),

\[\text{BO}_{\text{FC}}\leq n_{o}\Big{(}n_{i}b_{i}b_{w}+(n_{i}-1)(b_{i}+b_{w}+ \log_{2}(n_{i}))\Big{)}. \tag{11}\]

#### Iv-E2 Convolution Layers

Consider a one-dimensional convolutional layer, with an input of length \(n_{i}\) and per-element bit-width \(b_{i}\), and a filter with length \(n_{w}\) and per-element bit-width \(b_{w}\). It is assumed that the filter is padded with zeros on the boundaries so that the number of output features equals to the length of the input vector \(n_{i}\) ("same padding"). This layer requires \(n_{i}\) inner products between vectors of length \(n_{w}\). The BO is thus

\[\text{BO}_{\text{Conv}}\leq n_{i}\Big{(}n_{w}b_{i}b_{w}+(n_{w}-1)(b_{i}+b_{w}+ \log_{2}(n_{w})\Big{)}. \tag{12}\]

#### Iv-E3 LSTM Cells

Consider the LSTM cell described in [24, Eq. 13], with an input of length \(n_{i}\) and hidden state of size \(n_{h}\) at each time step. The cell has four augmented dense matrices with dimension \(n_{h}\times(n_{i}+n_{h}+1)\), in the three gates and the cell activation state. Suppose that the activations, and thus the hidden state, are quantized at \(b_{a}\) bits. The bit-width of the Cartesian product of the quantization grids is upper bounded by the sum of the individual bit-widths. Thus, from (11)

\[\text{BO}_{\text{LSTM}} \leq 4n_{h}\Big{\{}(n_{h}+n_{i}+1)\times(b_{i}+b_{a})b_{w}+(n_{h}+n_{i}) \tag{13}\] \[\times\big{(}b_{w}+b_{i}+b_{a}+\log_{2}(n_{h}+n_{i}+1))\Big{\}}.\]

Clearly, \(\text{BO}_{\text{BiLSTM}}=2\text{BO}_{\text{LSTM}}\).

Substituting \(b_{1}=b_{2}\) in (16), the storage and BO of the NN scale, respectively, linearly and quadratically with the bit-width. Therefore, quantization from FP32 at 4 bits reduces the memory by 8X, and complexity by 64X.

The BO of the Conv-FC and BiLSTM-FC models are obtained by combining (11), (12) and (13).
### _Quantization of NNs in Optical Fiber Communication_

The uniform and PoT PTQ (representing fixed-point numbers) have been naturally applied when demonstrating the NN equalizers in FPGA [50, 51] or ASIC [52], usually at 8 bits. PTQ has been applied to the NNs mitigating the nonlinear distortions in optical fiber [28, 53, 54, 55, 56, 48, 53], and the inter-symbol interference (ISI) in passive optical networks (PONs) with intensity-modulation direct-detection (IMDD) [51, 57, 58] and in general dispersive additive white Gaussian noise (AWGN) channels [59]. In particular, the authors of [51] show that an MLP-based many-to-many equalizer outperforms the maximum likelihood sequence estimator in mitigating the ISI in an IMDD 30 km PON link. They implement the NN in FPGA, and determine the impact of the weight resolution on the BER at 2-8 bits. In [54], a multi-layer perceptron equalizing a 1000 km SMF link is pruned and quantized with uniform PTQ-8, and the reduction in BO is reported. The authors of [52] implement the time-domain LDBP in ASIC, where the filter coefficients, as well as the signal in each step of SSFM, are quantized.

The APoT is considered in [28, 56, 60]. Fixed-point optimization-based PoT quantization is applied to an MLP equalizing an AWGN channel in [61]. The weights are quantized at 4 bits and activations at 14 bits. The authors of [60] represent the weights using a 2-term APoT expression, for multiplier-free NN nonlinearity mitigation in a 22x80 km SMF link. However, the quantization rate is not constrained.

The mixed-precision quantization is applied to a perturbation-based equalizer in [53] (similar to the Volterra equalizer) in a 18x100 km SMF link, in which the perturbation coefficients larger than a threshold are quantized at large bit-width, and the rest at one bit. Here, the quantization also simplifies the sum expressing the equalizer, combining the identical or similar terms [62].

In our prior work, we compared PTQ, QAT-STE, APoT [28] and SPQT [48] for the quantization of the NN equalizers. However, the best rate here is 5 bits. The authors of [56] study PTQ, QAT-STE and APoT, and demonstrate that the NN weights can be stored with a range of bit-widths and penalties, using pruning, quantization and compression.

The papers cited above mostly implement uniform, PoT, or APoT PTQ. In our experiments, these algorithms, and their combinations with the QAT-STE, did not achieve sufficiently small distortions in the low bit-width regime. The penalty due to the quantization depends on the size of the model. The current paper addresses the quantization error, using the SAB algorithm that lowers the rate markedly to 1-3 bits. Moreover, the activations are usually not quantized in the literature. In contrast, in this paper both weights and activations are quantized. Importantly, it will be shown in Section V that the quantization of activations impacts the performance considerably. Finally, quantization has been applied in the literature usually as an ingredient in a broader study, or combined with pruning and compression techniques. This paper provides a detailed analysis of the performance and complexity trade-off of different quantization algorithms, and goes beyond the previously reported results [28, 48] in technical advances, application, and discussions.

## V Demonstration of the Quantization Gains in Experiments

In this Section, we determine the performance and complexity trade-off of the several quantization algorithms. We compute the Q-factor penalty as a function of the launch power and quantization rate, as well as the reduction in the memory and computational complexity, in the three transmission experiments described in Section II.

### _TWC Fiber Experiment_

We consider the TWC fiber dual-polarization transmission experiment in Section II-2a, with the Conv-FC model in

\begin{table}
\begin{tabular}{c c c c c} \hline \hline \multicolumn{2}{c}{Bit-width} & \multicolumn{3}{c}{Q-factor} \\ Convolutional & Dense & Quantizer & -2 dBm & 2 dBm \\ \hline
32 & 32 & Unquantized & 8.6 & 7.54 \\
6 & 8 & Uniform & 8.1 & 6.34 \\
6 & 8 & ApoT & 8.4 & 7.4 \\ \hline \hline \end{tabular}
\end{table} TABLE II: UNIFORM VS NON-UNIFORM QUANTIZATION IN TWC FIBER EXPERIMENT

Fig. 6: Q-factor of the NN equalizer in the TWC fiber experiment. a) PTQ; b) QAT-STE; (c) SPTQ.
Section III-B2. The hyper-parameters of this model are the size of the convolutional filters \(K\) and the number of hidden neurons \(n_{h}\). The filters' length is set to be the residual channel memory, \(K=M\). This is estimated to be \(M=40\) complex symbols per polarization, through the auto-correlation function of the received symbols after CPE, and performance evaluation. The minimum number of hidden units is \(n_{h}=100\), below which the performance rapidly drops.

The NN is trained with 600,000 symbols from a 16-QAM constellation. A test set of 100,000 symbols is used to assess the performance of the NN. Each dataset is measured at a given power, during which the BER may fluctuate in time due to the environmental changes. The symbols on the boundary of the data frame are eliminated to remove the effects of the anomalies. The NN at each power is trained and tested with independent datasets of randomly chosen symbols at the same power. The NN is implemented in the Python's TensorFlow library. The loss function is the mean-squared error, and the learning algorithm is the Adam-Optimizer with the learning rate of 0.001. The libraries such as TensorFlow provide functions for basic PTQ and QAT-STE, however, at 8 bits or more. For quantization at an arbitrary bit-width \(b<8\), the algorithms have to be directly programmed. For benchmark models in deep learning, low bit-width implementations exist.

For quantization above 5 bits, PTQ and QAT-STE are applied, combined with APoT quantization, fixed- or mixed precision. In fixed-precision PTQ, the weights and activations of all layers are quantized at 6, 7 or 8 bits. In mixed-precision PTQ, 6 bits is assigned to the weights and activations of the convolutional layer, whereas the dense layer is given 8 bits due to its more significant impact on the performance. The Q-factor is nearly not impacted at 8 bits. Fig. 6(a) demonstrates that fixed-precision PTQ-6 incurs a penalty of 0.7 dB at -2 dBm compared to the unquantized NN, and 1.9 dB at 2 dBm. This comes with a gain of \(81\%\) reduction in the memory usage and a \(95\%\) reduction in the computational complexity.

The Q-factor improves using the QAT-STE, as depicted in Fig. 6(b). Here, the weights are initialized with random values, then trained and quantized at 5, 6, and 7 bits, and the activations at 6 bits. In this case, the drop is reduced to 0.5 dB at -2 dBm, and 1.2 dB at 2 dBm. As the transmission power is increased, the penalty due to the quantization increases.

The distribution of the weights of the dense layer is bell-shaped, as shown in Fig. 4. In consequence, assigning more quantization symbols around the mean is a reasonable strategy. The APoT quantization delivers a good performance, with a Q-factor penalty of less than \(0.2\) dB at \(-2\) and \(2\) dBm, as seen in Table II.

The uniform SPTQ is applied, by assigning 5 bits to the weights and activations of the dense layer. The convolutional layer is given 8 bits, but this layer has few weights, and little impact on the complexity. Fig. 6(c) shows that SPTQ at 5 bits leads to 0.2 dB Q-factor drop at -2 dBm, and 0.5 dB at 2 dBm. It can be seen that SPTQ outperforms the more complex QAT-STE by 2 bits at the same power [48]. Fig. 9(c) shows that increasing the partition size can notably enhance the Q-factor. Similar conclusions are drawn for SPTQ-4, as seen in Table III.

For quantization below 5 bits, we apply SAB. In a first study, we consider fixed-precision quantization, where the weights and activations are quantized at 4 bits successively over 4 partitions. The results in Table IV indicate that SAB outperforms SPTQ and AB, with a performance drop of 0.5 dB near optimal power. In contrast, SPTQ and AB quantization resulted in a 1.2 dB drop in performance. In a second study, we apply mixed-precision SAB, giving more bits to the last partition. We consider a partition of size 4 with the weights and activations in the first three partition sets quantized at 4 bits, and in the last set at 6 bits, averaging to 4.5 bits. The results are shown in Fig. 9(a), indicating the Q-factor drop of 0.17 dB at -2 dBm and 0.24 dB at 2 dBm. This comes with \(86\%\) reduction in memory usage, and \(94\%\) in computational complexity.

### _SMF Experiment_

We consider the SMF experiment described in Section II-2b, with the Conv-FC model. The NN parameters and the quantization algorithms are similar to those in the TWC fiber experiment.

For quantization above 5 bits, PTQ-6 led to a Q-factor drop of 0.3 dB at 1 dBm, and 0.4 dB at 4 dBm, as shown in Fig. 7

Fig. 7: Q-factor of the NN equalizer in the SMF experiment. a) PTQ; b) QAT-STE; c) uniform and companding PTQ.
(a). For QAT-STE-6, as shown in Fig. 7(b), the drop is 0.1 dB at 1 dBm, and 0.2 dB at 4 dBm.

For quantization below 5 bits, first the companding PTQ is applied. Fig. 7(c) shows that this quantizer outperforms the uniform quantization at 4 bits by about a dB, due to the non-uniform distribution of the weights of the dense layer. It is found that, while the APoT works well in the large bit-width regime \(b\geq 6\) (as in the TWC fiber experiment), it is uncompetitive at low bit-widths.

Next, we apply SAB quantization, in a partition of size 4, where the weights in the first 3 sets are quantized at 3 bits, and in the last set at 6 bits, with the average rate of 3.75 bits. The activations for all partition sets are quantized at 3 bits. The uniform and companding versions are both studied. Fig. 9(b) shows the results. Uniform SAB quantization results in a Q-factor drop of 0.3 dB at 1 dBm, and 0.6 dB at 4 dBm. This quantizer offers a reduction in memory usage and computational complexity, by \(88\%\) and \(94\%\), respectively. Applying the companding SAB quantization, the Q-factor drop is reduced to 0.2 dB at 1 dBm.

### _LEAF Experiment_

The NN in this experiment is the BiLSTM-FC equalizer, described in Section III-B2. There are \(n_{h}=100\) hidden neurons, and the input size is \(\bar{n}_{i}=4(M+1)\), \(M=40\). This model is found to be prone to the quantization error, because small errors can be amplified by the internal activations, and accumulate over long input temporal sequences. Thus, we quantize the weights and biases of the forget, input and output gates, as well as the activations at the output of the cell. However, the internal activations remain in full precision.

Fig. 8 (a) shows that PTQ-6 incurs a Q-factor penalty of \(0.9\) dB at 1 dBm, and \(1.2\) dB at \(-1\) dBm, respectively, while lowering the computational complexity by \(79\%\) and the memory usage by \(81\%\). QAT-STE significantly improves the Q-factor, as shown in Fig. 8 (b). At 6 bits, the drop is \(0.1\) dB at 1 dBm, and \(0.4\) dB at \(-1\) dBm. At 5 bits, the penalty is \(0.3\) dB at both \(1\) dBm and \(-1\) dBm, with \(82\%\) reduction in computational complexity and \(84\%\) in memory usage.

Fig. 8(c) shows that the AB quantizer at 4 and 5 bits outperforms PTQ and QAT Specifically, the Q-factor drop is only \(0.2\) dB at -1 dBm, and \(0.15\) dB at 1dBm.

### _Quantization of the Weights, but not Activations_

In the previous sections, the weights and activations were both quantized. It can be seen that there is a cut-off bit-width around 5-6 bits, below which the performance of the QAT-STE rapidly drops. Upon investigation, we noticed that the quantization of the activations substantially impacts the Q-factor. The activation functions are nonlinear, and could amplify the quantization error. In this section, we consider quantizing the weights of the NN but not activations. The bit-width of the activations can still be reduced from 32 to 8 with negligible performance drop. Therefore, the activations are quantized, at 8 bits.

In a first study, we quantize the weights of the Conv-FC model in the SMF experiment, using the fixed-precision SAB algorithm with a partition of size 4. The results are included in Table V, showing that the Q-factor drop at the optimal power is minimal, when the dense layer is quantized at as low as 3 bits. In a second study, we apply the mixed-precision SAB quantization with the same parameters. The first three partitions are quantized at 1 bit, and the last one at 4 bits. We obtain a quantization rate of 1.75 bits/weight, with 0.6 dB degradation in Q-factor, outperforming the state-of-the-art

\begin{table}
\begin{tabular}{c|c c} \hline Quantization scheme & bit-width & Q-factor \\ \hline Unquantized & 32 & 7.5 \\ SPTQ & 4 & 6.3 \\ AB & 4 & 6.3 \\ SAB & 4 & 7.0 \\ \hline \end{tabular}
\end{table} TABLE IV: FIXED-PRECISION QUANTIZATION, TWC FIBER EXPERIMENT

Fig. 8: Q-factor of the NN equalizer in the LEAF experiment. a PTQ; b) QAT-STE; (c) AB quantization.

\begin{table}
\begin{tabular}{c|c c c c c c c c} \hline \(N_{p}\) & \multicolumn{6}{c}{Q-factor} \\  & \(\mathcal{W}_{1}\) & \(\mathcal{W}_{2}\) & \(\mathcal{W}_{3}\) & \(\mathcal{W}_{4}\) & \(\mathcal{W}_{5}\) & \(\mathcal{W}_{6}\) & \(\mathcal{W}_{7}\) & \(\mathcal{W}_{8}\) \\ \hline
2 & 7.13 & **5.6** & & & & & & \\
4 & 7.5 & 7.33 & 7.33 & **6.3** & & & & \\
8 & 7.56 & 7.5 & 7.4 & 7.33 & 7.33 & 7.33 & **6.6** \\ \end{tabular}
\end{table} TABLE III: Q-FACTOR OF SPTQ-4, TWC FIBER EXPERIMENT
using the QAT-STE w/wo APoT by 2 dB. This important result demonstrates that low-complexity nearly-binary NNs can mitigate nonlinearities in optical fiber communication.

In the so called "extreme quantization," the NNs are quantized at 1 or 2 bits [63, 64, 14, 65]. Many approaches to the binary and ternary NNs have been proposed, e.g., based on better approximations to the derivative of the quantizer than in the STE. However, we tested some of these approaches in our experiments, and did not observe notable gains over the linear equalization. Consequently, while extreme quantization has shown success in large models in computer vision, further work is needed to determine if it can be adapted and successfully applied to the small NN equalizers in optical fiber communication.

## VI Conclusions

The paper shows that low-complexity quantized NNs can mitigate nonlinearities in optical fiber transmission. The QAT-STE partially mitigates the quantization error during the training, and is effective in the large bit-width regime with \(b>5\) bits. The companding quantization improves the Q-factor of the baseline schemes considerably, especially at low bit-widths. There is a cut-off bit-width of around 5 bits below which the penalty of the QAT-STE rapidly increases. In the low bit-width regime with \(b\leq 5\) bits, companding SAB quantization is the method of choice. There is a considerable performance penalty due to the quantization of activations. The weights of the NN can be quantized at 1.75 bits/parameter with \(\leq 0.5\) dB penalty, if the activations are quantized at \(b\geq 8\) bits. The weights and activations can be quantized at 3.75 bits/parameter, with minimal penalty. The LSTM-based receivers can be prone to the quantization error, due to the error amplification and propagation. Fully binary NN equalizers remain to be studied.

## Appendix A Bit-wise Operations for an Inner Product

The cost of computation is measured here by the required bit-wise operations AND \(\wedge\), OR \(\vee\), XOR \(\oplus\), NOT and SHIFT [49].

### _Addition and Multiplication of Integers_

The sum \(z=x+y\) of the integers \(x\) and \(y\) each with bit-width \(b\) is an integer with bit-width \(b+1\), with carry-over. Below, we show that \(z\) can be computed in \(\zeta b\) BO, where \(\zeta\) depends on the computing algorithm.

Denote the binary representation of \(x\), \(y\) and \(z\) with \(x_{1}x_{2}\cdots x_{b}\), \(y_{1}y_{2}\cdots y_{b}\), and \(z_{1}z_{2}\cdots z_{b+1}\), respectively. Let \(c_{1}c_{2}\cdots c_{b+1}\) be the carry-over binary sequence, initialized with \(c_{1}=0\). Then, for \(i\in\{1,2,\cdots,b+1\}\)

\[z_{i}=t\oplus c_{i},\quad c_{i+1}=(x_{i}\wedge y_{i})\vee(t\wedge c_{i}), \tag{14}\]

where \(t=x_{i}\oplus y_{i}\). Thus, computing \(z\) using (14) takes \(5b\) BO, i.e., \(\zeta=5\). This approach requires one bit storage for \(t\), and \(2b\) bits transmission for memory access.

Consider the multiplication of the integers \(\bar{z}=xy\), where \(x\) has bit-width \(b_{1}\) and \(y\) has \(b_{2}\) bits. Clearly, the bit-width of \(\bar{z}\) is \(b_{1}+b_{2}\). The multiplication \(2^{i}y\), \(i\in\mathbb{N}\), can be performed with one BO, by shifting the \(y\) in the binary form \(i\) positions to the left, and zero padding from right. The result is a binary sequence of the maximum length \(b_{1}+b_{2}\), and maximum \(b_{2}\) non-zero bits. Expanding \(x\) as a sum of \(b_{1}\) PoT numbers, \(\bar{z}\) is expressed as the sum of \(b_{1}\) binary sequences, each with up to \(b_{2}\) non-zero elements. Thus, \(\text{BO}=\zeta b_{1}b_{2}\).

The value of \(\zeta\) can change with the algorithm, and is immaterial. In this paper, we assume \(\zeta=1\). The computation of \(z\) and \(\bar{z}\) above may not be optimal; hence the BOs are upper bounds.

### _The Inner Product_

The sum of \(n\) numbers of bit-width \(b\) can be performed in \(\log_{2}n\) steps by pairwise addition (assuming for simplicity that \(n\) is a PoT number). The sum has bit-width \(b+\log_{2}(n)-1\) bits. The BO can be bounded as below, or obtained from [66].

\[\text{BO}_{\text{sum}} \leq b\times\frac{n}{2}+(b+1)\times\frac{n}{4}+\cdots+(b+\log_{2 }(n)-1)\times 1\] \[=\frac{n}{2}\Big{[}b\sum_{k=0}^{\log_{2}(n)-1}2^{-k}+\sum_{k=1}^ {\log_{2}(n)-1}k2^{-k}\Big{]}\] \[\leq\frac{n}{2}\Big{[}(b+\log_{2}n-1)\sum_{k=0}^{\log_{2}(n)-1}2^ {-k}\Big{]}\] \[=(b+\log_{2}n)(n-1). \tag{15}\]

Consider the inner product \(y=\mathbf{w}^{T}\mathbf{x}\), where \(\mathbf{w}=(w_{1},w_{2},\cdots,w_{n})\), \(\mathbf{x}=(x_{1},x_{2},\cdots,x_{n})\), and where \(w_{i}\) and \(x_{i}\) have, respectively, bit-width \(b_{1}\) and \(b_{2}\), \(\forall i\). Then, \(y\) has bit-width \(b_{1}+b_{2}+\log_{2}(n)-1\) bits. The products \(\{w_{i}x_{i}\}_{i=1}^{n}\) are calculated in \(nb_{1}b_{2}\) BO. Their sum is computed in BO given in (15) with \(b=b_{1}+b_{2}\). Thus

\[\text{BO}_{\text{inner}}\leq nb_{1}b_{2}+(n-1)(b_{1}+b_{2}+\log_{2}n). \tag{16}\]
