\section{Introduction}

Recent developments on pretrained contextualized embeddings have significantly improved the performance of structured prediction tasks in natural language processing. Approaches based on contextualized embeddings, such as ELMo \citep{peters-etal-2018-deep}, Flair \citep{akbik-etal-2018-contextual}, BERT \citep{devlin-etal-2019-bert}, and XLM-R \citep{conneau-etal-2020-unsupervised}, have been consistently raising the state-of-the-art for various structured prediction tasks. Concurrently, research has also showed that word representations based on the concatenation of multiple pretrained contextualized embeddings and traditional non-contextualized embeddings (such as word2vec \citep{mikolov2013distributed} and character embeddings \citep{santos2014learning}) can further improve performance \citep{peters-etal-2018-deep,akbik-etal-2018-contextual,strakova-etal-2019-neural,wang-etal-2020-more}. Given the ever-increasing number of embedding learning methods that operate on different granularities (e.g., word, subword, or character level) and with different model architectures, choosing the best embeddings to concatenate for a specific task becomes non-trivial, and exploring all possible concatenations can be prohibitively demanding in computing resources.

Neural architecture search (NAS) is an active area of research in deep learning to automatically search for better model architectures, and has achieved state-of-the-art performance on various tasks in computer vision, such as image classification \citep{real2019regularized}, semantic segmentation \citep{liu2019auto}, and object detection \citep{ghiasi2019fpn}. In natural language processing, NAS has been successfully applied to find better RNN structures  \citep{zoph2016neural,pham2018efficient} and recently better transformer structures \citep{so2019evolved,zhu2020autotrans}. In this paper, we propose Automated Concatenation of Embeddings (ACE) to automate the process of finding better concatenations of embeddings for structured prediction tasks. ACE is formulated as an NAS problem. In this approach, an iterative search process is guided by a controller based on its belief that models the effectiveness of individual embedding candidates in consideration for a specific task. At each step, the controller samples a concatenation of embeddings according to the belief model and then feeds the concatenated word representations as inputs to a task model, which in turn is trained on the task dataset and returns the model accuracy as a reward signal to update the belief model. We use the policy gradient algorithm \citep{williams1992simple} in reinforcement learning \citep{sutton2018reinforcement} to solve the optimization problem. In order to improve the efficiency of the search process, we also design a special reward function by accumulating all the rewards based on the transformation between the current concatenation and all previously sampled concatenations.

Our approach is different from previous work on NAS in the following aspects:
\begin{enumerate}[leftmargin=*]
	\item Unlike most previous work, we focus on searching for better word representations rather than better model architectures.
    \item We design a novel search space for the embedding concatenation search. Instead of using RNN as in previous work of \citet{zoph2016neural}, we design a more straightforward controller to generate the embedding concatenation. We design a novel reward function in the objective of optimization to better evaluate the effectiveness of each concatenated embeddings.
    \item ACE achieves high accuracy without the need for retraining the task model, which is typically required in other NAS approaches.
    \item Our approach is efficient and practical. Although ACE is formulated in a NAS framework, ACE can find a strong word representation on a single GPU with only a few GPU-hours for structured prediction tasks. In comparison, a lot of NAS approaches require dozens or even thousands of GPU-hours to search for good neural architectures for their corresponding tasks.

\end{enumerate}

Empirical results show that ACE outperforms strong baselines.
Furthermore, when ACE is applied to concatenate pretrained contextualized embeddings fine-tuned on specific tasks, we can achieve
state-of-the-art accuracy on 6 structured prediction tasks including Named Entity Recognition \citep{Sundheim1995NamedET}, Part-Of-Speech tagging \citep{derose-1988-grammatical}, chunking \citep{tjong-kim-sang-buchholz-2000-introduction}, aspect extraction \citep{10.1145/1014052.1014073}, syntactic dependency parsing \citep{arrive1969elements} and semantic dependency parsing \citep{oepensemeval} over 21 datasets.
Besides, we also analyze the advantage of ACE and reward function design over the baselines and show the advantage of ACE over ensemble models.

\section{Related Work}
\subsection{Embeddings}

Non-contextualized embeddings, such as word2vec \citep{mikolov2013distributed}, GloVe \citep{pennington2014glove}, and fastText \citep{bojanowski2017enriching}, help lots of NLP tasks.
Character embeddings \citep{santos2014learning} are trained together with the task and applied in many structured prediction tasks \citep{ma-hovy-2016-end,lample-etal-2016-neural,dozat-manning-2018-simpler}.
For pretrained contextualized embeddings, ELMo \citep{peters-etal-2018-deep}, a pretrained contextualized word embedding generated with multiple Bidirectional LSTM layers, significantly outperforms previous state-of-the-art approaches on several NLP tasks.
Following this idea, \citet{akbik-etal-2018-contextual} proposed Flair embeddings, which is a kind of contextualized character embeddings and achieved strong performance in sequence labeling tasks.
Recently, \citet{devlin-etal-2019-bert} proposed BERT, which encodes contextualized sub-word information by Transformers \citep{vaswani2017attention} and significantly improves the performance on a lot of NLP tasks.
Much research such as RoBERTa \citep{liu2019roberta} has focused on improving BERT model's performance through stronger masking strategies. Moreover, multilingual contextualized embeddings become popular. \citet{pires-etal-2019-multilingual} and \citet{wu-dredze-2019-beto} showed that Multilingual BERT (M-BERT) could learn a good multilingual representation effectively with strong cross-lingual zero-shot transfer performance in various tasks. \citet{conneau-etal-2020-unsupervised} proposed XLM-R, which is trained on a larger multilingual corpus and significantly outperforms M-BERT on various multilingual tasks.

\subsection{Neural Architecture Search}
Recent progress on deep learning has shown that network architecture design is crucial to the model performance.
However, designing a strong neural architecture for each task requires enormous efforts, high level of knowledge, and experiences over the task domain.
Therefore, automatic design of neural architecture is desired.
A crucial part of NAS is search space design, which defines the discoverable NAS space.
Previous work \citep{baker2016designing,zoph2016neural,xie2017genetic} designs a global search space \citep{elsken2019neural} which incorporates structures from hand-crafted architectures. For example, \citet{zoph2016neural} designed a chained-structured search space with skip connections.
The global search space usually has a considerable degree of freedom.
For example, the approach of \citet{zoph2016neural} takes 22,400 GPU-hours to search on CIFAR-10 dataset. Based on the observation that existing hand-crafted architectures contain repeated structures \citep{szegedy2016rethinking,he2016deep,huang2017densely}, \citet{zoph2018learning} explored cell-based search space which can reduce the search time to 2,000 GPU-hours.

In recent NAS research, reinforcement learning and evolutionary algorithms are the most usual approaches.
In reinforcement learning, the agent's actions are the generation of neural architectures and the action space is identical to the search space. Previous work usually applies an RNN layer \citep{zoph2016neural,zhong2018practical,zoph2018learning} or use Markov Decision Process \citep{baker2016designing} to decide the hyper-parameter of each structure and decide the input order of each structure.
Evolutionary algorithms have been applied to architecture search for many decades \citep{geoffrey1989designing,angeline1994evolutionary,stanley2002evolving,floreano2008neuroevolution,jozefowicz2015empirical}. The algorithm repeatedly generates new populations through recombination and mutation operations and selects survivors through competing among the population. Recent work with evolutionary algorithms differ in the method on parent/survivor selection and population generation. For example, \citet{real2017large}, \citet{liu2018hierarchical}, \citet{wistuba2018deep} and \citet{real2019regularized} applied tournament selection \citep{goldberg1991comparative} for the parent selection while \citet{xie2017genetic} keeps all parents. \citet{suganuma2017genetic} and \citet{elsken2017simple} chose the best model while \citet{real2019regularized} chose several latest models as survivors.

