\section{Introduction}
Conversational virtual assistants, such as Amazon Alexa, Google Home, and Apple Siri, have become increasingly popular in recent times. These systems can process queries from users and perform tasks such as playing music and finding locations. A core component in these systems is a task-oriented semantic parsing model that maps natural language expressions to structured representations containing intents and slots that describe the task to perform. For example, the expression \emph{Can you play some songs by Coldplay?} may be converted to \emph{Intent: PlaySong, Artist: Coldplay}, and the expression \emph{Turn off the bedroom light} may be converted to \emph{Intent: TurnOffLight, Device: bedroom}.

Task-oriented semantic parsing is traditionally approached as a joint intent classification and slot filling task. \citet{kamath2018survey} provide a comprehensive survey of models proposed to solve this task. Researchers have developed semantic parsers based on Recurrent Neural Networks \cite{mesnil2013investigation, liu2016attention, hakkani2016multi}, Convolutional Neural Networks \cite{xu2013convolutional, kim2014}, Recursive Neural Networks \cite{guo2014joint}, Capsule Networks \cite{sabour2017, zhang2019}, and slot-gated attention-based models \cite{goo2018slot}.

The current state-of-the-art models on SNIPS \cite{coucke2018snips}, ATIS \cite{price1990evaluation}, and Facebook TOP \cite{gupta2018semantic} datasets are all based on BERT-style \cite{bert, liu2019roberta} encoders and transformer architectures \cite{chen2019bert,castellucci2019multi,rongali}. It is challenging to deploy these large models on edge devices and enable the voice assistants to operate locally instead of relying on central cloud services, due to the limited memory budgets on these devices. However, there has been a growing push towards the idea of TinyAI \footnote{https://www.technologyreview.com/technology/tiny-ai/}.

In this paper, we aim to build space-efficient task-oriented semantic parsing models that produce near state-of-the-art performances by compressing existing large models. We propose to learn compositional code embeddings to significantly compress BERT-base and RoBERTa-base encoders with little performance loss. We further use ALBERT-base/large \cite{albert} and DistilBERT \cite{distilbert} to establish light baselines that achieve similar state-of-the-art performances, and apply the same code embedding technique. We show that our technique is complementary to the compression techniques used in ALBERT and DistilBERT. With all variants, we achieve 95.15\% $\sim$ 98.46\% embedding compression rates and 20.47\% $\sim$ 34.22\% encoder compression rates, with $>$97.5\% semantic parsing performance preservation.

\section{Related Compression Techniques}
\subsection{BERT Compression}
Many techniques have been proposed to compress BERT \cite{bert}. \citet{ganesh2020compressing} provide a survey on these methods. Most existing methods focus on alternative architectures in transformer layers or learning strategies.

In our work, we use DistilBERT and ALBERT-base as light pretrained language model encoders for semantic parsing. DistilBERT \cite{distilbert} uses distillation to pretrain a model that is 40\% smaller and 60\% faster than BERT-base, while retaining 97\% of its downstream performances. ALBERT \cite{albert} factorizes the embedding and shares parameters among the transformer layers in BERT and results in better scalability than BERT. ALBERT-xxlarge outperforms BERT-large on GLUE \cite{GLUEdataset}, RACE \cite{lai2017race}, and SQUAD \cite{rajpurkar2016squad} while using less parameters.

We use compositional code learning \cite{codebooks} to compress the model embeddings, which contain a substantial amount of model parameters. Previously ALBERT uses factorization to compress the embeddings. We find more compression possible with code embeddings.

\subsection{Embedding Compression}
Varied techniques have been proposed to learn compressed versions of non-contextualized word embeddings, such as, Word2Vec \cite{mikolov2013efficient} and GLoVE \cite{pennington2014glove}. \citet{spine} use denoising k-sparse autoencoders to achieve binary sparse intrepretable word embeddings. \citet{chen2016} achieve sparsity by representing the embeddings of uncommon words using sparse linear common combination of common words. \citet{word2bits} achieve compression by quantization of the word embeddings by using 1-2 bits per parameter. \citet{faruqui} use sparse coding in a dictionary learning setting to obtain sparse, non-negative word embeddings. \citet{raunak2017} achieve dense compression of word embeddings using PCA combined with a post-processing algorithm. \citet{codebooks} propose to represent word embeddings using compositional codes learnt directly in end-to-end fashion using neural networks. Essentially few common basis vectors are learnt and embeddings are reconstructed using their composition via a discrete code vector specific to each token embedding. This results in 98\% compression rate in sentiment analysis and 94\% - 99\% in machine translation tasks without performance loss with LSTM based models. All the above techniques are applied to embeddings such as WordVec and Glove, or LSTM models.

We aim to learn space-efficient embeddings for transformer-based models. We focus on compositional code embeddings \cite{codebooks} since they maintain the vector dimensions, do not require special kernels for calculating in a sparse or quantized space, can be finetuned with transformer-based models end-to-end, and achieve extremely high compression rate. \citet{chen2018learning} explores similar idea as \citet{codebooks} and experiment with more complex composition functions and guidances for training the discrete codes. \citet{chen2019differentiable} further show that end-to-end training from scratch of models with code embeddings is possible. Given various pretrained language models, we find that the method proposed by \citet{codebooks} is straightforward and perform well in our semantic parsing experiments.

\section{Method}
\subsection{Compositional Code Embeddings}
\citet{codebooks} apply additive quantization \cite{babenko2014additive} to learn compositional code embeddings to reconstruct pretrained word embeddings such as GloVe \cite{pennington2014glove}, or task-specific model embeddings such as those from an LSTM neural machine translation model. Compositional code embeddings $E^C$ for vocabulary $V$ consist of a set of $M$ codebooks $E_1^C$, $E_2^C$, ..., $E_M^C$, each with $K$ basis vectors of the same dimensionality $D$ as the reference embeddings $E$, and a discrete code vector ($C_w^1$, $C_w^2$, ..., $C_w^M$) for each token $w$ in the vocabulary. The final embedding for $w$ is composed by summing up the $C_w^i$th vector from the ith codebook as $E^C(C_w)=\sum_{i=1}^M E_i^C(C_w^i)$. Codebooks and discrete codes are jointly learned using the mean squared distance objective: $(C^*, E^{C*})=\mathop{\arg\min}_{C, E^C} \frac{1}{|V|}\sum_{w\in V}||E^C(C_w)-E(w)||^2.$ For learning compositional codes, the Gumbel-softmax reparameterization trick \cite{jang2016categorical, maddison2016concrete} is used for one-hot vectors corresponding to each discrete code.

\begin{table*}[t]
\centering
\resizebox{\textwidth}{!}{%
\begin{tabular}{lcccccccc}
\hline \textbf{Encoder} & \textbf{EncoderParam\# / Size} &\textbf{EmbParam\# / Size} & \textbf{SizeRatio}& \textbf{CCEmbSize} & \textbf{CCEncoderSize} & \textbf{EmbComp} & \textbf{EncoderComp}\\ \hline
RoBERTa-base            & 125.29M / 477.94MB  & 38.60M / 147.25MB   & 30.81\%   & 2.27MB    & 332.96MB & 98.46\%   & 30.33\% \\
BERT-base-uncased	    & 110.10M / 420.00MB  & 23.44M / 89.42MB    & 21.29\%   & 1.97MB    & 332.55MB & 97.80\%   & 20.82\% \\
DistilBERT-base-uncased	& 66.99M / 255.55MB   & 23.44M / 89.42MB    & 34.99\%   & 1.97MB    & 168.10MB & 97.80\%   & 34.22\% \\
ALBERT-large-v2         & 17.85M / 68.09MB    & 3.84M / 14.65MB     & 21.52\%   & 0.71MB    & 54.15MB  & 95.15\%   & 20.47\% \\
ALBERT-base-v2          & 11.81M / 45.05MB    & 3.84M / 14.65MB     & 32.52\%   & 0.71MB    & 31.11MB  & 95.15\%   & 30.94\% \\
\hline
\end{tabular}
}
\caption{\label{adv} Model compression with compositional code (``cc") embeddings. The embedding layers are compressed by more than 95\% with compositional code embeddings in all of the BERT variants.}
\end{table*}

\subsection{Transformer-Based Models with Compositional Code Embeddings}
In this work, we learn compositional code embeddings to reduce the size of the embeddings in pretrained contextualized language models. We extract the embedding tables from pretrained RoBERTa-base \cite{liu2019roberta}, BERT-base \cite{bert}, DistilBERT-base \cite{distilbert}, ALBERT-large-v2 and ALBERT-base-v2 \cite{albert} from the huggingface transformers library \cite{Wolf2019HuggingFacesTS} and follow the approach presented by \citet{codebooks} to learn the code embeddings. We then replace the embedding tables in the transformer models with the compositional code approximations and evaluate the compressed language models by finetuning on downstream tasks. When \citet{codebooks} feed compositional code embeddings into the LSTM neural machine translation model, they fix the embedding parameters and train the rest of the model from random initial values. In our experiments, we fix the discrete codes, initialize the transformer layers with those from the pretrained language models, initialize the task-specific output layers randomly, and finetune the codebook basis vectors with the rest of the non-discrete parameters.

\subsection{Size Advantage of Compositional Code Embeddings}
An embedding matrix $E\in \mathbb{R}^{|V|\times D}$ stored as 32-bit float point numbers, where $|V|$ is the vocabulary size and $D$ is the embedding dimension, requires $32|V|D$ bits. Its compositional code reconstruction requires $32MKD$ bits for $MK$ basis vectors, and $M\log_2K$ bits for codes of each of $|V|$ tokens. Since each discrete code takes an integer value in $[1, K]$, it can be represented using $\log_2K$ bits.

Table \ref{adv} illustrates the size advantage of compositional code embeddings for various pretrained transformer models \cite{Wolf2019HuggingFacesTS} used in our experiments. While the technique focuses on compressing the embedding table, it is compatible with other compression techniques for transformer models, including parameter sharing among transformer layers and embedding factorization used in ALBERT and distillation for learning DistilBERT. In our experiments, we apply the code learning technique to compress embeddings in five pretrained BERT variants by 95.15\% $\sim$ 98.46\% to build competitive but significantly lighter semantic parsing models.

\begin{table}
\centering
\resizebox{0.45\textwidth}{!}{%
\begin{tabular}{lccccc}
\hline \textbf{Dataset} & \textbf{Train} &\textbf{Valid} & \textbf{Test} & \textbf{\#Intent} & \textbf{\#Slot} \\ \hline
ATIS	& 4,478 & 500 & 893 & 26 & 83 \\
SNIPS	& 13,084 & 700 & 700 & 7 & 39 \\
Facebook TOP    & 31,279 & 4,462 & 9,042 & 25 & 36 \\
\hline
\end{tabular}
}
\caption{\label{table-data} Statistics for semantic parsing datasets.}
\end{table}

\section{Datasets}
Following \citet{rongali}, we evaluate our models on SNIPS \cite{coucke2018snips}, Airline Travel Information
System (ATIS) \cite{price1990evaluation}, and Facebook TOP \cite{gupta2018semantic} datasets for task-oriented semantic parsing (Table \ref{table-data}). For SNIPS and ATIS, we use the same train/validation/test split as \citet{goo2018slot}.

\section{Experiments and Analyses}
For transformer model training, we base our implementation on the huggingface transformers library v2.6.0 \cite{Wolf2019HuggingFacesTS}. We use the AdamW optimizer \cite{loshchilov2017decoupled} with 10\% warmup steps and linear learning rate decay to 0. Forr code embedding learning, we base our implementation on that of \citet{codebooks}. By default we learn code embeddings with 32 codebooks and 16 basis vectors per codebook. Unless otherwise specified, hyperparameters are found according to validation performances from one random run. We conduct our experiments on a mixture of Tesla M40, TITAN X, 1080 Ti, and 2080 Ti GPUs. We use exact match (EM) and intent accuracy as evaluation metrics. Exact match requires correct predictions for all intents and slots in a query, and is our primary metric.

\begin{table}[t]
\centering
\resizebox{0.5\textwidth}{!}{%
\begin{tabular}{lcccccc}
\hline \multicolumn{5}{l}{\textbf{Model}} & \textbf{EM} &\textbf{Intent}  \\ \hline
\multicolumn{5}{l}{Joint BiRNN \cite{hakkani2016multi}}	& 73.2 & 96.9 \\
\multicolumn{5}{l}{Attention BiRNN \cite{liu2016attention}}	&  74.1 & 96.7 \\
\multicolumn{5}{l}{Slot Gated Full Attention \cite{goo2018slot}}	& 75.5 & 97.0 \\
\multicolumn{5}{l}{CapsuleNLU \cite{zhang2019}}	&  80.9 & 97.3 \\
\multicolumn{5}{l}{BERT-Seq2Seq-Ptr \cite{rongali}}	& 86.3 & 98.3 \\
\multicolumn{5}{l}{RoBERTa-Seq2Seq-Ptr \cite{rongali}}	& 87.1 & 98.0 \\
\multicolumn{5}{l}{BERT-Joint \cite{castellucci2019multi}}	& 91.6 & \textbf{99.0} \\
\multicolumn{5}{l}{Joint BERT \cite{chen2019bert}}	& \textbf{92.8} & 98.6  \\
\hline \textbf{Ours} & \textbf{epo} & \textbf{lr} & \textbf{wd} & \textbf{EM-v} &\textbf{EM} & \textbf{Intent} \\ \hline
ALBERT-base	&   & 5e-5 & 0.05 & 90.71 & \textbf{91.29} & 98.86 \\
ALBERT-base\_cc & 1100 & 5e-5 & 0.01 & 90.00 & 89.14 & 98.14 \\\hline
ALBERT-large	&   & 3e-5 & 0.05 & 91.29 & \textbf{92.43} & 98.14 \\
ALBERT-large\_cc & 1100 & 2e-5 & 0.05 & 91.14 & \textbf{92.43} & 98.71 \\\hline
DistilBERT-base	&   & 3e-5 & 0.05 & 90.29 & 91.14 & 98.57 \\
DistilBERT-base\_cc	& 900 & 6e-5 & 0.01 & 90.14 & \textbf{91.24} & 98.43 \\\hline
BERT-base	&  & 3e-5 & 0.05 & 92.14 & \textbf{92.29} & 99.14  \\
BERT-base\_cc	& 900 & 6e-5  & 0.05 & 91.29 & 90.71 & 98.71 \\
\hline
\end{tabular}
}
\caption{\label{table-SNIPS} Results on SNIPS. ``cc" indicate models with code embeddings. ``epo" is the epoch number for offline code embedding learning. ``lr" and ``wd" are the peak learning rate and weight decay for whole model finetuning. ``EM-v", ``EM", ``Intent" indicate validation exact match, test exact match, and test intent accuracy.}
\end{table}

\subsection{SNIPS and ATIS}
We implement a joint sequence-level and token-level classification layer for pretrained transformer models. The intent probabilities are predicted as $\mathrm{y^i=softmax(W^ih_0+b^i)}$, where $\mathrm{h_0}$ is the hidden state of the [CLS] token. The slot probabilities for each token j are predicted as $\mathrm{y^s_j=softmax(W^sh_j+b^s)}$. We use the cross entropy loss to maximize $\mathrm{p(y^i|x)\prod p(y^s_j|x)}$ where j is the first piece-wise token for each word in the query. We learn code embeddings for \{500, 700, 900, 1100, 1300\} epochs. We train transformer models with original and code embeddings all for 40 epochs with batch size 16 and sequence length 128. Uncased BERT and DistilBERT perform better than the cased versions. We experiment with peak learning rate \{2e-5, 3e-5, ..., 6e-5\} and weight decay \{0.01, 0.05, 0.1\}. As shown in Table \ref{table-SNIPS} and \ref{table-ATIS}, we use different transformer encoders to establish strong baselines which achieve EM values that are within 1.5\% of the state-of-the-art.

On both datasets, models based on our compressed ALBERT-large-v2 encoder (54MB) perserves $>$99.6\% EM of the previous state-of-the-art model \cite{chen2019bert} which uses a BERT encoder (420MB). In all settings, our compressed encoders perserve $>$97.5\% EM of the uncompressed counterparts under the same training settings. We show that our technique is effective on a variety of pretrained transformer encoders.

\begin{table}[t]
\centering
\resizebox{0.5\textwidth}{!}{%
\begin{tabular}{lcccccc}
\hline \multicolumn{5}{l}{\textbf{Model}} & \textbf{EM} &\textbf{Intent}  \\ \hline
\multicolumn{5}{l}{Joint-BiRNN \cite{hakkani2016multi}}	& 80.7 & 92.6\\
\multicolumn{5}{l}{Attention-BiRNN \cite{liu2016attention}}	& 78.9 & 91.1\\
\multicolumn{5}{l}{Slot-Gated \cite{goo2018slot}}	& 82.2 & 93.6\\
\multicolumn{5}{l}{CapsuleNLU \cite{zhang2019}}	& 83.4 & 95.0\\
\multicolumn{5}{l}{BERT-Seq2Seq-Ptr \cite{rongali}}	& 86.4 & 97.4 \\
\multicolumn{5}{l}{RoBERTa-Seq2Seq-Ptr \cite{rongali}}	& 87.1 & 97.4\\
\multicolumn{5}{l}{BERT-Joint \cite{castellucci2019multi}}	& \textbf{88.2} & \textbf{97.8} \\
\multicolumn{5}{l}{Joint-BERT \cite{chen2019bert}}	& \textbf{88.2} & 97.5 \\
\hline \textbf{Ours} & \textbf{epo} & \textbf{lr} & \textbf{wd} & \textbf{EM-v} &\textbf{EM} & \textbf{Intent} \\ \hline
ALBERT-base	&   & 5e-5 & 0.05 & 93.4 & 86.90 & 97.42 \\
ALBERT-base\_cc & 900 & 6e-5 & 0.1 & 94.2 & \textbf{87.23} & 96.75 \\\hline
ALBERT-large	&   & 5e-5 & 0.05 & 93.8 & \textbf{88.02} & 97.54 \\
ALBERT-large\_cc & 1100 & 5e-5 & 0.1 & 94.0 & 87.91 & 97.54 \\\hline
DistilBERT-base	&   & 4e-5 & 0.05 & 93.6 & \textbf{88.13} & 97.42 \\
DistilBERT-base\_cc	& 1100 & 6e-5 & 0.05 & 93.2 & 87.12 & 97.54 \\\hline
BERT-base	&   &  4e-5 & 0.01 & 93.4 & \textbf{88.13} & 97.54 \\
BERT-base\_cc	& 700 & 6e-5  & 0.1 & 93.0 & 87.35 & 97.20 \\
\hline
\end{tabular}
}
\caption{\label{table-ATIS} Results on ATIS. Refer to the caption of Table \ref{table-SNIPS} for abbreviation explanations.}
\end{table}

\subsection{Facebook TOP}
\begin{table}[t]
\centering
\resizebox{0.5\textwidth}{!}{%
\begin{tabular}{lccc}
\hline \multicolumn{2}{l}{\textbf{Model}} & \textbf{EM} &\textbf{Intent}  \\ \hline
\multicolumn{2}{l}{RNNG \cite{gupta2018semantic}} & 78.51 & - \\
\multicolumn{2}{l}{Shift Reduce (SR) Parser} & 80.86 & - \\
\multicolumn{2}{l}{SR with ELMo embeddings} &83.93 & - \\
\multicolumn{2}{l}{SR ensemble + ELMo + SVMRank} & \textbf{87.25} & - \\
\multicolumn{2}{l}{BERT-Seq2Seq-Ptr \cite{rongali}}	& 83.13 & 97.91 \\
\multicolumn{2}{l}{RoBERTa-Seq2Seq-Ptr \cite{rongali}}	& 86.67 & \textbf{98.13} \\
\hline \textbf{Ours} & \textbf{EM-v} &\textbf{EM} & \textbf{Intent} \\ \hline
ALBERT-Seq2Seq-Ptr & 84.56 & \textbf{85.41} & 98.47 \\
ALBERT-Seq2Seq-Ptr\_cc & 83.48 & 84.42 & 98.05 \\ \hline
DistilBERT-Seq2Seq-Ptr & 84.25 & \textbf{85.12} & 98.50 \\
DistilBERT-Seq2Seq-Ptr\_cc & 82.76 & 83.42 & 98.09 \\ \hline
BERT-Seq2Seq-Ptr & 83.83 & \textbf{85.01} & 98.59 \\
BERT-Seq2Seq-Ptr\_cc & 82.36 & 83.34 & 98.25 \\ \hline
RoBERTa-Seq2Seq-Ptr	 & 85.00 & \textbf{85.67} & 98.59 \\
RoBERTa-Seq2Seq-Ptr\_cc	& 83.51 & 83.78 & 98.17 \\ \hline
\end{tabular}
}
\caption{\label{table-BaselineTOP} Results on Facebook TOP. The SR models are by \citet{einolghozati2019improving}. Refer to the caption of Table \ref{table-SNIPS} for abbreviation explanations.}
\end{table}

Table \ref{table-BaselineTOP} presents results on Facebook TOP. We follow \citet{rongali} and experiment with Seq2Seq models. We use different pretrained BERT-variants as the encoder, transformer decoder layers with $d_{model}=768$ \cite{attention}, and a pointer generator network \cite{vinyals2015pointer} which uses scaled dot-product attention to score tokens. The model is trained using the cross-entropy loss with label smoothing of 0.1. For simplicity, we always train code embeddings for 900 epochs offline. Learning rate 2e-5 and weight decay 0.01 are used for transformer training. BERT and DistilBERT are cased in these experiments. During inference, we employ beam decoding with width 5. Our greatly compressed models present 98$\sim$99\% performances of the original models.

\begin{table*}
\centering
\tiny
\resizebox{.75\textwidth}{!}{%
\begin{tabular}{ccccccc}
\hline \textbf{Epoch} & \textbf{MeanEucDist} &\textbf{NN-cos} & \textbf{NN-Euc} & \textbf{SNIPS} & \textbf{ATIS} & \textbf{TOP}\\ \hline
100	    & 0.3677$\pm$0.25\%  & 0.66$\pm$1.90\%        & 0.65$\pm$2.00\% & 79.29 & 82.31 & 78.09 \\
200	    & 0.3254$\pm$0.08\%  & 2.20$\pm$0.69\%         & 2.30$\pm$0.84\% & 85.43 & 84.99 & 81.59 \\
300	    & 0.3023$\pm$0.09\%  & 3.66$\pm$0.92\%         & 3.96$\pm$0.55\%         & 86.86 & 86.11 & 83.17 \\
400	    & 0.2841$\pm$0.23\%  & 4.84$\pm$0.58\%         & 5.26$\pm$0.83\%         & \textbf{89.71} & 87.01 & 83.45 \\
500	    & 0.2685$\pm$0.26\%  & 5.72$\pm$0.48\%         & 6.21$\pm$0.78\%         & 87.71 & 87.23 & 83.82 \\
600	    & 0.2573$\pm$0.12\%  & 6.20$\pm$0.39\%         & 6.72$\pm$0.18\%         & 88.14 & 85.69 & 83.41 \\
700	    & 0.2499$\pm$0.20\%  & 6.42$\pm$0.49\%         & 6.94$\pm$0.33\%         & 88.00 & \textbf{87.35} & 84.27 \\
800	    & 0.2444$\pm$0.07\%  & 6.54$\pm$0.39\%         & 7.07$\pm$0.15\%         & 88.57 & 86.90 & 84.09 \\
900	    & 0.2407$\pm$0.10\%  & 6.62$\pm$0.31\%         & 7.14$\pm$0.14\%         & 88.57 & 86.56 & \textbf{84.42} \\
1000	& 0.2380$\pm$0.07\%  & 6.65$\pm$0.39\%         & 7.16$\pm$0.10\%         & 89.14 & 87.12 & 83.86 \\
\hline
\end{tabular}
}
\caption{\label{table-emp} Analyses for the code embedding learning process (M=32, K=16). MeanEucDist, NN-cos, and NN-Euc are averaged across 5 runs. ``SNIPS", ``ATIS", and ``TOP" are the test exact match achieved on the three datasets.}
\end{table*}

\subsection{Analysis for Code Convergence}
We study the relationship among a few variables during code learning for the embeddings from pretrained ALBERT-base (Table \ref{table-emp}). During the first 1000 epochs, the mean Euclidean distance between the original and reconstructed embeddings decrease with a decreasing rate. The average number of shared top-20 nearest neighbours according to cosine similarity and Euclidean distances between the two embeddings increase with a decreasing rate. We apply code embeddings trained for different numbers of epochs to ALBERT-base-v2 and finetune on semantic parsing. On SNIPS and ATIS, we find the best validation setting among learning rate \{2,3,4,5,6\}e-5 and weight decay \{0.01, 0.05, 0.01\}. We observe that the test exact match plateaus for code embeddings trained for more than 400 epochs. On Facebook TOP, we use learning rate 2e-5 and weight decay 0.01, and observe the similar trend.

\subsection{Effects of M and K}
We use embeddings from pretrained ALBERT-base-v2 as reference to learn code embeddings with M in \{8, 16, 32, 64\} and K in \{16, 32, 64\}. As shown in Table \ref{table-mk}, after 700 epochs, the MSE loss for embeddings with \textit{larger} M and K converges to \textit{smaller} values in general. With M=64, more epochs are needed for convergence to smaller MSE losses compared to those from smaller M. We apply the embeddings to ALBERT-base-v2 and finetune on SNIPS. In general, larger M yields better performances. Effects of K are less clear when M is large.

\begin{table}[t]
\centering
\tiny
\resizebox{.4\textwidth}{!}{%
\begin{tabular}{cccccc}
\hline \textbf{M} & \textbf{K} & \textbf{epo} & \textbf{MSE} & \textbf{EM}\\ \hline
8 & 16 & 700 & 0.3155$\pm$0.05\% & 85.43 \\
8 & 32 & 700 & 0.3032$\pm$0.04\% & 87.43 \\
8 & 64 & 700 & 0.2944$\pm$0.04\% & 87.43\\
16 & 16 & 700 & 0.2855$\pm$0.05\% & 88.57\\
16 & 32 & 700 & 0.2727$\pm$0.09\% & 88.00\\
16 & 32 & 700 & 0.2669$\pm$0.08\% & 88.14\\
32 & 16 & 700 & 0.2499$\pm$0.20\% & 89.00\\
32 & 32 & 700 & 0.2421$\pm$0.20\% & 89.14\\
32 & 64 & 700 & 0.2396$\pm$0.27\% & 88.29\\
64 & 16 & 700 & 0.2543$\pm$0.47\% & 88.29\\
64 & 16 & 1000 & 0.2256$\pm$1.06\% & 89.71\\
64 & 32 & 700 & 0.2557$\pm$0.37\% & 89.86\\
64 & 32 & 1000 & 0.2159$\pm$0.43\% & 89.71\\
\hline
\end{tabular}
}
\caption{\label{table-mk} Effects of M and K. Mean squared errors (MSE) are averaged over 5 runs. Best validation exact match (EM) is presented for compressed transformer models trained with 0.05 weight decay and \{3,4,5,6,7\}e-5 peak learning rates on SNIPS.}
\end{table}

\section{Conclusion}
Current state-of-the-art task-oriented semantic parsing models are based on pretrained RoBERTa-base (478MB) or BERT-base (420MB). We apply DistilBERT (256MB), ALBERT-large (68MB), and ALBERT-base (45MB), and observe near state-of-the-art performances. We learn compositional code embeddings to compress the model embeddings by 95.15\% $\sim$ 98.46\%, the pretrained encoders by 20.47\% $\sim$ 34.22\%, and observe 97.5\% performance preservation on SNIPS, ATIS, and Facebook TOP. Our compressed ALBERT-large is 54MB and can achieve 99.6\% performances of the previous state-of-the-art models on SNIPS and ATIS. Our technique has potential to be applied to more tasks including machine translation in the future.

\section*{Acknowledgement}
This project is part of the data science industry mentorship program initiated by Andrew McCallum at University of Massachusetts Amherst. We thank the teaching assistants Rajarshi Das and Xiang Lorraine Li for helpful discussion and the instructor Andrew McCallum for valuable feedback. Experiments in this project are conducted on the Gypsum cluster at UMassAmherst. The cluster is purchased with funds from the Massachusetts Technology Collaborative.

\bibliographystyle{acl_natbib}


\section*{Acknowledgements}

The research of P.A. was supported by Sistema Nacional de
Investigadores (M\'exico).

\section*{Addendum}

According to the reviewer: ``Another point to be observed is the
dependence of the Rayleigh-Ritz variational method on the choice
of the wave function. Despite not being mentioned by the authors
of this comment, the wave function used in the Rayleigh-Ritz
variational method is obtained from the asymptotic analysis made
by Vit\'{o}ria et al. If one uses another wave function that
differs from the wave function obtained from the asymptotic
analysis made by Vit\'{o}ria et al, therefore, the results will be
different. In addition, no mathematical proof has been shown in
this comment that clarifies the relation of the approximate
solutions to the biconfluent Heun equation.''

This comment is surprising. In order to apply the Ritz variational
method it is mandatory that the basis functions satisfy the
correct boundary conditions at $x=0$ and $x\rightarrow \infty $.
Therefore, we have chosen the simplest basis set that satisfy such
boundary conditions. The  set of Gaussian functions chosen here is
complete and, for this reason it should give the actual
eigenvalues of the problem at hand. This fact is clearly revealed
in the convergence of the approximate eigenvalues shown in Tables
\ref{tab:theta1} and \ref{tab:theta2}. The Ritz variational method
is well known and and has been widely used for the study of many
quantum-mechanical problems.


\section{Introduction}
\label{sec:introduction}

Structured prediction often relies on models that train on maximum likelihood and use beam search for approximate decoding.
This procedure leads to two significant mismatches between the training and testing settings:
the model is trained on oracle trajectories and therefore does not learn about its own mistakes;
the model is trained without beam search and therefore does not learn how to use the beam effectively for search.

Previous algorithms have addressed one or the other of these mismatches.
For example, DAgger~\cite{ross2011reduction} and scheduled sampling~\cite{bengio2015scheduled} use the learned model to visit non-oracle states at training time, but do not use beam search (i.e., they keep a single hypothesis).
Early update~\cite{collins_incremental_2004}, LaSO~\cite{daume2005learning}, and BSO~\cite{wiseman2016sequence} are trained with beam search, but do not expose the model to beams without a gold hypothesis (i.e., they either stop or reset to beams with a gold hypothesis).

Recently, \citet{negrinho2018learning} proposed a meta-algorithm that instantiates beam-aware algorithms as a result of choices for the surrogate loss (i.e., which training loss to incur at each visited beam) and data collection strategy (i.e., which beams to visit during training).
A specific instantiation of their meta-algorithm addresses both mismatches by relying on an insight on how to induce training losses for beams without the gold hypothesis: \emph{for any beam}, its lowest cost neighbor should be scored sufficiently high to be kept in the successor beam.
To induce these training losses it is sufficient to be able to compute the best neighbor of any state (often called a dynamic oracle~\cite{goldberg2012dynamic}).
Unfortunately, \citet{negrinho2018learning} do not provide empirical results, leaving open questions such as whether instances can be trained robustly, when is beam-aware training most useful, and what is the impact on performance of the design choices.

\paragraph{Contributions}
We empirically study beam-aware algorithms instantiated through the meta-algorithm of~\citet{negrinho2018learning}.
We tackle supertagging as it is a sequence labelling task with an easy-to-compute dynamic oracle and a moderately-sized label set (approximately $1000$) which may require more effective search.
We examine two supertagging models (one from \citet{vaswani2016supertagging} and a simplified version designed to be heavily reliant on search) and train them with instantiations of the meta-algorithm.
We explore how design choices influence performance, and give recommendations based on our empirical findings.
For example, we find that perceptron losses perform consistently worse than margin and log losses.
We observe that beam-aware training can have a large impact on performance, particularly when the model must use the beam to manage uncertainty during prediction.
Code for reproducing all results in this paper is available at \url{https://github.com/negrinho/beam_learn_supertagging}.

\begin{figure*}[tbh]
  \centering
  \includegraphics[width=0.75\textwidth]{figures_paper/step.pdf}
  \caption{Beam $b$ has neighborhood $A_b$, where $k = |b| = |b'| = 3$ and $n = |A_b| = 5$.
  Edges from elements in $b$ to elements in $A_b$ encode neighborhood relationships, e.g., $v_3$ has a single neighbor $v'_5$.
  Permutation $\hat \sigma : [n] \to [n]$ sorts hypotheses in decreasing order of score, and permutation $\sigma^* : [n] \to [n]$ sorts them in increasing order of cost, i.e, $v'_{\sigma^*(1)}$ is the lowest cost neighbor and $v'_{\hat \sigma(1)}$ is the highest scoring neighbor.
  The successor beam $b'$ keeps the neighbor states in $A_b$ with highest score according to vector $s$, or equivalently highest rank according to $\hat \sigma$. %
  %
  }
  \label{fig:local_notation}
\end{figure*}

\begin{figure*}[tbh]
    \centering
    \includegraphics[width=0.68\textwidth]{figures_paper/trajectory.pdf}
    \caption{Sampling a trajectory through the beam search space at training time.
    A loss $\ell(b_i, \theta)$ is incurred at each visited beam $b_i$, $i \in [h - 1]$, resulting in total accumulated loss $\ell(b_{1:h}, \theta)$ for beam trajectory $b_{1:h}$.
    The terminal beam $b_h$ corresponds to a complete output $y(b_h) \in \mathcal Y$.
    Transitions between beams are sampled according to a data collection policy $\pi' : V_k \to \Delta(V_k)$.
    We consider $\pi'$ induced by a scoring function $s(\cdot, \theta) : V \to \mathbb R$ or cost function $c^* : V \to \mathbb R$.
    Parameters $\theta$ parametrize the scoring function of the model.
    Losses $\ell(b_i, \theta)$ are low if the scores of the neighbors of $b_i$ comfortably keep the lowest cost elements in the successor beam (see Section~\ref{ssec:surrogate_losses}), and high otherwise.
    See Figure~\ref{fig:local_notation} for the notation to describe the surrogate loss $\ell(b_i, \theta)$ at each beam $b_i$.
    }
    \label{fig:trajectory}
\end{figure*}

\section{Background on learning to search and beam-aware methods}
\label{sec:background}

For convenience, we reuse notation introduced in~\citet{negrinho2018learning} to describe their meta-algorithm and its components (e.g., scoring function, surrogate loss, and data collection strategy).
See Figure~\ref{fig:local_notation} and Figure~\ref{fig:trajectory} for an overview of the notation.
When relevant, we instantiate notation for left-to-right sequence labelling under the Hamming cost, which supertagging is a special case of.

\paragraph{Input and output spaces}
Given an input structure $x \in \mathcal X$, the output structure $y \in \mathcal Y_x$, is generated through a sequence of incremental decisions.
An example $x \in \mathcal X$ induces a tree $G_x = (V_x, E_x)$ encoding the sequential generation of elements in $\mathcal Y_x$, where $V_x$ is the set of nodes and $E_x$ is the set of edges.
The leaves of $G_x$ correspond to elements of $\mathcal Y_x$ and the internal nodes correspond to incomplete outputs.
For left-to-right sequence labelling, for a sequence $x \in \mathcal X$, each decision assigns a label to the current position of $x$ and the nodes of tree encode labelled prefixes of $x$, with terminal nodes encoding complete labellings of $x$.

\paragraph{Cost functions}
Given a golden pair $(x, y) \in \mathcal X \times \mathcal Y$, the cost function $c _{x, y} : \mathcal Y_x \to \mathbb R$ measures how bad the prediction $\hat y \in \mathcal Y_x$ is relative to the target output structure $y \in \mathcal Y_x$.
Using $c _{x, y} : \mathcal Y_x \to \mathbb R$, we define a cost function $c^*_{x,y} : V _x \to \mathbb R$ for partial outputs by assigning to each node $v \in V_x$ the cost of its best reachable complete output, i.e., $c _{x, y} ^*(v) = \min _{v' \in T_v} c _{x, y}(v')$, where $T_v \subseteq \mathcal Y_x$ is the set of complete outputs reachable from $v$.
For a left-to-right search space for sequence labelling, if $c_{x, y} : \mathcal Y_x \to \mathbb R$ is Hamming cost, the optimal completion cost $c^*_{x, y} : \mathcal Y_x \to \mathbb R$ is the number of mistakes in the prefix as the optimal completion matches the remaining suffix of the target output.

\paragraph{Dynamic oracles}
An oracle state is one for which the target output structure can be reached.
Often optimal actions can only be computed for oracle states.
Dynamic oracles compute optimal actions even for non-oracle states.
Evaluations of $c^* _{x, y} : V _x \to \mathbb R$ for arbitrary states allows us to induce the dynamic oracle---at a state $v \in V_x$, the optimal action is to transition to the neighbor $v' \in N_v$ with the lowest completion cost.
For sequence labelling, this picks the transition that assigns the correct label.
For other tasks and metrics, more complex dynamic oracles may exist, e.g., in dependency parsing~\citep{goldberg2012dynamic,goldberg2013training}.
For notational brevity, from now on, we omit the dependency of the search spaces and cost function on $x \in \mathcal X$, $y \in \mathcal Y$, or both.

\paragraph{Beam search space}
Given a search space $G = (V, E)$, the beam search space $G _k  = (V_k, E_k)$ is induced by choosing a beam size $k \in \mathbb N$ and a strategy for generating the successor beam out of the current beam and its neighbors.
In this paper, we expand all the elements in the beam and score the neighbors simultaneously.
The highest scoring $k$ neighbors are used to form the successor beam.
For $k = 1$, we recover the greedy search space $G$.

\paragraph{Beam cost functions}
The natural cost function $c^* : V_k \to \mathbb R$ for $G _k$ is created from the element-wise cost function on $G$, and assigns to each beam the cost of its best element, i.e., for $b \in V_k$, $c^*(b) = \min_{v \in b} c^*(v)$.
For a transition $(b, b') \in E_k$, we define the transition cost $c(b, b') = c^*(b') - c^*(b)$, where $b' \in N_b$, i.e., $b'$ can be formed from the neighbors of the elements in $b$.
A cost increase happens when $c(b, b') > 0$, i.e., the best complete output reachable in $b$ is no longer reachable in $b'$.

\paragraph{Policies}
Policies operate in beam search space $G _k$ and are induced through a learned scoring function $s(\cdot, \theta) : V \to \mathbb R$ which scores elements in the original space $G$.
A policy $\pi: V_k \to \Delta(V_k)$, i.e., mapping states (i.e., beams) to distributions over next states.
We only use deterministic policies where the successor beam is computed by sorting the neighbors in decreasing order of score and taking the top $k$.

\paragraph{Scoring function}
In the non-beam-aware case, the scoring function arises from the way probabilities of complete sequences are computed with the locally normalized model, namely
  $
  p(y|x, \theta) = \prod _{j = 1}^h p(y_i | y_{1:i-1}, x, \theta)
  $,
where we assume that all outputs for $x \in \mathcal X$ have $h$ steps.
For sequence labelling, $h$ is the length of the sentence.
The resulting scoring function $s(\cdot, \theta) : V \to \mathbb R$ is
  $
  s(v, \theta) = \sum _{i = 1}^j \log p(y_i | y_{1:i-1}, x, \theta)
  $,
where $v = y_{1:j}$ and $j \in [h]$.
Similarly, the scoring function that we learn in the beam-aware case is
  $
  s(v, \theta) = \sum _{i = 1}^j \tilde s (v, \theta)
  $,
where $x$ has been omitted, $v = y_{1:j}$, and $\tilde s(\cdot, \theta) : V \to \mathbb R$ is the learned incremental scoring function.
In Section~\ref{sec:discussion}, we observe that this cumulative version performs uniformly better than the non-cumulative one.

\section{Meta-algorithm for learning beam search policies}
\label{sec:meta_algorithm}

We refer the reader to \citet{negrinho2018learning} for a discussion of how specific choices for the meta-algorithm recover algorithms from the literature.

\subsection{Data collection strategies}
\label{ssec:data_collection}

The data collection strategy determines which beams are visited at training time (see Figure~\ref{fig:trajectory}).
Strategies that use the learned model differ on how they compute the successor beam $b' \in N_b$ when $s(\cdot, \theta)$ leads to a beam without the gold hypothesis, i.e., $c(b, b') > 0$, where $b' = \{v_{\hat \sigma(1)}, \ldots, v_{\hat \sigma (k)} \} \subset A_b$ and $A_b = \{v_1, \ldots, v_n\} = \cup_{v \in b} N_v$.
We explore several data collection strategies:

\paragraph{stop}
If the successor beam does not contain the gold hypothesis, stop collecting the trajectory.
Structured perceptron training with early update~\cite{collins_incremental_2004} use this strategy.

\paragraph{reset}
If the successor beam does not contain the gold hypothesis, reset to a beam with only the gold hypothesis\footnote{Any beam with the gold hypothesis would be valid, e.g., the top $k - 1$ according to the scores plus the gold hypothesis, which we call \textit{reset (multiple)}}.
LaSO~\cite{daume2005learning} use this strategy.
For $k = 1$, we recover teacher forcing as only the oracle hypothesis is kept in the beam.

\paragraph{continue}
Ignore cost increases, always using the successor beam.
DAgger~\cite{ross2011reduction} take this strategy, but does not use beam search.
\citet{negrinho2018learning} suggest this for beam-aware training but do not provide empirical results.

\paragraph{reset (multiple)}
Similar to reset, but keep $k - 1$ hypothesis from the transition, i.e., $b' = \{v_{\sigma^*(1)}, v_{\hat \sigma(1)}, \ldots v_{\hat \sigma(k - 1)}\}$.
We might expect this data collection strategy to be closer to \textit{continue} as a large fraction of the elements of the successor beam are induced by the learned model.

\paragraph{oracle}
Always transition to the beam induced by $\sigma^*: [n] \to [n]$, i.e., the one obtained by sorting the costs in increasing order.
For $k = 1$, this recovers teacher forcing.
In Section~\ref{ssec:comp_data_collection}, we observe that \textit{oracle} dramatically degrades performance due to increased exposure bias with increased $k$.

\subsection{Surrogate losses}
\label{ssec:surrogate_losses}

Surrogate losses encode that the scores produced by the model for the neighbors must score the best neighbor sufficiently high for it to be kept comfortably in the successor beam.
For $k = 1$, many of these losses reduce to losses used in non-beam-aware training.
Given scores $s \in \mathbb R^n$ and costs $c \in \mathbb R^n$ over neighbors in $A_b = \{v_1, \ldots, v_n\}$, we define permutations $\hat \sigma: [n] \to [n]$ and $\sigma^*: [n] \to [n]$ that sort the elements in $A_b$ in decreasing order of scores and increasing order of costs, respectively, i.e., $s_{\hat \sigma(1)} \geq \ldots \geq s_{\hat \sigma(n)}$ and $c_{\sigma^*(1)} \leq \ldots \leq s_{\sigma^*(n)}$.
See Figure~\ref{fig:local_notation} for a description of the notation used to describe surrogate losses.
Our experiments compare the following surrogate losses:

\paragraph{perceptron (first)} Penalize failing to score the best neighbor at the top of the beam (regardless of it falling out of the beam or not).
\begin{align*}
\ell(s, c) =
        \max
          \left( 0, s _{ \hat \sigma (1) } - s _{ \sigma ^*(1) } \right).
  %
\end{align*}

\paragraph{perceptron (last)} If this loss is positive at a beam, the successor beam induced by the scores does not contain the golden hypothesis.
\begin{align*}
\ell(s, c) = \max
    \left(
        0, s _{ \hat \sigma (k) } - s _{ \sigma ^*(1) }
    \right).
\end{align*}

\paragraph{margin (last)} Penalize margin violations of the best neighbor of the hypothesis in the current beam.
Compares the correct neighbor $s_{\sigma^*(1)}$ with the neighbor $v _{ \hat \sigma(k)}$ last in the beam.
\begin{align*}
\ell(s, c) = \max
    \left(
        0, s _{ \hat \sigma(k) } - s _{ \sigma ^*(1) } + 1
    \right)
\end{align*}

\paragraph{cost-sensitive margin (last)}
Same as \textit{margin (last)} but weighted by the cost difference of the pair.
\citet{wiseman2016sequence} use this loss.
\begin{align*}
\ell(s, c) =
    (c _{\hat \sigma (k) } - c _{ \sigma ^* (1) })
      \max
        ( 0, s _{ \hat \sigma (k) } - s _{ \sigma ^*(1) } +1 ).
\end{align*}

\paragraph{log loss (neighbors)}
    Normalizes over all elements in $A_b$. For beam size $k = 1$, it reduces to the usual log loss.
    \begin{align*}
          \ell(s, c)
          &= - s _{ \sigma ^*(1) } + \log
          \left(
              { \sum _{ i = 1} ^n \exp ( s _i ) }
          \right)
      %
  \end{align*}

\paragraph{log loss (beam)}
Normalizes only over the top $k$ neighbors of a beam according to the scores $s$.
  \begin{align*}
        \ell(s, c)
          &= - s _{ \sigma ^*(1) } + \log
          \left(
              \sum _{ i \in I }  \exp ( s _i )
          \right) ,
      %
  \end{align*}
 where $I =\{ \sigma ^*(1), \hat \sigma(1), \ldots, \hat \sigma(k) \}$.
 The normalization is only over the golden hypothesis $v _{\sigma ^*(1)}$ and the elements included in the beam.
 \citet{andor2016globally} use this loss.

\begin{figure*}[tbp]
  \centering
  \includegraphics[width=0.8\textwidth]{figures_paper/models_two.pdf}
  \caption{
    High-level structure of the two models used in the experiments.
    The model on the left is from \citet{vaswani2016supertagging}.
    The model on the right is a simplification of the one on the left, namely, it does not have an encoding of the complete sentence at the start of prediction.
  }
  \label{fig:model}
\end{figure*}

\subsection{Training}
\label{ssec:training}

The meta-algorithm of \citet{negrinho2018learning} is instantiated by choosing a surrogate loss, data collection strategy, and beam size.
Training proceeds by sampling an example $(x, y) \in \mathcal X \times \mathcal Y$ from the training set.
A trajectory through the beam search space $G_k$ is collected using the chosen data collection strategy.
A surrogate loss is induced at each non-terminal beam in the trajectory (see Figure~\ref{fig:trajectory}).
Parameter updates are computed based on the gradient of the sum of the losses of the visited beams.

\section{Experiments}
\label{sec:experiments}

We explore different configurations of the design choices of the meta-algorithm to understand their impact on training behavior and performance.

\subsection{Task details}
\label{ssec:task_details}

We train our models for supertagging, a sequence labelling where accuracy is the performance metric of interest.
Supertagging is a good task for exploring beam-aware training, as contrary to other sequence labelling datasets such as named-entity recognition~\citep{tjong_kim_sang_introduction_2003}, chunking~\citep{sang2000introduction}, and part-of-speech tagging~\citep{marcus_building_1993}, has a moderate number of labels and therefore it is likely to require effective search to achieve high performances.
We used the standard splits for CCGBank~\citep{hockenmaier2007ccgbank}: the training and development sets have, respectively, $39604$ and $1913$ examples.
Models were trained on the training set and used the development set to compute validation accuracy at the end of each epoch to keep the best model.
As we are performing an empirical study, similarly to \citet{vaswani2016supertagging}, we report validation accuracies.
Each configuration is ran three times with different random seeds and the mean and standard deviation are reported.
We replace the words that appear at most once in the training set by \texttt{UNK}.
By contrast, no tokenization was done for the training supertags.

\subsection{Model details}
\label{ssec:model_details}

We have implemented the model of~\citet{vaswani2016supertagging} and a simpler model designed by removing some of its components.
The two main differences between our implementation and theirs are that we do not use pretrained embeddings (we train the embeddings from scratch) and we use the gold POS tags (they use only the pretrained embeddings).

\paragraph{Main model}
For the model of \citet{vaswani2016supertagging} (see Figure~\ref{fig:model}, left), we use $64$, $16$, and $64$ for the dimensions of the word, part-of-speech, and supertag embeddings, respectively.
All LSTMs (forward, backward, and LM) have hidden dimension $256$.
We refer the reader to \citet{vaswani2016supertagging} for the exact description of the model.
Briefly, embeddings for the words and part-of-speech tags are concatenated and fed to a bi-directional LSTM, the outputs of both directions are then fed into a combiner (dimension-preserving linear transformations applied independently to both inputs, added together, and passed through a ReLU non-linearity).
The output of the combiner and the output of the LM LSTM (which tracks the supertag prefix up to a prediction point) is then passed to another combiner that generates scores over supertags.

\paragraph{Simplified model}
We also consider a simplified model that drops the bi-LSTM encoder and the corresponding combiner (see Figure~\ref{fig:model}, right).
The concatenated embeddings are fed directly into the second combiner with the LM LSTM output.
Values for the hyperparameters are the same when possible.
This model must leverage the beam effectively as it does not encode the sentence with a bi-LSTM.
Instead, only the embeddings for the current position are available, giving a larger role to the LM LSTM over supertags.
While supertagging can be tackled with a stronger model, this restriction is relevant for real-time tasks, e.g., the complete input might not be known upfront.

\paragraph{Training details}
Models are trained for $16$ epochs with SGD with batch size $1$ and cosine learning rate schedule~\citep{loshchilov2016sgdr}, starting at $10^{-1}$ and ending at $10^{-5}$.
No weight decay or dropout was used.
Training examples are shuffled after each epoch.
Results are reported for the model with the best validation performance across all epochs.
We use $16$ epochs for all models for simplicity and fairness.
This number was sufficient, e.g.,
we replicated Table~\ref{tab:results_strategies} by training with $32$ epochs and observed minor performance differences (see Table~\ref{tab:results_strategies_train32epochs}).

\subsection{Non-beam-aware training}
\label{ssec:nonbeam_aware_training}

We first train the models with $k = 1$ and then use beam search to decode.
Crucially, the model does not train with a beam and therefore does not learn to use it effectively.
We vary the data collection strategy.
The results are presented in Table~\ref{tab:results_nonbeam_aware} and should be used as a reference when reading the other tables to evaluate the impact of beam-aware training.
Tables are formatted such that the first and second horizontal halves contain the results for the main model and simplified model, respectively.
Each position contains the mean and the standard deviation of running that configuration three times.
We use this format in all tables presented.

\begin{table}[tb]
  %
	\centering
  %
  \resizebox{\columnwidth}{!}{
	\begin{tabular}{lcccc}
		\toprule
	    & 1 & 2 & 4 & 8 \\
		\midrule
      oracle/reset &  $93.78_{0.12}$ & $93.81_{0.11}$ & $93.82_{0.10}$ & $93.82_{0.10}$ \\
      continue &      $94.04_{0.07}$ & $94.05_{0.07}$ & $94.05_{0.07}$ & $94.06_{0.07}$ \\
      stop &          $93.86_{0.09}$ & $93.90_{0.07}$ & $93.90_{0.07}$ & $93.91_{0.07}$ \\
    \midrule
      oracle/reset &  $73.20_{0.31}$ & $76.55_{0.24}$ & $77.42_{0.27}$ & $77.54_{0.22}$ \\
		  continue &      $81.99_{0.04}$ & $82.30_{0.03}$ & $82.37_{0.08}$ & $82.41_{0.08}$ \\
      stop &          $74.35_{0.23}$ & $77.06_{0.14}$ & $77.73_{0.13}$ & $77.82_{0.09}$ \\
    \bottomrule
  \end{tabular}
  }
	\caption{
		Development accuracies for models trained with different data collection strategies in a non-beam-aware way (i.e., $k = 1$) and decoded with beam search with varying beam size.
    \textit{continue} performs best, showing the importance of exposing the model to its mistakes.
    Differences are larger for the simplified model.
    }
	\label{tab:results_nonbeam_aware}
\end{table}

The \textit{continue} data collection strategy (i.e., DAgger for $k = 1$) results in better models than training on the oracle trajectories.
Beam search results in small gains for these settings.
In this experiment, training with oracle is the same as training with reset as the beam always contains only the oracle hypothesis.
The performance differences are small for the main model but much larger for the simplified model, underscoring the importance of beam search when there is greater uncertainty about predictions.
For the stronger model, the encoding of the left and right contexts with the bi-LSTM provides enough information at each position to predict greedily, i.e., without search.
This difference appears consistently in all experiments, with larger gains for the weaker model.

The gains achieved by the main model by decoding with beam search post-training are very small (from $0.02$ to $0.05$).
This suggests that training the model in a non-beam-aware fashion and then using beam search does not guarantee improvements.
The model must be learned with search to improve on these results.
For the simpler model, larger improvements are observed (from $0.42$ to $4.34$).
Despite the gains with beam search for \textit{reset} and \textit{stop}, they are not sufficient to beat the greedy model trained on its own trajectories, yielding $81.99$ for \textit{continue} with $k = 1$ versus $77.54$ for \textit{oracle} and $77.82$ for \textit{reset}, both with $k = 8$.
These results show the importance of the data collection strategy, even when the model is not trained in a beam-aware fashion.
These gains are eclipsed by beam-aware training, namely, compare Table~\ref{tab:results_nonbeam_aware} with Table~\ref{tab:results_strategies}.
See Figure~\ref{fig:non_beam_vaswani} for the evolution of the validation and training accuracies with epochs.

\begin{table}[tb]
	\centering
  %
  \resizebox{\columnwidth}{!}{
	\begin{tabular}{lcccc}
		\toprule
	    & 1 & 2 & 4 & 8 \\
		\midrule
      oracle              & $94.10_{0.08}$ & $92.98_{0.07}$ & $91.66_{0.22}$ & $85.95_{0.79}$ \\
      reset               & $94.20_{0.11}$ & $94.34_{0.06}$ & $94.33_{0.01}$ & $94.42_{0.04}$ \\
      reset (mult.)       & $94.15_{0.07}$ & $93.98_{0.08}$ & $94.06_{0.06}$ & $94.16_{0.05}$ \\
      continue            & $94.15_{0.02}$ & $94.35_{0.05}$ & $94.37_{0.04}$ & $94.33_{0.04}$ \\
      stop                & $93.95_{0.09}$ & $94.11_{0.05}$ & $94.24_{0.07}$ & $94.25_{0.06}$ \\
    \midrule
      oracle              & $75.09_{0.17}$ & $80.67_{0.40}$ & $78.69_{1.27}$ & $47.38_{1.79}$ \\
      reset               & $75.06_{0.16}$ & $87.21_{0.14}$ & $91.24_{0.02}$ & $92.46_{0.09}$ \\
      reset (mult.)       & $75.04_{0.18}$ & $86.19_{0.12}$ & $90.76_{0.11}$ & $92.16_{0.03}$ \\
		  continue            & $82.01_{0.06}$ & $89.17_{0.08}$ & $91.80_{0.12}$ & $92.69_{0.01}$ \\
      stop                & $75.08_{0.54}$ & $87.16_{0.08}$ & $90.98_{0.13}$ & $92.18_{0.06}$ \\
    \bottomrule
  \end{tabular}
  }
	\caption{
		Development accuracies for beam-aware training with varying data collection strategies.}
	\label{tab:results_strategies}
\end{table}

\subsection{Comparing data collection strategies}
\label{ssec:comp_data_collection}

We train both models using the \textit{log loss (neighbors)}, described in Section~\ref{ssec:surrogate_losses}, and vary the data collection strategy, described in Section~\ref{ssec:data_collection}, and beam size.
Results are presented in Table~\ref{tab:results_strategies}
Contrary to Section~\ref{ssec:nonbeam_aware_training}, these models are trained to use beam search rather than it being an artifact of approximate decoding.
Beam-aware training under \emph{oracle} worsens performance with increasing beam size (due to increasing exposure bias).
During training, the model learns to pick the best neighbors for beams containing only close to optimal hypotheses, which are likely very different from the beams encountered when decoding.
The results for the simplified model are similar---with increasing beam size, performance first improves but then degrades.
For the main model, we observe modest but consistent improvements with larger beam sizes across all data collection strategies except \textit{oracle}.
By comparing the results with those in the first row of Table~\ref{tab:results_nonbeam_aware}, we see that we improve on the model trained with maximum likelihood and decoded with beam search.

\begin{table}[tb]
	\centering
  %
  \resizebox{\columnwidth}{!}{
	\begin{tabular}{lcccc}
		\toprule
	    & 1 & 2 & 4 & 8 \\
		\midrule
        percep. (first)         & $92.81_{0.06}$ & $93.22_{0.04}$ & $93.44_{0.02}$ & $93.52_{0.06}$ \\
        percep. (last)          & $92.84_{0.11}$ & $93.57_{0.06}$ & $93.86_{0.09}$ & $93.77_{0.04}$ \\
        m. (last)               & $94.10_{0.07}$ & $94.29_{0.07}$ & $94.27_{0.03}$ & $94.43_{0.04}$ \\
        cost-s. m. (last)       & $93.98_{0.03}$ & $94.32_{0.10}$ & $94.37_{0.03}$ & $94.33_{0.13}$ \\
        log loss (beam)         & $92.29_{0.07}$ & $92.09_{0.11}$ & $94.24_{0.08}$ & $94.32_{0.02}$ \\
        log loss (neig.)        & $94.22_{0.00}$ & $94.29_{0.03}$ & $94.27_{0.06}$ & $94.38_{0.01}$ \\
        \midrule
        percep. (first)         & $77.62_{0.14}$ & $86.32_{0.05}$ & $89.83_{0.05}$ & $91.00_{0.07}$ \\
        percep. (last)          & $77.67_{0.07}$ & $87.62_{0.03}$ & $90.82_{0.16}$ & $91.98_{0.11}$ \\
        m. (last)               & $81.75_{0.04}$ & $88.80_{0.02}$ & $91.91_{0.05}$ & $92.81_{0.05}$ \\
        cost-s. m. (last)       & $81.76_{0.05}$ & $88.92_{0.06}$ & $91.81_{0.03}$ & $92.81_{0.03}$ \\
        log loss (beam)         & $77.50_{0.07}$ & $88.25_{0.08}$ & $91.46_{0.06}$ & $92.56_{0.11}$ \\
        log loss (neig.)        & $81.94_{0.02}$ & $89.01_{0.10}$ & $91.75_{0.03}$ & $92.60_{0.03}$ \\
	    \bottomrule
  \end{tabular}
  }
	\caption{Development accuracies for the loss functions in Section~\ref{ssec:surrogate_losses}.}
	\label{tab:results_losses}
\end{table}

\begin{figure*}[tbp]
  \centering
  \includegraphics[width=0.24\textwidth]{figures_paper/table1/dev_acc_vaswani.pdf}
  \includegraphics[width=0.24\textwidth]{figures_paper/table1/train_acc_vaswani.pdf}
  \includegraphics[width=0.24\textwidth]{figures_paper/table1/dev_acc_lm.pdf}
  \includegraphics[width=0.24\textwidth]{figures_paper/table1/train_acc_lm.pdf}
  \caption{Validation and training accuracies for non-beam-aware training (i.e., $k = 1$) with different data collection strategies for the main (left half) and simplified (right half) models.
  \textit{continue} achieves higher accuracies.
  %
  %
  }
  \label{fig:non_beam_vaswani}
\end{figure*}

\begin{figure*}[tbp]
  \centering
  \includegraphics[width=0.24\textwidth]{figures_paper/table2/dev_acc_vaswani.pdf}
  \includegraphics[width=0.24\textwidth]{figures_paper/table2/train_acc_vaswani.pdf}
  \includegraphics[width=0.24\textwidth]{figures_paper/table2/dev_acc_lm.pdf}
  \includegraphics[width=0.24\textwidth]{figures_paper/table2/train_acc_lm.pdf}
  \caption{Validation and training accuracies for beam-aware training with different data collection strategies and beam sizes for the main (left half) and simplified (right half) models.
  Larger beam sizes achieve higher performances while overfitting less, and are crucial for the simplified model to achieve higher training and validation accuracies.
  For smaller beams \textit{continue} performs better than \textit{reset}.
  All models can be trained stably from scratch.
  Three runs were aggregated by showing the mean and the standard deviation for each epoch.
  %
  }
  \label{fig:beam_vaswani}
\end{figure*}

The data collection strategy has a larger impact on performance for the simplified model.
\textit{continue} achieves the best performance.
Compare these performances with those for the simplified model in Table~\ref{tab:results_nonbeam_aware}.
For larger beams, the improvements achieved by beam-aware training are much larger than those achieved by non-beam-aware ones.
For example, $92.69$ versus $82.41$ for \textit{continue} with $k = 8$, where in the first case it is trained in a beam-aware manner ($k=8$ for both training and decoding), while in the second case, beam search is used only during decoding ($k = 1$ during training but $k = 8$ during decoding).
This shows the importance of training with beam search and exposing the model to its mistakes.
Without beam-aware training, the model is unable to learn to use the beam effectively.
Check Figure~\ref{fig:beam_vaswani} for the evolution of the training and validation accuracies with training epoch for beam-aware training.

\subsection{Comparing surrogate losses}
\label{ssec:comp_surrogate_losses}

We train both models with \textit{continue} and vary the surrogate loss and beam size.
Results are presented in Table~\ref{ssec:surrogate_losses}.
Perceptron losses (e.g., \textit{perceptron (first)} and \textit{perceptron (last)}) performed worse than their margin-based counterparts (e.g., \textit{margin (last)} and \textit{cost-sensitive margin (last)}).
\textit{log loss (beam)} yields poor performances for small beam sizes (e..g, $k=1$ and $k = 2$).
This is expected due to small contrastive sets (i.e., at most $k+1$ elements are used in \textit{log loss (beam)}).
For larger beams, the results are comparable with \textit{log loss (neighbors)}.

\subsection{Additional design choices}
\label{sec:discussion}

\paragraph{Score accumulation}

The scoring function was introduced as a sum of prefix terms.
A natural alternative is to produce the score for a neighbor without adding it to a running sum, i.e., $s(y_{1:j}, \theta) = \tilde s(y_{1:j}, \theta)$ rather than $s(y_{1:j}, \theta) = \sum_{i = 1} ^j \tilde s(y_{1:i}, \theta)$.
Surprisingly, score accumulation performs uniformly better across all configurations.
For the main model, beam-aware training degraded performance with increasing beam size.
For the simplified model, beam-aware training improved on the results in Table~\ref{tab:results_nonbeam_aware}, but gains were smaller than those with score accumulation.
We observed that the LM LSTM failed to keep track of differences earlier in the supertag sequence, leading to similar scores over their neighbors.
Accumulating the scores is a simple memory mechanism that does not require the LM LSTM to learn to propagate long-range information.
This performance gap may not exist for models that access information more directly (e.g., transformers~\citep{vaswani2017attention} and other attention-based models~\citep{bahdanau2014neural}).
See the appendix for Table~\ref{tab:results_score_accumulation} which compares configurations with and without score accumulation.
Performance differences range from $1$ to $5$ absolute percentage points.

\paragraph{Update on all beams}
The meta-algorithm of~\citet{negrinho2018learning} suggests inducing losses on every visited beam as there is always a correct action captured by appropriately scoring the neighbors.
This leads to updating the parameters on every beam.
By contrast, other beam-aware work updates only on beams where the transition leads to increased cost (e.g., \citet{daume2005learning} and \citet{andor2016globally}).
We observe that always updating leads to improved performance, similar to the results in Table~\ref{tab:results_losses} for perceptron losses.
We therefore recommend inducing losses on every visited beam.
See the appendix for Table~\ref{tab:results_only_on_cost_increase}, which compares configurations trained with and without updating on every beam.

\section{Related work}
\label{sec:related_work}

Related work uses either imitation learning (often called learning to search when applied to structured prediction) or beam-aware training.
Learning to search~\cite{daume2009search-based, chang2015learning, goldberg2012dynamic, bengio2015scheduled, negrinho2018learning} is a popular approach for structured prediction.
This literature is closely related to imitation learning~\cite{ross2010efficient, ross2011reduction, ross2014reinforcement}.
\citet{ross2011reduction} addresses exposure bias by collecting data with the learned policy at training time.
\citet{collins_incremental_2004} proposes a structured perceptron variant that trains with beam search, updating the model parameters when the correct hypothesis falls out of the beam.
\citet{huang2012structured} introduces a theoretical framework to analyze the convergence of early update.
\citet{zhang2008tale} develops a beam-aware algorithm for dependency parsing that uses early update and dynamic oracles.
\citet{goldberg2012dynamic, goldberg2013training} introduce dynamic oracles for dependency parsing.
\citet{ballesteros2016training} observes that exposing the model to mistakes during training improves a dependency parser.
\citet{bengio2015scheduled} makes a similar observation and present results on image captioning, constituency parsing, and speech recognition.
Beam-aware training has also been used for speech recognition~\cite{collobert2019fully, baskar2019promising}.
\citet{andor2016globally} proposes an early update style algorithm for learning models with a beam, but use a log loss rather than a perceptron loss as in \citet{collins_incremental_2004}.
Parameters are updated when the golden hypothesis falls out of the beam or when the model terminates with the golden hypothesis in the beam.
\citet{wiseman2016sequence} use a similar algorithm to \citet{andor2016globally} but they use a margin-based loss and reset to a beam with the golden hypothesis when it falls out of the beam.
\citet{edunov2017classical} use beam search to find a contrastive set to define sequence-level losses.
\citet{goyal2017continuous, goyal2019empirical} propose a beam-aware training algorithm that relies on a continuous approximation of beam search.
\citet{negrinho2018learning} introduces a meta-algorithm that instantiates beam-aware algorithms based on choices for beam size, surrogate loss function, and data collection strategy.
They propose a DAgger-like algorithm for beam search.

\section{Conclusions}
\label{sec:conclusions}

Maximum likelihood training of locally normalized models with beam search decoding is the default approach for structured prediction.
Unfortunately, it suffers from exposure bias and does not learn to use the beam effectively.
Beam-aware training promises to address some of these issues, but is not yet widely used due to being poorly understood.
In this work, we explored instantiations of the meta-algorithm of \citet{negrinho2018learning} to understand how design choices affect performance.
We show that beam-aware training is most useful when substantial uncertainty must be managed during prediction.
We make recommendations for instantiating beam-aware algorithms based on the meta-algorithm, such as inducing losses at every beam, using log losses (rather than perceptron-style ones), and preferring the \textit{continue} data collection strategy (or \textit{reset} if necessary).
We hope that this work provides evidence that beam-aware training can greatly impact performance and be trained stably, leading to their wider adoption.

\section*{Acknowledgements}

We gratefully acknowledge support from 3M | M*Modal.
This work used the Bridges system, which is supported by NSF award number ACI-1445606, at the Pittsburgh Supercomputing Center (PSC).

\bibliographystyle{acl_natbib}