\section{Automated Concatenation of Embeddings}
In ACE, a task model and a controller interact with each other repeatedly. The task model predicts the task output, while the controller searches for better embedding concatenation as the word representation for the task model to achieve higher accuracy. Given an embedding concatenation generated from the controller, the task model is trained over the task data and returns a reward to the controller. The controller receives the reward to update its parameter and samples a new embedding concatenation for the task model.
Figure \ref{fig:architecture} shows the general architecture of our approach.
\subsection{Task Model}
For the task model, we emphasis on sequence-structured and graph-structured outputs. Given a structured prediction task with input sentence ${\bm{x}}$ and structured output ${\bm{y}}$, we can calculate the probability distribution $P({\bm{y}}|{\bm{x}})$ by:
\begin{align}
P({\bm{y}}|{\bm{x}}) = \frac{\exp{(\text{Score}({\bm{x}},{\bm{y}})})}{\sum_{{\bm{y}}^{\prime} \in {\mathbb{Y}}({\bm{x}})} \exp{(\text{Score}({\bm{x}},{\bm{y}}^{\prime})})} \nonumber
\end{align}
where ${\mathbb{Y}}({\bm{x}})$ represents all possible output structures given the input sentence ${\bm{x}}$. Depending on different structured prediction tasks, the output structure ${\bm{y}}$ can be label sequences, trees, graphs or other structures. In this paper, we use sequence-structured and graph-structured outputs as two exemplar structured prediction tasks. We use BiLSTM-CRF model \citep{ma-hovy-2016-end,lample-etal-2016-neural} for sequence-structured outputs and use BiLSTM-Biaffine model \citep{dozat2016deep} for graph-structured outputs:
\begin{align}
P^{\text{seq}}({\bm{y}}|{\bm{x}}) &= \text{BiLSTM-CRF}({\bm{V}},{\bm{y}})\nonumber\\
P^{\text{graph}}({\bm{y}}|{\bm{x}}) &= \text{BiLSTM-Biaffine}({\bm{V}},{\bm{y}})\nonumber
\end{align}
where ${\bm{V}} = [{\bm{v}}_1; \cdots; {\bm{v}}_n]$, ${\bm{V}} \in \mathbb{R}^{d\times n}$ is a matrix of the word representations for the input sentence ${\bm{x}}$ with $n$ words, $d$ is the hidden size of the concatenation of all embeddings. The word representation ${\bm{v}}_i$ of $i$-th word is a concatenation of $L$ types of word embeddings:
\begin{align}
{\bm{v}}_i^l &= \text{embed}_i^l ({\bm{x}}); \;\; {\bm{v}}_i = [{\bm{v}}_i^1;{\bm{v}}_i^2; \dots; {\bm{v}}_i^L] \nonumber
\end{align}
where $\text{embed}^l$ is the model of $l$-th embeddings, ${\bm{v}}_i\in \mathbb{R}^d$, ${\bm{v}}_i^l\in \mathbb{R}^{d^l}$. $d^l$ is the hidden size of $\text{embed}^l$.

\subsection{Search Space Design}
\label{sec:search_space}
The neural architecture search space can be represented as a set of neural networks \citep{elsken2019neural}.
A neural network can be represented as a directed acyclic graph with a set of nodes and directed edges.
Each node represents an operation, while each edge represents the inputs and outputs between these nodes.
In ACE, we represent each embedding candidate as a node.
The input to the nodes is the input sentence ${\bm{x}}$, and the outputs are the embeddings ${\bm{v}}^l$. Since we concatenate the embeddings as the word representation of the task model, there is no connection between nodes in our search space. Therefore, the search space can be significantly reduced.
For each node, there are a lot of options to extract word features.
Taking BERT embeddings as an example, \citet{devlin-etal-2019-bert} concatenated the last four layers as word features while \citet{kondratyuk-straka-2019-75} applied a weighted sum of all twelve layers. However, the empirical results \citep{devlin-etal-2019-bert} do not show a significant difference in accuracy. We follow the typical usage for each embedding to further reduce the search space. As a result, each embedding only has a fixed operation and the resulting search space contains $2^L{-}1$ possible combinations of nodes.

In NAS, weight sharing \citep{pmlr-v80-pham18a} shares the weight of structures in training different neural architectures to reduce the training cost. In comparison, we fixed the weight of pretrained embedding candidates in ACE except for the character embeddings.
Instead of sharing the parameters of the embeddings, we share the parameters of the task models at each step of search.
However, the hidden size of word representation varies over the concatenations, making the weight sharing of structured prediction models difficult. Instead of deciding whether each node exists in the graph, we keep all nodes in the search space and add an additional operation for each node to indicate whether the embedding is masked out. To represent the selected concatenation, we use a binary vector ${\bm{a}}=[a_1, \cdots, a_l, \cdots, a_L]$ as an mask to mask out the embeddings which are not selected:
\begin{align}
{\bm{v}}_i = [{\bm{v}}_i^1a_1;\dots;{\bm{v}}_i^la_l; \dots; {\bm{v}}_i^La_L] \label{eq:vector}
\end{align}
where $a_l$ is a binary variable. Since the input ${\bm{V}}$ is applied to a linear layer in the BiLSTM layer, multiplying the mask with the embeddings is equivalent to directly concatenating the selected embeddings:
\begin{align}
{\bm{W}}^{\top}{\bm{v}}_i = \sum_{l=1}^L {\bm{W}}_l^{\top}{\bm{v}}_i^l a_l \label{eq:linear}
\end{align}
where ${\bm{W}} {=} [{\bm{W}}_1;{\bm{W}}_2;\dots ;{\bm{W}}_L]$ and ${\bm{W}}{\in} \mathbb{R} ^ {d \times h}$ and ${\bm{W}}_l{\in} \mathbb{R} ^ {d^l \times h}$. Therefore, the model weights can be shared after applying the embedding mask to all embedding candidates' concatenation. Another benefit of our search space design is that we can remove the unused embedding candidates and the corresponding weights in ${\bm{W}}$ for a lighter task model after the best concatenation is found by ACE.

\begin{figure*}
	\centering
	\includegraphics[scale=0.8]{rl.pdf}
	\caption{The main paradigm of our approach is shown in the middle, where an example of reward function is represented in the left and an example of a concatenation action is shown in the right.}
	\label{fig:architecture}
\end{figure*}

\subsection{Searching in the Space}
During search, the controller generates the embedding mask for the task model iteratively.
We use parameters ${\bm{\theta}}=[\theta_1;\theta_2;\dots;\theta_L]$ for the controller instead of using the RNN structure applied in previous approaches \citep{zoph2016neural,zoph2018learning}. The probability distribution of selecting an concatenation ${\bm{a}}$ is $P^{\text{ctrl}}({\bm{a}};{\bm{\theta}})=\prod_{l=1}^LP_l^{\text{ctrl}}(a_l;\theta_l)$. Each element $a_l$ of ${\bm{a}}$ is
sampled independently from a Bernoulli distribution,
which is defined as:
\begin{align}
P_l^{\text{ctrl}}(a_l;\theta_l) {=}
    \begin{cases}
    \sigma(\theta_l) &a_l{=}1\\
    1{-} P_l^{\text{ctrl}}(a_l{=}1;\theta_l) &a_l{=}0
    \end{cases}
    \label{eq:prob}
\end{align}
where $\sigma$ is the sigmoid function.
Given the mask, the task model is trained until convergence and returns an accuracy $R$ on the development set.
As the accuracy cannot be back-propagated to the controller, we use the reinforcement algorithm for optimization.
The accuracy $R$ is used as the reward signal to train the controller.
The controller's target is to maximize the expected reward $J({\bm{\theta}})=\mathbb{E}_{P^{\text{ctrl}}({\bm{a}};{\bm{\theta}})}[R]$ through the policy gradient method \citep{williams1992simple}.
In our approach, since calculating the exact expectation is intractable, the gradient of $J({\bm{\theta}})$ is approximated by sampling only one selection following the distribution $P^{\text{ctrl}}({\bm{a}};{\bm{\theta}})$ at each step for training efficiency:
\begin{align}
\nabla_{\bm{\theta}} J({\bm{\theta}}) \approx  \sum_{l=1}^L  \nabla_{\bm{\theta}} \log P_l^{\text{ctrl}}({a}_l;\theta_l) (R-b) \label{eq:rl}
\end{align}
where $b$ is the baseline function to reduce the high variance of the update function. The baseline usually can be the highest accuracy during the search process.
Instead of merely using the highest accuracy of development set over the search process as the baseline, we design a reward function on how each embedding candidate contributes to accuracy change by utilizing all searched concatenations' development scores.
We use a binary vector $|{\bm{a}}^t-{\bm{a}}^{i}|$ to represent the change between current embedding concatenation ${\bm{a}}^t$ at current time step $t$ and ${\bm{a}}^{i}$ at previous time step $i$.
We then define the reward function as:
\begin{align}
{\bm{r}}^t = \sum_{i=1}^{t-1} (R_t-R_i) |{\bm{a}}^t-{\bm{a}}^{i}| \label{eq:reward_a}
\end{align}
where ${\bm{r}}^t$ is a vector with length $L$ representing the reward of each embedding candidate. $R_t$ and $R_i$ are the reward at time step $t$ and $i$.
When the Hamming distance of two concatenations $Hamm({\bm{a}}^t,{\bm{a}}^{i})$ gets larger, the changed candidates' contribution to the accuracy becomes less noticeable. The controller may be misled to reward a candidate that is not actually helpful.
We apply a discount factor to reduce the reward for two concatenations with a large Hamming distance to alleviate this issue.
Our final reward function is:
\begin{align}
{\bm{r}}^t {=} \sum_{i=1}^{t-1} (R_t{-}R_i) \gamma ^{Hamm({\bm{a}}^t,{\bm{a}}^{i})-1} |{\bm{a}}^t{-}{\bm{a}}^{i}| \label{eq:reward}
\end{align}
where $\gamma \in (0,1)$. Eq. \ref{eq:rl} is then reformulated as:
\begin{align}
\nabla_{\bm{\theta}} J_t({\bm{\theta}}) \approx \sum_{l=1}^L \nabla_{\bm{\theta}} \log P_l^{\text{ctrl}}({a}_l^t;\theta_l) {r}^t_{l} \label{eq:gradient}
\end{align}

\subsection{Training}
To train the controller, we use a dictionary ${\mathbb{D}}$ to store the concatenations and the corresponding validation scores. At $t=1$, we train the task model with all embedding candidates concatenated. From $t=2$, we repeat the following steps until a maximum iteration $T$:
\begin{enumerate}[leftmargin=*]
    \item Sample a concatenation ${\bm{a}}^t$ based on the probability distribution in Eq. \ref{eq:prob}.
    \item Train the task model with ${\bm{a}}^t$ following Eq. \ref{eq:vector} and evaluate the model on the development set to get the accuracy $R_t$.
    \item Given the concatenation ${\bm{a}}^t$, accuracy $R_t$ and ${\mathbb{D}}$, compute the gradient of the controller following Eq. \ref{eq:gradient} and update the parameters of controller.
    \item Add ${\bm{a}}^t$ and $R_t$ into ${\mathbb{D}}$, set $t=t+1$.
\end{enumerate}
When sampling ${\bm{a}}^t$, we avoid selecting the previous concatenation ${\bm{a}}^{t-1}$ and the all-zero vector (i.e., selecting no embedding). If ${\bm{a}}^t$ is in the dictionary ${\mathbb{D}}$, we compare the $R_t$ with the value in the dictionary and keep the higher one.

\begin{table*}[!ht]
\small
\centering
\setlength\tabcolsep{2.25pt}
\begin{tabular}{l||cccc|ccc|cccccccc}
\hlineB{4}
      & \multicolumn{4}{c|}{\bf \textsc{NER}}     & \multicolumn{3}{c|}{\bf \textsc{POS}}   & \multicolumn{8}{c}{\bf \textsc{AE}}          \\
\hhline{~||----|---|--------}
      & de       & en         & es   & nl  & Ritter & ARK    & TB-v2   & 14Lap  & 14Res   & 15Res   & 16Res & es   & nl   & ru   & tr   \\
\hline\hline
\bf \textsc{All}    & 83.1        & 92.4        & \textbf{88.9} & 89.8  & 90.6   & 92.1   & 94.6    & 82.7   & 88.5    & 74.2    & 73.2  & 74.6 & 75.0 & 67.1 & 67.5 \\
\bf \textsc{Random} & 84.0        & 92.6        & 88.8 & 91.9  & 91.3   & 92.6   & 94.6    & 83.6   & 88.1    & 73.5    & 74.7  & 75.0 & 73.6 & 68.0 & 70.0 \\
\bf \textsc{ACE}    & \textbf{84.2}        & \textbf{93.0}        & \textbf{88.9} & \textbf{92.1}  & \textbf{91.7}   & \textbf{92.8}   & \textbf{94.8}    & \textbf{83.9}   & \textbf{88.6}    & \textbf{74.9}    & \textbf{75.6}  & \textbf{75.7} & \textbf{75.3} & \textbf{70.6} & \textbf{71.1} \\
\hline
\end{tabular}
\begin{tabular}{l||c|cc|cccccc||c}
\hline
      & \multicolumn{1}{c|}{\bf \textsc{Chunk}} & \multicolumn{2}{c|}{\bf \textsc{DP}}   & \multicolumn{6}{c||}{\bf \textsc{SDP}}   & \multirow{2}{*}{\bf \textsc{Avg}}\\
\hhline{~||-|--|------||~}
      & CoNLL 2000         & UAS & LAS  & DM-ID & DM-OOD & PAS-ID & PAS-OOD & PSD-ID & PSD-OOD &           \\
\hline\hline
\bf \textsc{All}    & 96.7    & 96.7     & 95.1 & 94.3  & 90.8   & \textbf{94.6}   & 92.9    & 82.4   & 81.7    & 85.3      \\
\bf \textsc{Random} & 96.7    & 96.8    & 95.2 & 94.4  & 90.8   & \textbf{94.6}   & 93.0    & 82.3   & 81.8    & 85.7      \\
\bf \textsc{ACE}    & \textbf{96.8} & \textbf{96.9}        & \textbf{95.3} & \textbf{94.5}  & \textbf{90.9}   & 94.5   & \textbf{93.1}    & \textbf{82.5}   & \textbf{82.1}    & \textbf{86.2}     \\
\hlineB{4}
\end{tabular}
\caption{Comparison with concatenating all embeddings and random search baselines on 6 tasks.}
\label{tab:baseline}
\end{table*}

\section{Experiments}
We use ISO 639-1 language codes to represent languages in the table\footnote{\url{https://en.wikipedia.org/wiki/List_of_ISO_639-1_codes}}.

\subsection{Datasets and Configurations}
\label{sec:datasets}
To show ACE's effectiveness, we conduct extensive experiments on a variety of structured prediction tasks varying from syntactic tasks to semantic tasks. The tasks are named entity recognition (NER), Part-Of-Speech (POS) tagging, Chunking, Aspect Extraction (AE), Syntactic Dependency Parsing (DP) and Semantic Dependency Parsing (SDP).
The details of the 6 structured prediction tasks in our experiments are shown in below:
\begin{itemize}[leftmargin=*]
    \item {\bf NER}: We use the corpora of 4 languages from the CoNLL 2002 and 2003 shared task \citep{tjong-kim-sang-2002-introduction,tjong-kim-sang-de-meulder-2003-introduction} with standard split.
    \item {\bf POS Tagging}: We use three datasets, Ritter11-T-POS \citep{ritter-etal-2011-named}, ARK-Twitter \citep{gimpel-etal-2011-part,owoputi-etal-2013-improved} and Tweebank-v2 \citep{liu-etal-2018-parsing} datasets (Ritter, ARK and TB-v2 in simplification). We follow the dataset split of \citet{nguyen2020bertweet}.
    \item {\bf Chunking}: We use CoNLL 2000 \citep{tjong-kim-sang-buchholz-2000-introduction} for chunking. Since there is no standard development set for CoNLL 2000 dataset, we split 10\% of the training data as the development set.
    \item {\bf Aspect Extraction}: Aspect extraction is a subtask of aspect-based sentiment analysis \citep{pontiki-etal-2014-semeval,pontiki-etal-2015-semeval,pontiki-etal-2016-semeval}. The datasets are from the laptop and restaurant domain of SemEval 14, restaurant domain of SemEval 15 and restaurant domain of  SemEval 16 shared task (14Lap, 14Res, 15Res and 16Res in short). Additionally, we use another 4 languages in the restaurant domain of SemEval 16 to test our approach in multiple languages. We randomly split 10\% of the training data as the development set following \citet{li-etal-2019-exploiting}.
    \item {\bf Syntactic Dependency Parsing}: We use Penn Tree Bank (PTB) 3.0 with the same dataset pre-processing as \citep{ma-etal-2018-stack}.
    \item {\bf Semantic Dependency Parsing}: We use DM, PAS and PSD datasets for semantic dependency parsing \citep{oepensemeval} for the SemEval 2015 shared task \citep{oepen2015semeval}. The three datasets have the same sentences but with different formalisms. We use the standard split for SDP. In the split, there are in-domain test sets and out-of-domain test sets for each dataset.
\end{itemize}
Among these tasks, NER, POS tagging, chunking and aspect extraction are sequence-structured outputs while dependency parsing and semantic dependency parsing are the graph-structured outputs. POS Tagging, chunking and DP are syntactic structured prediction tasks while NER, AE, SDP are semantic structured prediction tasks.

We train the controller for $30$ steps and save the task model with the highest accuracy on the development set as the final model for testing. Please refer to Appendix \ref{app:detail} for more details of other settings

\subsection{Embeddings}
\paragraph{Basic Settings:} For the candidates of embeddings on English datasets, we use the language-specific model for ELMo, Flair, base BERT, GloVe word embeddings, fastText word embeddings, non-contextual character embeddings \citep{lample-etal-2016-neural}, multilingual Flair (M-Flair), M-BERT and XLM-R embeddings. The size of the search space in our experiments is $2^{11}{-}1{=}2047$\footnote{Flair embeddings have two models (forward and backward) for each language.}. For language-specific models of other languages, please refer to Appendix \ref{sec:embed} for more details. In AE, there is no available Russian-specific BERT, Flair and ELMo embeddings and there is no available Turkish-specific Flair and ELMo embeddings. We use the corresponding English embeddings instead so that the search spaces of these datasets are almost identical to those of the other datasets. All embeddings are fixed during training except that the character embeddings are trained over the task. The empirical results are reported in Section \ref{sec:exp:base}.

\paragraph{Embedding Fine-tuning:} A usual approach to get better accuracy is fine-tuning transformer-based embeddings. In sequence labeling, most of the work follows the fine-tuning pipeline of BERT that connects the BERT model with a linear layer for word-level classification.
However, when multiple embeddings are concatenated, fine-tuning a specific group of embeddings becomes difficult because of complicated hyper-parameter settings and massive GPU memory consumption.
To alleviate this problem, we first fine-tune the transformer-based embeddings over the task and then concatenate these embeddings together with other embeddings in the basic setting to apply ACE. The empirical results are reported in Section \ref{sec:exp:finetune}.

\subsection{Results}
We use the following abbreviations in our experiments: \textbf{UAS}: Unlabeled Attachment Score; \textbf{LAS}: Labeled Attachment Score; \textbf{ID}: In-domain test set; \textbf{OOD}: Out-of-domain test set. We use language codes for languages in NER and AE.

\subsubsection{Comparison With Baselines} \label{sec:exp:base}
To show the effectiveness of our approach, we compare our approach with two strong baselines. For the first one, we let the task model learn by itself the contribution of each embedding candidate that is helpful to the task. We set ${\bm{a}}$ to all-ones (i.e., the concatenation of all the embeddings) and train the task model ({\tt All}). The linear layer weight ${\bm{W}}$ in Eq. \ref{eq:linear} reflects the contribution of each candidate. For the second one, we use the random search ({\tt Random}), a strong baseline in NAS \citep{li2020random}. For {\tt Random}, we run the same maximum iteration as in ACE. For the experiments, we report the averaged accuracy of 3 runs. Table \ref{tab:baseline} shows that ACE outperforms both baselines in 6 tasks over 23 test sets with only two exceptions. Comparing {\tt Random} with {\tt All}, {\tt Random} outperforms {\tt All} by 0.4 on average and surpasses the accuracy of {\tt All} on 14 out of 23 test sets, which shows that concatenating all embeddings may not be the best solution to most structured prediction tasks.
In general, searching for the concatenation for the word representation is essential in most cases, and our search design can usually lead to better results compared to both of the baselines.

\begin{table*}[t]
\small
\centering
\begin{tabular}{l|ccccc||l|ccc}
\hlineB{4}
 & \multicolumn{5}{c||}{\bf \textsc{NER}} &       &\multicolumn{3}{c}{\bf \textsc{POS}}\\
 \hhline{~|-----||~|---}
  & de   & de$_\text{06}$ & en   & es   & nl   & & Ritter & ARK  & TB-v2 \\
\hline\hline
\citet{baevski-etal-2019-cloze}      & -    & -    & 93.5 & -    & -    & \citet{owoputi-etal-2013-improved} & 90.4   & 93.2 & 94.6  \\
\citet{strakova-etal-2019-neural}   & 85.1 & -    & 93.4 & 88.8 & 92.7     & \citet{gui-etal-2017-part}          & 90.9   & -    & 92.8  \\
\citet{yu-etal-2020-named} & 86.4 & 90.3 & 93.5 & 90.3 & 93.7 & \citet{gui-etal-2018-transferring} & 91.2   & 92.4 & -     \\
\citet{yamada-etal-2020-luke} & - & - & 94.3 & - & - & \citet{nguyen2020bertweet}         & 90.1   & 94.1 & 95.2  \\
\hline
XLM-R+Fine-tune & 87.7    & 91.4 & 94.1 & 89.3    & 95.3     & XLM-R+Fine-tune    & 92.3   & 93.7 & 95.4 \\
ACE+Fine-tune & \textbf{88.3} & \textbf{91.7} & \textbf{94.6} & \textbf{95.9} & \textbf{95.7} & ACE+Fine-tune  & \textbf{93.4}   & \textbf{94.4} & \textbf{95.8} \\
\hlineB{4}
\end{tabular}
\caption{Comparison with state-of-the-art approaches in NER and POS tagging. $^{\dagger}$: Models are trained on both train and development set.}
\label{tab:ner_pos}
\end{table*}

\begin{table*}[t]
\small
\centering
\setlength\tabcolsep{4pt}
\begin{tabular}{l|c||l|cccccccc}
\hlineB{4}
 & \multicolumn{1}{c||}{\bf \textsc{Chunk}} & &\multicolumn{8}{c}{\bf \textsc{AE}}        \\
\hhline{~|-||~|--------}
 &  CoNLL 2000 & & 14Lap  & 14Res & 15Res & 16Res & es   & nl   & ru   & tr   \\
\hline\hline
\citet{akbik-etal-2018-contextual}  & 96.7    & \citet{xu-etal-2018-double}$^\dagger$ & 84.2   & 84.6 & 72.0  & 75.4  & -    & -    & -    & -    \\
\citet{clark-etal-2018-semi}        & 97.0    & \citet{xu-etal-2019-bert}          & 84.3   & -    & -     & 78.0  & -    & -    & -    & -    \\
\citet{liu-etal-2019-gcdt}              & \textbf{97.3}    & \citet{wang-etal-2020-structure}   & -      & -    & -     & 72.8  & 74.3 & 72.9 & 71.8 & 59.3 \\
\citet{chen-etal-2020-seqvat} & 95.5       & \citet{wei-etal-2020-dont}         & 82.7   & 87.1 & 72.7  & 77.7  & -    & -    & -    & -    \\
\hline
XLM-R+Fine-tune  & 97.0 & XLM-R+Fine-tune & 85.9  & 90.5  & 76.4  & 78.9  & 77.0 & 77.6 & 77.7 & 74.1  \\
ACE+Fine-tune  &  \textbf{97.3}    & ACE+Fine-tune   & \textbf{87.4}   & \textbf{92.0} & \textbf{80.3}  & \textbf{81.3}  & \textbf{79.9} & \textbf{80.5} & \textbf{79.4} & \textbf{81.9} \\
\hlineB{4}
\end{tabular}
\caption{Comparison with state-of-the-art approaches in chunking and aspect extraction. $^\dagger$: We report the results reproduced by \citet{wei-etal-2020-dont}.}
\label{tab:chunk_ae}
\end{table*}

\begin{table*}[t]
\small
\centering
\begin{tabular}{l|cc||l|cccccc}
\hlineB{4}
& \multicolumn{2}{c||}{\bf \textsc{DP}}     &            & \multicolumn{6}{c}{\bf \textsc{SDP}}       \\
\hhline{~|--||~|------}
& \multicolumn{2}{c||}{\bf \textsc{PTB}}     &            & \multicolumn{2}{c}{\bf \textsc{DM}} & \multicolumn{2}{c}{\bf \textsc{PAS}} & \multicolumn{2}{c}{\bf \textsc{PSD}}       \\
     & UAS     & LAS  &            & ID   & OOD & ID  & OOD & ID  & OOD \\
\hline\hline
\citet{zhou-zhao-2019-head}$^\dagger$            & 97.2    & 95.7 &\citet{he2019establishing}$^\ddagger$      &  94.6 & 90.8    & 96.1 & 94.4     & 86.8 & 79.5  \\
\hhline{~|~~||-|------}
\citet{mrini-etal-2020-rethinking}$^\dagger$ & 97.4 & 96.3 &  D \& M \shortcite{dozat-manning-2018-simpler} &93.7 & 88.9    & 93.9 & 90.6     & 81.0 & 79.4  \\
\hhline{-|--||~|~~~~~~}
\citet{li2020global}       & 96.6    & 94.8 &    \citet{wang-etal-2019-second}  & 94.0 & 89.7    & 94.1 & 91.3     & 81.4 & 79.6      \\
 \citet{zhang-etal-2020-efficient}       & 96.1    & 94.5 &      \citet{jia-etal-2020-semi}  &93.6 & 89.1    & - & -    & - & -  \\
  \citet{wang-tu-2020-second} & 96.9    & 95.3& F \& G \shortcite{fernandez-gonzalez-gomez-rodriguez-2020-transition} & 94.4 & 91.0    & 95.1 & 93.4     & 82.6 & 82.0     \\
\hline
XLNET+Fine-tune           & 97.0        & 95.6   &XLNet+Fine-tune    & 94.2       & 90.6      & 94.8       & 93.4       & 82.7       & 81.8\\
ACE+Fine-tune  & \textbf{97.2}    & \textbf{95.8} & ACE+Fine-tune  & \textbf{95.6} & \textbf{92.6}    & \textbf{95.8} & \textbf{94.6}     & \textbf{83.8} & \textbf{83.4}    \\
\hlineB{4}
\end{tabular}
\caption{Comparison with state-of-the-art approaches in DP and SDP. $^\dagger$: For reference, they additionally used constituency dependencies in training. We also find that the PTB dataset used by \citet{mrini-etal-2020-rethinking} is not identical to the dataset in previous work such as \citet{zhang-etal-2020-efficient} and \citet{wang-tu-2020-second}. $^\ddagger$: For reference, we confirmed with the authors of \citet{he2019establishing} that they used a different data pre-processing script with previous work.}.
\label{tab:dp_sdp}
\end{table*}

\subsubsection{Comparison With State-of-the-Art approaches} \label{sec:exp:finetune}
As we have shown, ACE has an advantage in searching for better embedding concatenations. We further show that ACE is competitive or even stronger than state-of-the-art approaches. We additionally use XLNet \citep{yang2019xlnet} and RoBERTa as the candidates of ACE.  In some tasks, we have several additional settings to better compare with previous work. In NER, we also conduct a comparison on the revised version of German datasets in the CoNLL 2006 shared task \citep{buchholz-marsi-2006-conll}. Recent work such as \citet{yu-etal-2020-named} and \citet{yamada-etal-2020-luke} utilizes document contexts in the datasets. We follow their work and extract document embeddings for the transformer-based embeddings. Specifically, we follow the fine-tune process of \citet{yamada-etal-2020-luke} to fine-tune the transformer-based embeddings over the document except for BERT and M-BERT embeddings. For BERT and M-BERT, we follow the document extraction process of \citet{yu-etal-2020-named} because we find that the model with such document embeddings is significantly stronger than the model trained with the fine-tuning process of \citet{yamada-etal-2020-luke}. In SDP, the state-of-the-art approaches used POS tags and lemmas as additional word features to the network. We add these two features to the embedding candidates and train the embeddings together with the task. We use the fine-tuned transformer-based embeddings on each task instead of the pretrained version of these embeddings as the candidates.\footnote{Please refer to Appendix for more details about the embeddings.

We additionally compare with fine-tuned XLM-R model for NER, POS tagging, chunking and AE, and compare with fine-tuned XLNet model for DP and SDP, which are strong fine-tuned models in most of the experiments. Results are shown in Table \ref{tab:ner_pos}, \ref{tab:chunk_ae}, \ref{tab:dp_sdp}.
Results show that ACE with fine-tuned embeddings achieves state-of-the-art performance in all test sets, which shows that finding a good embedding concatenation helps structured prediction tasks. We also find that ACE is stronger than the fine-tuned models, which shows the effectiveness of concatenating the fine-tuned embeddings\footnote{We compare ACE with other fine-tuned embeddings in Appendix.}.

\begin{filecontents}{max_dev.dat}
episode	ACE	Random
1	96.2	95.96666667
2	96.21	96.00333333
3	96.37666667	96.00333333
4	96.37666667	96.05
5	96.39	96.05
6	96.39	96.05
7	96.39333333	96.05
8	96.39333333	96.17666667
9	96.40333333	96.18
10	96.43666667	96.18
11	96.52666667	96.23333333
12	96.52666667	96.23333333
13	96.52666667	96.23333333
14	96.52666667	96.28666667
15	96.52666667	96.28666667
16	96.53666667	96.28666667
17	96.53666667	96.3
18	96.53666667	96.35666667
19	96.55333333	96.35666667
20	96.55333333	96.35666667
21	96.59	96.44
22	96.61	96.44
23	96.64	96.44
24	96.64	96.44
25	96.64	96.44
26	96.64	96.44
27	96.64333333	96.44
28	96.64333333	96.45
29	96.64333333	96.45
30	96.64333333	96.45
\end{filecontents}

\begin{filecontents}{avg_dev.dat}
episode	ACE	Random
1	96.2	95.96666667
2	95.82666667	93.76333333
3	95.95333333	93.82333333
4	95.75666667	95.99666667
5	95.83333333	94.23333333
6	95.39	95.95333333
7	96.16333333	95.32
8	95.70333333	94.94666667
9	95.86	95.82333333
10	96.18666667	95.07333333
11	96.41666667	95.85666667
12	95.57	92.43
13	96.19	95.75333333
14	96.42666667	96.18666667
15	96.3	95.67333333
16	96.46	95.46
17	96.2	95.86666667
18	96.46	96.12333333
19	96.45333333	94.70333333
20	96.51	96.08
21	96.55333333	96.23666667
22	96.44	93.74
23	96.49	95.50333333
24	96.46333333	96.29
25	96.49666667	94.07333333
26	96.54666667	93.87
27	96.6	94.09333333
28	96.55	96.12666667
29	96.56666667	96.06666667
30	96.51333333	95.07333333
\end{filecontents}

\begin{figure*}[t!]
\begin{minipage}{1.0\linewidth}
\centering
\begin{tikzpicture}

    \begin{axis}[
        xshift=-6cm,

        name=ner,

        width=0.5\textwidth,
        height=0.3\textwidth,

        ylabel=Best Accuracy,
        legend columns=2,
        legend pos=south east,
        legend style={font=\small},
        tick label style={font=\small},

        ylabel style={font=\small,yshift=-0.2cm},
        ]
        \addplot[red] table[x=episode,y=ACE] {max_dev.dat};
        \addplot[blue,dashed] table[x=episode,y=Random] {max_dev.dat};
        \legend{ACE, Random}
    \end{axis}
    \begin{axis}[

        at={(ner.south west)},
        xshift=8.0cm,

        width=0.5\textwidth,
        height=0.3\textwidth,

        ylabel=Sample Accuracy,
        legend columns=2,
        legend pos=south east,
        legend style={font=\small},
        tick label style={font=\small},

        ytick={91,92,93,94,95,96},
        ylabel style={font=\small,yshift=-0.5cm},
        ]
        \addplot[red] table[x=episode,y=ACE] {avg_dev.dat};
        \addplot[blue,dashed] table[x=episode,y=Random] {avg_dev.dat};
        \legend{ACE, Random}
    \end{axis}
\end{tikzpicture}
\caption{Comparing the efficiency of random search ({\tt Random}) and ACE. The x-axis is the number of time steps. The left y-axis is the averaged best validation accuracy on CoNLL English NER dataset. The right y-axis is the averaged validation accuracy of the current selection.}
\label{fig:dev_curve}
\end{minipage}
\end{figure*}

\section{Analysis}

\subsection{Efficiency of Search Methods}
To show how efficient our approach is compared with the random search algorithm, we compare the algorithm in two aspects on CoNLL English NER dataset. The first aspect is the best development accuracy during training. The left part of Figure \ref{fig:dev_curve} shows that ACE is consistently stronger than the random search algorithm in this task. The second aspect is the searched concatenation at each time step. The right part of Figure \ref{fig:dev_curve} shows that the accuracy of ACE gradually increases and gets stable when more concatenations are sampled.

\subsection{Ablation Study on Reward Function Design}
\label{sec:ablation}
To show the effectiveness of the designed reward function, we compare our reward function (Eq. \ref{eq:reward}) with the reward function without discount factor (Eq. \ref{eq:reward_a}) and the traditional reward function (reward term in Eq. \ref{eq:rl}). We sample 2000 training sentences on CoNLL English NER dataset for faster training and train the controller for $50$ steps. Table \ref{tab:ablation} shows that both the discount factor and the binary vector $|{\bm{a}}^t-{\bm{a}}^{i}|$ for the task are helpful in both development and test datasets.
\begin{table}[h!]
\small
\centering
\begin{tabular}{l||cc}
\hlineB{4}
 & \bf \textsc{Dev} &   \bf \textsc{Test}\\
\hline
ACE & \textbf{93.18} & \textbf{90.00}\\
No discount (Eq. \ref{eq:reward_a}) &  92.98 & 89.90\\
Simple (Eq. \ref{eq:rl}) & 92.89 & 89.82 \\
\hlineB{4}
\end{tabular}
\caption{Comparison of reward functions.}
\label{tab:ablation}
\end{table}

\begin{table}[!ht]
\small
\centering
\setlength\tabcolsep{1.2pt}
\begin{tabular}{l||cccccccc}
\hlineB{4}
             & \multirow{2}{*}{\bf \textsc{NER}}  & \multirow{2}{*}{\bf \textsc{POS}}  &  \multirow{2}{*}{\bf \textsc{AE}}   & \multirow{2}{*}{\bf \textsc{CHK}} &\multicolumn{2}{c}{\bf \textsc{DP}} & \multicolumn{2}{c}{\bf \textsc{SDP}}  \\
&  &  & &   & UAS & LAS & ID  & OOD \\
\hline
\hline
All          & 92.4 & 90.6 & 73.2 & 96.7  & 96.7   & 95.1   & 94.3 & 90.8    \\
Random          & 92.6 & 91.3  & 74.7 & 96.7 & 96.8   & 95.2   & 94.4 & 90.8    \\
ACE          & \textbf{93.0} & \textbf{91.7} & \textbf{75.6}& \textbf{96.8}   & \textbf{96.9}   & \textbf{95.3}   & \textbf{94.5} & \textbf{90.9}   \\
All+Weight & 92.7  & 90.4 & 73.7 & 96.7 & 96.7 & 95.1 & 94.3 & 90.7 \\
Ensemble & 92.2 & 90.6  & 68.1 & 96.5 & 96.1   &94.3   & 94.1 & 90.3    \\
Ensemble$_{\text{dev}}$ & 92.2 & 90.8   & 70.2 & 96.7 & 96.8   &95.2   & 94.3 & 90.7    \\
\hline
Ensemble$_{\text{test}}$ & 92.7 & 91.4  & 73.9 & 96.7 & 96.8   &95.2   & 94.4 & 90.8    \\
\hlineB{4}
\end{tabular}
\caption{A comparison among \texttt{All}, \texttt{Random}, ACE, \texttt{All+Weight} and \texttt{Ensemble}. \textbf{\sc CHK}: chunking.}
\label{tab:ensemble}
\end{table}

\subsection{Comparison with Embedding Weighting \& Ensemble Approaches}
\label{sec:ensemble}
We compare ACE with two more approaches to further show the effectiveness of ACE. One is a variant of \texttt{All}, which uses a weighting parameter ${\bm{b}}=[b_1, \cdots, b_l, \cdots, b_L]$ passing through a sigmoid function to weight each embedding candidate. Such an approach can explicitly learn the weight of each embedding in training instead of a binary mask. We call this approach \texttt{All+Weight}. Another one is model ensemble, which trains the task model with each embedding candidate individually and uses the trained models to make joint prediction on the test set. We use voting for ensemble as it is simple and fast. For sequence labeling tasks, the models vote for the predicted label at each position. For DP, the models vote for the tree of each sentence. For SDP, the models vote for each potential labeled arc. We use the confidence of model predictions to break ties if there are more than one agreement with the same counts. We call this approach \texttt{Ensemble}. One of the benefits of voting is that it combines the predictions of the task models efficiently without any training process. We can search all possible $2^L{-}1$ model ensembles in a short period of time through caching the outputs of the models. Therefore, we search for the best ensemble of models on the development set and then evaluate the best ensemble on the test set (\texttt{Ensemble$_{\texttt{dev}}$}). Moreover, we additionally search for the best ensemble on the test set for reference (\texttt{Ensemble$_{\texttt{test}}$}), which is the upper bound of the approach. We use the same setting as in Section \ref{sec:exp:base} and select one of the datasets from each task. For NER, POS tagging, AE, and SDP, we use CoNLL 2003 English, Ritter, 16Res, and DM datasets, respectively. The results are shown in Table \ref{tab:ensemble}. Empirical results show that ACE outperforms all the settings of these approaches and even \texttt{Ensemble$_\texttt{test}$}, which shows the effectiveness of ACE and the limitation of ensemble models. \texttt{All}, \texttt{All+Weight} and \texttt{Ensemble$_\texttt{dev}$} are competitive in most of the cases and there is no clear winner of these approaches on all the datasets. These results show the strength of embedding concatenation. Concatenating the embeddings incorporates information from all the embeddings and forms stronger word representations for the task model, while in model ensemble, it is difficult for the individual task models to affect each other.

\section{Discussion: Practical Usability of ACE}
Concatenating multiple embeddings is a commonly used approach to improve accuracy of structured prediction. However, such approaches can be computationally costly as multiple language models are used as input. ACE is more practical than concatenating all embeddings as it can remove those embeddings that are not very useful in the concatenation. Moreover, ACE models can be used to guide the training of weaker models through techniques such as knowledge distillation in structured prediction \cite{kim-rush-2016-sequence,kuncoro-etal-2016-distilling,wang-etal-2020-structure,wang2020structural}, leading to models that are both stronger and faster.

\section{Conclusion}
In this paper, we propose Automated Concatenation of Embeddings, which automatically searches for better embedding concatenation for structured prediction tasks. We design a simple search space and use the reinforcement learning with a novel reward function to efficiently guide the controller to search for better embedding concatenations. We take the change of embedding concatenations into the reward function design and show that our new reward function is stronger than the simpler ones. Results show that ACE outperforms strong baselines. Together with fine-tuned embeddings, ACE achieves state-of-the-art performance in 6 tasks over 21 datasets.

\subsubsection*{Acknowledgments}
This work was supported by the National Natural Science Foundation of China (61976139) and by Alibaba Group through Alibaba Innovative Research Program. We thank Chengyue Jiang for his comments and suggestions on writing.

\bibliographystyle{acl_natbib}


\section{Introduction}
\IEEEPARstart{F}{iber} nonlinearity compensation \cite{winzer}--\cite{am1} and probabilistic shaping \cite{ps2}--\cite{KarimJLT} have been proven to be effective means to increase the spectral efficiency of optical communication systems.
Digital back propagation (DBP) is considered as the benchmark nonlinear compensation technique due to its high performance and accuracy, when applied with small step size \cite{ip}. In wavelength-division multiplexing (WDM) transmission systems, multi-channel DBP provides the best performance and mitigates both intra-channel and inter-channel nonlinear effects \cite{mcdbp}. However, multi-channel DBP is impractical for real-time implementation due to its high complexity, its requirement for high-speed analog-to-digital converters (ADC), and also the unavailability of the information of the adjacent WDM channels.

Probabilistic shaping based on distribution matching (DM) via constant composition (CC) \cite{ps0}, multiset partitioning \cite{mp2}, sphere shaping via shell mapping \cite{R1}, and enumerative sphere shaping (ESS) \cite{KarimJLT, KarimECOC} has been considered in optical communication systems. A combination of probabilistic shaping and DBP have been also investigated in \cite{psdbp}. In \cite{snl}, and recently in \cite{KarimJLT} and \cite{m2}, it has been shown that short blocklength shaping based on sphere shaping and CCDM provides a nonlinear tolerance gain in comparison with uniform signaling and with long blocklengths shaping. It is also known that ESS has a lower rate loss than CCDM at finite block lengths, which translates into a higher shaping gain \cite{gultekin2,GISIT}. This makes short blocklength ESS an interesting approach for achieving both shaping gain and nonlinear tolerance.

In this paper, we propose to exploit the nonlinear tolerance gain that short block length ESS provides \cite{KarimJLT, KarimECOC}, as a way to avoid the use of high complexity nonlinearity compensation techniques like DBP. We investigate the performance of uniform signaling with fiber nonlinearity compensation techniques and ESS with and without nonlinearity compensation, and explore the possibility of complexity reduction in this context.
We compare the performance of ESS with different blocklengths to uniform signaling with nonlinearity compensation via single-channel DBP.
We also consider the case of ESS with Volterra based nonlinear equalization (VNLE) \cite{VNLE1, VNLE2}, which reduces the complexity of nonlinear compensation by half, when compared to DBP \cite{VNLE1}. An evaluation of the complexity and storage of the proposed shaping and nonlinearity compensation approaches is also performed.

For a dense WDM system, we show that ESS provides better performance in terms of finite length bit-metric decoding (BMD) rate when compared to uniform with DBP for different number of steps per span. In terms of nonlinear tolerance, ESS exhibits the highest effective signal-to-noise ratio (SNR) at the shortest block length. In this context, ESS with VNLE and ESS with DBP applied at one step per span exceeds the performance of uniform with DBP per span and with $4$ steps per span DBP, respectively.
ESS also has lower complexity than DBP. However, ESS introduces additional latency and storage requirements. These results identify the potential for complexity reduction when considering short blocklength ESS instead of fiber nonlinearity compensation via DBP.
\section{System model and performance metrics}\label{sec:SP}
\subsection{System Model}\label{sec:SPmodel}
\begin{figure*}[tbp]
	\centering
		\includegraphics[width=0.75\linewidth]{0}
		\caption{Transmission diagram.

		$N_s$: number of spans. Dashed box is used for only CD compensation. Green color and cyan color correspond to uniform and shaping signaling respectively.
	}
		\label{fig:2}
		\vspace{-0.3cm}
\end{figure*}

The block diagram of the considered system is shown in  Fig.~\ref{fig:2}. The information bits are shaped based on ESS and passed through a rate $R_\text{c}$ low-density parity-check (LDPC) encoder following the probabilistic amplitude shaping (PAS) framework \cite{pas}.
The ESS shaper generates bounded energy sequences by fixing a maximum energy constraint \cite{gultekin2, KarimJLT}. The shaping rate $R_\text{s}=k/N$~[bits/amp], where $k$ and $N$ are the number of input information bits and output amplitudes respectively, is obtained by adjusting the maximum energy constraint.
 The shaping rate used in this work is $R_\text{s}=1.85$ [bits/amp.]. The LDPC parity bits and a part of the information bits are used as sign bits for the shaped amplitudes, as explained in [4]. 64-quadrature amplitude modulation (QAM) (8-Pulse-amplitude modulation (PAM) for shaping signaling with amplitude alphabet cardinality $a=4$, applied per real dimension) is used as modulation format for both shaped and uniform signaling due to its high spectral efficiency and the considerable possible shaping gain, which increases with the modulation format. Different LDPC rates $R_\text{c} = 3/4$ and $R_\text{c}=4/5$ should be used for uniform and shaped signaling, respectively, to ensure a fair comparison with the same information rate $R =2.25$ [bits/1D-sym]. In this work we focus on the SNR and AIR performances.

At the receiver side, nonlinear compensation is applied via DBP or VNLE. We also consider the case of only linear chromatic dispersion (CD) compensation.

\vspace{-0.2cm}
\subsection{Performance Metrics}\label{sec:SPerf}
 It has been shown that short blocklengths shaping provides interesting performance in terms of nonlinear tolerance in the optical fiber channel \cite{KarimJLT, KarimECOC, m2}. Thus, suitable performance metrics like finite length BMD rate, taking into account the rate loss with respect to infinite blocklengths, and effective SNR are used in this work. The finite length BMD rate gives an indication of the overall performance including the nonlinear gain and the linear shaping gain, and the effective SNR measures only the nonlinear tolerance gain.
The finite length BMD rate is defined as \cite{ps2}
 \begin{equation}\label{gmi}
     \text{AIR}_{\textit{N}} = \underbrace{\left[  H(\boldsymbol{C}) - \sum_{i=1}^m H(C_i \mid Y) \right]}_{\text{BMD Rate}} - \underbrace{\left[H(A)-\frac{k}{N}\right]}_{\text{Rate loss}},
 \end{equation}
where $H(\cdot)$ denotes entropy, $m$ is the number of bits per symbol, and \textit{$A$} are the shaped amplitudes. $\boldsymbol {C}= ( C_1, C_2, ..., C_m )$ are the bit levels of the transmitted symbol, and $Y$ corresponds to the received symbol.
The rate loss corresponds to the gap between the entropy and the shaping rate, and vanishes at infinite blocklengths.

The effective SNR includes both amplified spontaneous emission (ASE) noise and nonlinear noise contributions. It is calculated per QAM symbol taking into account the probability of each constellation point, and defined as \cite{ps3}
\begin{align}\label{eff.snr}
\text{SNR}_{\text{eff}} = \frac{\mathbb{E}[|{X}|^2]}{\mathbb{E}[|{Y}-{X}|^2]},
\end{align}
where $\mathbb{E}[\cdot]$ represents expectation, and $X$ and $Y$ are the transmitted and received symbols respectively (see Fig.~\ref{fig:2}).

\section{Simulation setup and results}\label{sec:simus}
\subsection{Simulation Setup}\label{sec:simusset}
 We consider two simulation scenarios: dual-polarization single-channel system, where the simulation setup is shown in Fig.~\ref{fig:2}, and dual-polarization 11 WDM channels system to quantify the effect of intra-channel and inter-channel nonlinear effects, respectively.
We compare the performance of ESS with linear electronic dispersion compensation (EDC) and ESS with nonlinearity compensation via VNLE and DBP against uniform signaling with DBP applied with different number of steps per span. The VNLE is performed in parallel and applied once per span \cite{VNLE1}.
In this work, polarization mode dispersion and linear phase noise are neglected to focus on the impact of the nonlinear effects. When EDC only is performed, we assume an ideal compensation of the common phase rotation of the entire constellation due to nonlinearity.
We consider a dense WDM scenario with large symbol rate and high order modulation format to ensure a transmission with high data rate. The symbol rate is $45$ Gbaud. The modulation format is $64$ QAM. We use a root-raised cosine (RRC) filter with a roll-off factor $\rho = 0.1$.

We consider a dispersion unmanaged system with multi-span standard single-mode fiber (SSMF). Concerning the SSMF parameters, the attenuation coefficient is $\alpha=0.2~\mathrm{dB\cdot km^{-1}}$, the dispersion parameter is $D=17 ~\mathrm{ps \cdot nm^{-1} \cdot km^{-1}}$, and the nonlinear coefficient is \mbox{$\gamma=1.3~ \mathrm{W^{-1} \cdot km^{-1}}$}. The signal is amplified after each $L=80$~km span by an erbium-doped fiber amplifier (EDFA) with a $5$ dB noise figure and $16$ dB gain.
At the receiver side, the signal is passed by a channel selection, and the nonlinear compensation is applied after downsampling to $2$ samples/symbol.

\subsection{Simulation Results}\label{sec:simusres}
We firstly consider the dual-polarization single-channel system. In Fig.~\ref{fig:3}, we plot the effective SNR versus shaping blocklength $N$ for a transmission reach of $2800$ km at optimal input power. Fig.~\ref{fig:3} shows that uniform signaling with DBP applied at $8$ steps per span exhibits the best performance and the gain is $1.24$ dB, $2.55$ dB, and $3.96$ dB in comparison with DBP at $4$ steps per span, DBP applied per span and EDC-only, respectively. ESS exhibits its best performance at the shortest blocklength, and for $N=100$, it shows a gain of $0.22$~dB in comparison with uniform signaling. At the same blocklength, ESS with one step per span DBP exhibits a gain of $0.1$~dB in comparison with uniform with DBP applied per span.

\begin{figure}[!t]
		\includegraphics[width=0.83\linewidth]{1}

		\caption{Effective SNR vs. shaping blocklength $N$ for a single WDM channel.}
		\label{fig:3}
				\vspace{-0.2cm}
\end{figure}
Fig.~\ref{fig:4} shows the finite length BMD rate performances as a function of the shaping blocklength.
The highest performance of ESS with VNLE is obtained at $N=360$. It corresponds to the optimal trade-off between the nonlinear tolerance, which is inversely proportional to the blocklength, and the linear shaping gain, which increases with the blocklength. For this blocklength, ESS with nonlinearity compensation via VNLE provides similar performance to uniform with DBP applied per span, which allows a reduction of nonlinear compensation complexity by half \cite{am1}. DBP at $8$ steps per span exhibits the best performance and shows a gain of $0.52$~dB in comparison with DBP at $4$ steps per span, due to its higher accuracy.

\begin{figure}[!]
	\includegraphics[width=0.83\linewidth]{2}

		\caption{Finite length AIR vs. shaping blocklength $N$ for a single WDM channel.}
		\label{fig:4}
		\vspace{-0.2cm}
\end{figure}

Next, we extend the scenario to a $11$ dual-polarization WDM channels system with $45$ Gb in $50$ GHz grid configuration. The results for the middle WDM channel will be shown because it is the most affected by the nonlinear impairments.

In Fig.~\ref{fig:5}, the effective SNR is plotted as a function of the shaping block length for $2000$ km as transmission distance at optimal input power. It is observed that $64$-QAM uniform signaling with DBP at $8$ steps per span provides the best performance, and the gain is about $0.4$ dB, $0.18$ dB and $0.12$ dB in comparison with CD compensation, DBP applied per span and DBP at $4$ steps per span, respectively. Again, it is shown that ESS gives the best performance at the shortest blocklength, while at a blocklength $N=2400$, uniform signaling shows better performance than ESS. It is also observed that the perfromance gap between ESS and ESS with DBP applied per span is lower than the case of uniform and uniform with DBP applied per span. This can be explained by the fact that ESS with short bloclengths mitigates a part of the nonlinearity and also has different statistics than uniform signaling, which results on different behaviors of DBP for both signaling.
For $N=100$, ESS with DBP and ESS with VNLE, applied per span, provide gains of about $0.02$ dB and $0.01$~dB when compared to uniform signaling with DBP at $4$ per span and DBP per span, respectively. This means that in the presence of probabilistic shaping, by using short blocklength shaping, the complexity of the nonlinearity compensation can be significantly reduced in dense WDM transmission systems.
\begin{figure}[!t]

		\includegraphics[width=0.83\linewidth]{3}
		\caption{Effective SNR vs. shaping blocklength $N$ for $11$ WDM channels.}
		\label{fig:5}
						\vspace{-0.2cm}
\end{figure}

\begin{table*}[t]
\begin{centering}
\caption{Computational and storage complexity.}
\par\end{centering}
\begin{centering}
{\scriptsize{}}%
\begin{tabular}{|>{\centering}m{0.12\paperwidth}|>{\centering}m{0.2\paperwidth}|>{\centering}m{0.3\paperwidth}|>{\centering}m{0.1\paperwidth}|}
\cline{2-4} \cline{3-4} \cline{4-4}
\multicolumn{1}{>{\centering}m{0.12\paperwidth}|}{} & {\scriptsize{}No. of real-valued multiplications} & {\scriptsize{}No. of real-valued additions} & {\scriptsize{}Storage requirement}\tabularnewline
\hline
\centering{}{\scriptsize{}Uniform} & {\scriptsize{}$(8N_\text{FFT}\log_{2}\left(N_\text{FFT}\right)+8N_\text{FFT})/N_\text{sym}$} & {\scriptsize{}$(8N_\text{FFT}\log_{2}\left(N_\text{FFT}\right))/N_\text{sym}$} & {\scriptsize{}-}\tabularnewline
\hline

\centering{}{\scriptsize{}Uniform-DBP} & {\scriptsize{}$(8N_\text{s}sN_\text{FFT}\log_{2}\left(N_\text{FFT}\right)+21N_\text{s}sN_\text{FFT})/N_\text{sym}$} & {\scriptsize{}$(8N_\text{s}sN_\text{FFT}\log_{2}\left(N_\text{FFT}\right)+4N_\text{s}sN_\text{FFT})/N_\text{sym}$} & {\scriptsize{}-}\tabularnewline
\hline
\centering{}{\scriptsize{}ESS} & {\scriptsize{}$(8N_\text{FFT}\log_{2}\left(N_\text{FFT}\right)+8N_\text{FFT})/N_\text{sym}$} & {\scriptsize{}$(8aN_\text{sym}$+$8N_\text{FFT}\log_{2}\left(N_\text{FFT}\right))/N_\text{sym}$} & {\scriptsize{} $O(N^{2}\,\log\,N)$}\tabularnewline
\hline
\centering{}{\scriptsize{}ESS-DBP} & {\scriptsize{}$(8N_\text{s}sN_\text{FFT}\log_{2}\left(N_\text{FFT}\right)+21N_\text{s}sN_\text{FFT})/N_\text{sym}$} & {\scriptsize{}$8aN_\text{sym}$+$8N_\text{s}sN_\text{FFT}\log_{2}\left(N_\text{FFT}\right)+4N_\text{s}sN_\text{FFT})/N_\text{sym}$} & {\scriptsize{}$O(N^{2}\,\log\,N)$}\tabularnewline
\hline
\centering{}{\scriptsize{}ESS-VNLE} & {\scriptsize{}$(8N_\text{s}N_\text{FFT}\log_{2}\left(N_\text{FFT}\right)+8.5N_\text{s}N_\text{FFT}+4N_\text{FFT}+4N_\text{FFT}\log_{2}\left(N_\text{FFT}\right))/N_\text{sym}$} & {\scriptsize{}$(8aN_\text{sym}$+$8N_\text{s}N_\text{FFT}\log_{2}\left(N_\text{FFT}\right)+8N_\text{s}N_\text{FFT}+8N_\text{FFT}\log_{2}\left(N_\text{FFT}\right)+4N_\text{FFT})/N_\text{sym}$} & {\scriptsize{} $O(N^{2}\,\log\,N)$}\tabularnewline
\hline

\end{tabular}{\scriptsize\par}
\par\end{centering}
\vspace{1mm}

\centering{}{\scriptsize{}$N_\text{s}:$ Number of fiber spans, $s:$ Number
of steps per fiber span, $N_\text{FFT}:$ Fast Fourier transform size, $a$: amplitude alphabet cardinlaity, $N_\text{sym}:$ Total number of symbols.}{\scriptsize\par}
\end{table*}

In terms of the finite length BMD rate, as shown in Fig.~\ref{fig:6}, ESS with DBP applied per span exhibits the best performance. The shaping blocklength that provides the optimal trade-off between linear shaping gain and nonlinear tolerance is around $N=600$. It is also observed that DBP with a high number of steps per span, i.e., high accuracy, still shows lower performance than ESS with only linear CD compensation.
\begin{figure

			\includegraphics[width=0.83\linewidth]{4}
	\caption{Finite length AIR vs. shaping blocklength $N$ for $11$ WDM channel.}
	\label{fig:6}
\end{figure}
\vspace{-0.2cm}
\subsection{Complexity Analysis}

Table I summarizes the computational complexity and the required storage for the considered techniques. The computational complexity is evaluated for a 4-dimensional symbol (i.e., dual-polarization symbols). It is important to mention that the uniform signaling and ESS only cases require the CD compensation, and the significant portion of the complexity for such techniques comes from the CD compensation part. On the other hand, when these techniques (i.e., uniform signaling and ESS) are combined with DBP or VNLE, the CD compensation part is already included in the DBP and VNLE implementation. Bounded-precision ESS \cite{GISIT} is used in this work to reduce the storage requirements \cite{cmpx}. The CD compensation is implemented in frequency-domain using a fast Fourier transform (FFT)/Inverse-FFT method, as in \cite{VNLE1}. For both DBP and VNLE methods, we follow the same frequency-domain approach for CD compensation.
The ESS can be implemented with a smaller computational complexity than nonlinearity compensation due to its lower number of real-valued multiplications and additions. In addition, ESS complexity does not depend on the number of spans, unlike nonlinear compensation via DBP and VNLE.  However, its realization requires additional storage. There is a trade-off between the computational complexity and the required storage for the shaping and nonlinearity compensation techniques. With short blocklength ESS, the nonlinearity compensation complexity can be significantly reduced with increased performance, especially for the  WDM systems.
\section{Conclusion}\label{sec:con}
 We have investigated the performance of fiber nonlinearity compensation in comparison with finite blocklength ESS in single-channel and dense WDM transmission systems. We have shown that ESS exceeds the performance of uniform signaling with higher complexity nonlinearity compensation in terms of  finite length bit-metric decoding rate. Furthermore, shorth blocklengths ESS, which provide nonlinear tolerance gain, has lower complexity in terms of real-valued multiplications and additions than DBP, but it introduces storage requirements. This make short blocklength probabilistic shaping more suitable for high data rate dense WDM systems than nonlinear effects compensation via DBP.