DeepXML: A Deep Extreme Multi-Label Learning Framework Applied to Short Text Documents

Abstract: Scalability and accuracy are well recognized challenges in deep extreme multi-label learning where the objective is to train architectures for automatically annotating a data point with the most relevant subset of labels from an extremely large label set. This paper develops the DeepXML framework that addresses these challenges by decomposing the deep extreme multi-label task into four simpler sub-tasks each of which can be trained accurately and efficiently. Choosing different components for the four sub-tasks allows DeepXML to generate a family of algorithms with varying trade-offs between accuracy and scalability. In particular, DeepXML yields the Astec algorithm that could be 2-12% more accurate and 5-30x faster to train than leading deep extreme classifiers on publically available short text datasets. Astec could also efficiently train on Bing short text datasets containing up to 62 million labels while making predictions for billions of users and data points per day on commodity hardware. This allowed Astec to be deployed on the Bing search engine for a number of short text applications ranging from matching user queries to advertiser bid phrases to showing personalized ads where it yielded significant gains in click-through-rates, coverage, revenue and other online metrics over state-of-the-art techniques currently in production. DeepXML's code is available atthis https URL

\section{Introduction}

\textbf{Objective}: This paper develops the DeepXML\xspace framework for deep extreme multi-label learning where the goal is to train architectures for automatically annotating a data point with the most relevant {\it subset} of labels from an extremely large label set. Note that multi-label learning generalizes multi-class classification which aims to predict a single mutually exclusive label. The objectives for developing DeepXML\xspace are threefold. First, DeepXML\xspace provides a framework for how to think about deep extreme multi-label learning that can not only be used to analyze seemingly disparate algorithms such as XML-CNN~\cite{Liu17} and MACH~\cite{Medini19} but which can also be used to derive significantly improved versions of such state-of-the-art deep extreme classifiers. Second, DeepXML\xspace can generate a family of new state-of-the-art algorithms obtained by combining various types of feature architectures with different classifiers in a scalable and accurate manner. In particular, DeepXML\xspace was used to derive Astec\xspace (standing for an Accelerated Short Text Extreme Classifier) which is specialized for extremely low-latency and high throughput short text applications as it can make billions of predictions per day and handle peak rates of up to a hundred and twenty thousand queries per second. Third, DeepXML\xspace provides an easy-to-use modular framework in which practitioners can design architectures for diverse applications by making minimal changes and simply plugging in the components of their choice rather than going back to the drawing board and designing a specialized architecture for each application from scratch.

\textbf{Short text applications}: This paper focuses on the classification of short text documents, having just 3-10 words on average, into millions of labels. In particular, this paper considers a range of short text applications including matching user search engine queries to advertiser bid phrases, predicting Wikipedia tags from a document's title, predicting frequently bought together products from a given retail product's name and showing personalized ads based on the set of webpage titles in a user's browsing history. Such applications pose a number of additional challenges to deep extreme classifiers as compared to long text documents having up to 200 words on average. First, the deep extreme classifier is forced to make predictions on the basis of just 3-10 words on average which is a significantly harder task than long text document classification. Second, short text corpora have fewer occurrences of each word than long text corpora thereby leading to a paucity of training data. For instance, the Bing datasets have millions of words which occur at most twice in the training set and thus learning good quality embeddings for such rare words can be challenging. Third, low-latency and high-throughput short text applications with billions of users require predictions in milliseconds on a CPU to keep operating and energy costs low. While these challenges seem daunting at the extreme scale, it is nevertheless important to design solutions for short text applications as they can benefit billions of people.

\textbf{Challenges in deep extreme classification}: Deep extreme classifiers jointly learn a feature architecture with an extremely large  classification layer leading to the following challenges. First, training and fine-tuning the feature architecture for millions of labels can be computationally expensive and can also lead to learning poor quality representations when training data is scarce. Second, both the forward prediction pass as well as gradient back-propagation become infeasible if the classification layer has costs that are linear in the number of labels (such as for a fully connected output layer). One might be tempted to address these challenges by replacing the classification layer with highly scalable non-deep learning based extreme classifiers~\cite{Prabhu18b, Jain19} which reduce the costs of the forward and backward pass to logarithmic in the number of labels. This is achieved by learning a sub-linear search data structure based on graphs~\cite{Jain19}, trees~\cite{Prabhu18b,Khandagale19,Wydmuch18,Prabhu14,Jain16,Prabhu18}, hashes~\cite{Tagami17,Siblini18a} or clustering~\cite{Bhatia15} in a fixed feature space such as a bag-of-words representation or fixed pre-trained embeddings. Unfortunately, such algorithms cannot be directly used to replace the classification layer in the deep learning setting as the feature representation changes with every mini-batch update. This necessitates the frequent recomputation or updation of the sub-linear search data structure on the updated features which can be prohibitively expensive. As such, deep learning presents additional statistical and computational challenges to extreme classification.

\textbf{DeepXML\xspace}: DeepXML\xspace addresses these challenges by decomposing the deep extreme learning task into the following four sub-tasks or modules each of which can be constrained to be learnt in log-time while maintaining accuracy. In Module I, an intermediate feature representation is learnt using an application-appropriate feature architecture trained on a simpler surrogate task. The objective in Module I is to efficiently learn a near-final feature representation which can be fixed and used to eliminate all but a logarithmic number of the hardest negative labels for each data point. Then, in Module II, a graph, tree, hash or clustering based sub-linear search data structure is trained just once on the fixed intermediate features to shortlist the hardest negative labels for each data point in log-time. The motivation is to reduce the problem for any given data point from an extreme task with millions of labels to a traditional classification task with just hundreds of labels. Note that this can be achieved with minimal loss in accuracy as uninformative negative labels can be discarded since they don't contribute to the final solution~\cite{Jain19}. Module III then transfers the intermediate features to learn final features for the given extreme task subject to the constraint that the final features don't lie very far away from the intermediate features. This is done so as to get all the accuracy gains of fine-tuning for the task at hand while ensuring that the hardest negatives continue to lie in the shortlisted label set. Module IV jointly learns an extreme classifier along with the final features in log-time using just the shortlisted labels. Varying the choices for the component including the feature architecture, the surrogate task, the sub-linear search data structure, the negative sampling procedure, the transfer mechanism and the extreme classifier leads to a family of state-of-the-art algorithms including Astec\xspace, DECAF~\cite{Mittal21}, GalaXC~\cite{Saini21}, ECLARE~\cite{Mittal21b}, {\it etc}.

\textbf{Astec\xspace}: Astec\xspace was derived from DeepXML\xspace specifically for short text applications. It employed a low capacity feature architecture which could be learnt accurately from limited training data. This allowed Astec\xspace to be 2-12\% more accurate than leading deep extreme classifiers on publicly available benchmark datasets while also being up to 20\% more accurate than state-of-the-art techniques for matching user queries to advertiser bid phrases on Bing datasets. Furthermore, by leveraging the DeepXML\xspace framework, Astec\xspace could be 5-30$\times$ faster to train than leading deep extreme classifiers and could efficiently scale to problems involving 62 million labels. Finally, Astec\xspace could make predictions in milliseconds on a CPU and could therefore make billions of predictions per day, with peak rates of 120,000 queries per second, using commodity hardware. As a result, Astec\xspace increased the click-through-rate by 6.5\% over state-of-the-art techniques in production for showing personalized ads based on the webpage titles and URLs in a user's browsing history. Similarly, Astec\xspace yielded an increase of 1.6\% in revenue per thousand queries, a 2.9\% increase in match quality and an 8.6\% increase in query coverage over leading techniques in production for matching user queries to advertiser bid phrases. These represent just two of the multiple short text applications for which Astec\xspace resulted in a significant increase in key metrics on Bing.

\section{Related Work} \label{sec:lit}

\textbf{DeepXML\xspace and extreme Classification}: Much progress has been made in extreme classification~\citep{Agrawal13,Prabhu14,Mineiro15,Jain16,Jain19,Yen16,Yen18a,Babbar17,Babbar19,Prabhu18,Prabhu18b,Khandagale19,Bhatia15,Tagami17,Wydmuch18,Jasinska16} for fixed representations such as bag-of-words features and pre-trained embeddings. Unfortunately, these algorithms cannot be used directly for deep extreme classification as the sub-linear search data structure they rely on for scalability needs to be frequently recomputed due to the change in features with every mini-batch update. As a result, specialized deep extreme classifiers~\citep{Liu17, Medini19, You18, Jernite17, Ye20, Yuan20, Chalkidis19} have been developed and many of these can be analysed and improved in the DeepXML\xspace framework.

\textbf{Astec\xspace and short text extreme classification}: Of particular relevance to Astec\xspace and this paper are specialized extreme classifiers that have been developed for short text applications including Slice~\citep{Jain19} for recommending related queries on Bing and MACH~\citep{Medini19} for searching Amazon retail products. Slice is a highly scalable classifier for fixed pre-trained embeddings and, unfortunately, cannot be used for deep extreme classification as already mentioned. Astec\xspace could therefore be up to 20\% more accurate than Slice on pre-trained CDSSM~\citep{Huang13}, FastText~\citep{Joulin17} or BERT~\citep{Devlin19} embeddings (see Section~\ref{sec:results}). Also note that Slice could always be incorporated into DeepXML\xspace's second module if desired. Furthermore, Astec\xspace can be seen as a generalization of MACH when analyzed through the DeepXML\xspace framework. In particular, MACH stops training after the first DeepXML\xspace module and therefore has to learn a large ensemble of base classifiers to compensate. Astec\xspace could be up to 12\% more accurate and orders of magnitude faster to train as it learnt a single base classifier by efficiently and accurately leveraging the DeepXML\xspace  framework with all four modules.

\begin{figure}
    \centering
    \includegraphics[width=0.92\columnwidth, clip]{Figs/astec_fig.pdf}
    \caption{Astec\xspace's architecture: Please refer to the text for details. Best viewed under magnification and in color.}
    \label{fig:arch}
\end{figure}

\textbf{Short text applications}: Apart from extreme classification, four classes of techniques have been developed for matching user queries to advertiser bid phrases and the other short text applications considered in this paper. The first class of techniques leverage additional sources of information such as landing pages~\cite{Yejin10}, web search results and other queries~\cite{Broder09} and are therefore beyond the scope of this paper. The second class of techniques are based on generative models~\citep{Gao12,Zhou19,Lee18,Lian19} that synthesize bid phrases for a given query. Unfortunately, unconstrained synthesis can result in poor quality bid phrases being generated while constraining them through tries~\cite{Lian19} or other approaches can limit bid phrase coverage. The third class of techniques use graph neural networks, random walks and other graph processing methods on the query-bid phrase, query-url, query-token or query-query graphs~\cite{Ioannis08,Mei08}. Features based on such graphs can be readily incorporated into the DeepXML\xspace framework by leveraging graph neural networks in the first and third modules~\cite{Saini21,Mittal21b}. The fourth class of techniques embed both queries and bid phrases into the same space using a Siamese network~\cite{Huang13,Bai18,Changlarge20,Reimers19} or two-tower model~\cite{Yi19,Krichene18}, and make predictions for a novel query by retrieving the nearest embedded bid phrases with highest cosine similarity or other metrics. Section~\ref{sec:results} demonstrates that Astec\xspace could yield significant improvements in online metrics for multiple short text applications over large  ensembles of state-of-the-art methods for each class of techniques, that are currently in production in Bing.

\section{The DeepXML\xspace framework}\label{sec:fm}
\textbf{Notation}: Let $L$ be the number of labels and $V$ be the vocabulary size if the input is a text document. Each of the $N$ training points is then represented as $(\vx_i,\vy_i)$, where $\vx_i$ is a data-point represented  either as a dense vector, sequence of tokens or sparse bag-of-tokens depending on the application and $\vy_i \in \bc{-1,+1}^L$ is the ground truth label vector with $y_{il} = +1$ if label $l \in [L]$ is relevant to data point $i$ and $y_{il} = -1$ otherwise.

\textbf{Components}: DeepXML\xspace has the following components. First, a feature architecture $\cZ$ that maps a data point $\vx_i$ onto a dense $D$ dimensional representation $\hat\vx_i$, {\it i.e.,} $\cZ: \vx_i \rightarrow \hat\vx_i \in \bR^D$. Second, a surrogate objective to train intermediate feature representations. Third, a sub-linear search structure and negative sampling procedure. Fourth, a transfer mechanism to obtain final feature representations and, finally, parameters $\vW$ of a classifier model to make final predictions. Section~\ref{sec:alg} includes efficient choices made by the Astec\xspace algorithm for all these components.

\textbf{Summary}: Given a task-specific loss function $\ell$ that measures the accuracy of a classifier model on a specific label, an ideal training strategy would train $\cZ, \vW$ jointly, taking into account all positive and all negative labels of all $N$ training points, {\it i.e.,} solve $\argmin_{\cZ,\vW} \cL(\cZ,\vW)$
\[
 \cL(\cZ,\vW) = \sum_{i=1}^N\sum_{l=1}^L \ell(\vx_i, y_{il}; \cZ, \vW)
\]
However, this strategy requires jointly learning $\Om{VD + LD}$ parameters (to describe $\cZ$ and $\vW$) using the objective $\cL$ for which computing even a single gradient takes $\Om{NLD}$ time (as the objective has $NL$ terms) and is prohibitive for even moderate scale datasets. This remains true even if the objective uses a loss that does not decompose over the labels such as the F-measure, {\it etc}. To remedy this, DeepXML\xspace proposes a \emph{modular} strategy that effectively scales to tasks with millions of labels and data points. The core idea is to solve a much cheaper ``surrogate'' task (in \textbf{Module I}) and use this solution to identify (in \textbf{Module II}) shortlists $\hat\cN_i$ of say $\bigO{\log L}$ negative labels for each data point $i \in [N]$ ({\it i.e.,} $y_{il} = -1$ for all $l \in \hat\cN_i$) that offer a good approximation to the objective, {\it i.e.,} for all $i \in [N]$, it is the case that
\[
\sum_{l: y_{il} = -1} \ell(\vx_i, y_{il}; \cZ, \vW) \approx \sum_{l \in \hat\cN_i} \ell(\vx_i, y_{il}; \cZ, \vW).
\]
For example, if using a margin loss function such as the hinge loss, a shortlist of all margin violator negative labels would suffice. More generally, since these labels are expected to incur high loss values and also be most likely to be confused for a positive label, they are commonly referred to as \emph{hard negatives}~\cite{Jain19}. After performing a feature transfer from the surrogate to the original task (in \textbf{Module III}), final training is done (in \textbf{Module IV}) to learn $\cZ, \vW$ by solving a much less expensive optimization problem with the objective $\hat\cL$
\[
\hat\cL(\cZ,\vW) = \sum_{i=1}^N\sum_{l \in \hat\cN_i \cup \cP_i} \ell(\vx_i, y_{il}; \cZ, \vW),
\]
where for any $i \in [N]$, $\cP_i := \bc{l: y_{il} = +1}$ denotes the set of positive labels of that data point. Minimizing $\hat\cL(\cZ,\vW)$ is expected to yield parameters that resemble those obtained by minimizing $\cL(\cZ,\vW)$ due to the way hard negatives are designed. However, each module in the DeepXML\xspace framework can be executed in no more than $\bigO{ND\log L}$ time if careful choices are made. Finally, note that DeepXML\xspace supports learning an ensemble of multiple learners by training them from scratch or else simply training a re-ranker to cut down the training costs~(See section~\ref{sec:alg}).

\textbf{Flexibility}: DeepXML\xspace offers the flexibility of tackling a range of disparate applications with diverse inputs ranging from documents to images to graphs by letting practitioners plug in the components of their choice with minimal effort and without having to redesign the entire architecture from scratch for each application. Various feature architectures, ranging from convolutions to transformers, as well as classifier architectures, ranging from 1-vs-All to trees, can be plugged in depending on the accuracy and scalability requirements. Furthermore, DeepXML\xspace offers the flexibility of using diverse types of deep extreme classifiers, ranging from XML-CNN to MACH, by casting them in the proposed framework. Finally, DeepXML\xspace also offers the flexibility to incorporate metadata such as label features~\cite{Mittal21} or label correlations~\cite{Mittal21b} into the various modules to obtain superior performance. This allowed DeepXML\xspace to be used for a number of applications including text ads, product ads, rich ads, native ads, retail product recommendation, news recommendation, personalized query recommendation, {\it etc}.

\subsection{Module I: Intermediate representation}
In this module, an intermediate feature architecture $\cZ^0$ is trained using a \emph{surrogate} task. Several considerations need to be kept in mind while choosing a feature architecture and surrogate task. The feature architecture $\cZ^0$ should be trainable from available data especially  with respect to rare tokens and rare labels. $\cZ^0$ should also be able to efficiently embed data points $\vx_i \mapsto \hat\vx_i^0$, ideally in time $\bigO{cND\log L}$ where $c$ is some constant depending on the architecture, to satisfy the requirements of low-latency and high-throughput applications. Simultaneously, the surrogate task should be solvable faster than solving $\argmin_{\cZ,\vW} \cL(\cZ,\vW)$, ideally offering gradient computations in time $\bigO{cND\log L}$. It should also promote learning of a feature architecture $\cZ^0$ whose data point representations, say $\hat\vx^0_i = \cZ^0(\vx_i)$, closely resemble those offered by $\cZ$, say $\hat\vx_i = \cZ(\vx_i)$. Recall that $\cZ$ is the feature architecture that could have been learnt by directly solving $\argmin_{\cZ,\vW} \cL(\cZ,\vW)$. Doing so ensures that $\hat\vx^0_i$ and $\hat\vx_i$ have approximately the same nearest neighbors and label shortlists generated using $\hat\vx^0_i$ (in Module II) are apt proxies for those that could have been generated using $\hat\vx_i$.

The feature architecture $\cZ^0$ may be learnt in several ways. Unsupervised surrogate tasks include skip-gram models~\cite{Joulin17,Mikolov13}, next sentence prediction~\cite{Devlin19}, masked language modeling~\cite{Devlin19}, and multi-task learning~\cite{Liu19}. While scalable, architectures learnt using unsupervised training may lie far away from, and were empirically found to be 4-5\% less accurate than, those learnt using supervised training~(see Table~6 in the \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace). Supervised learning techniques cut down the training cost by reducing the effective number of labels to $\hat{L} \ll L$. This can be implicitly done by sampling a small sub-set of labels in an online manner, say by mini-batch negative sampling~\cite{Guo19}, or explicitly by label selection, label projection or label clustering. Label selection methods~\cite{Bi13,Barezi19} select a subset of labels $\kL$ but were found to offer sub-optimal accuracies. These methods suffer from poor token coverage: if the set of data points tagged with at least one label in $\kL$, {\it i.e.,} $\bc{i : y_{ij} = +1 \text{ for any } j \in \kL}$ do not cover all vocabulary tokens, then either pre-trained token embeddings have to be externally sourced or else performance may suffer. Ensuring perfect token coverage is usually challenging -- nearly 1M labels had to be selected in order to cover 95\% of the vocabulary on the Q2B-3M dataset which defeats the very purpose of creating a surrogate task. Low-rank projection methods~\cite{Medini19,Bhatia15,Tagami17,Mineiro15} project labels on to a low-dimensional space as $\hat{\vy}_i=\vP {\vy}_i$, where $\vP \in \bR^{\hat{L}\times L}$ is a projection matrix. These are theoretically well understood but offered accuracies similar to label selection approaches in our experiments. Label clustering approaches~\cite{Prabhu18b,Khandagale19,You18,Siblini18a,Mittal21} cluster labels and treat each cluster as a \emph{meta label} using either explicit label features (using state-of-the-art encoders~\cite{Devlin19,Joulin17}) or else using indirect label representations~\cite{Jain19,Prabhu18b}. Label clustering-based approaches were found to offer the best performance in the experiments reported in section~\ref{sec:results} but other approaches might be more suitable for other applications.

\subsection{Module II: Negative sampling}
\label{sec:phase-2}
In this module, $\cZ^0$ is used to obtain intermediate representations $\hat\vx^0_i = \cZ^0(\vx_i)$ for all data points which are then used to obtain shortlists $\hat\cN_i \subset [L]$ of the $\bigO{\log L}$ most confusing or ``hardest'' negative labels, for every data point $i \in [N]$. Such deliberate hard negative mining outperforms cheaper alternatives such as mini-batch sampling~\cite{Guo19}, or sampling negatives from the power law distribution~\cite{Mikolov13} (see Fig.~2 in the \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace). This is because the probability of choosing the most confusing negative labels is negligible when the number of labels is in the millions. Two main considerations need to be kept in mind while creating these shortlists. First, the shortlist should contain negative labels most likely to be confused for positive labels to offer concise and directed signals while training the classifiers. Second, the shortlist should be computable for every data point $i \in [N]$ in time sub-linear in $L$. Several options exist including those based on graphs~\cite{Jain19,MalkovY16}, trees~\cite{Khandagale19,Prabhu18b,Siblini18a}, clusters~\cite{Mittal21,You18} or hashing~\cite{Rawat19,Medini19,Bhatia15} that achieve sub-linear time negative sampling when working with fixed features. Note that it is possible to reuse the same data structure to shortlist $\bigO{\log L}$ labels during prediction as well.

\subsection{Module III: Transfer learning}
In this module, the final form of the feature architecture $\cZ$ is created by adapting the intermediate architecture $\cZ^0$. A non-trivial transfer may be required since $\cZ^0$ is tuned for the surrogate task and not the original task. If this transfer involves any re-parametrization, any free parameters thus introduced are either trained separately or else fine-tuned jointly in Module-IV (see below). Three main considerations need to be kept in mind while performing this feature transfer. First, $\cZ$ should not impose significant computational overhead as compared to $\cZ^0$. Second, additional parameters introduced in $\cZ$ should be trainable from available training data, especially with respect to rare tokens, in time $\bigO{ND\log L}$. Finally, feature representations offered by $\cZ$ should not lie too far away from those offered by $\cZ^0$ so that hard negatives discovered in Module II continue to remain relevant for classifier training. Sophisticated transfer learning techniques~\cite{Pan2010,Wang17,Wei19,Saini21} could be deployed in this module and in particular, a method aiming for higher accuracies could choose to fine-tune $\cZ$ in its entirety across Modules I-IV, albeit at greater computational expense.

\subsection{Module IV: Classifier learning}
\label{sec:phase-4}
In this module, the classifier's parameters $\vW$ (and optionally, any free parameters in $\cZ$) are (jointly) learnt using an approximate objective that considers only the positive labels of a data point, {\it i.e.,} $\cP_i := \bc{l: y_{il} = +1}$ and the shortlisted negative labels $\hat\cN_i$.
	\[
	\hat\cL(\cZ,\vW) = \sum_{i=1}^N\sum_{l \in \hat\cN_i \cup \cP_i} \ell(\vx_i, y_{il}; \cZ, \vW)
	\]
Various classifiers including 1-vs-All~\cite{Babbar17,Babbar19,Prabhu18b,Jain19,Khandagale19,You18}, tree~\cite{Prabhu14,Jain16,Jernite17,Wydmuch18} or $k$-NN~\cite{Tagami17,Bhatia15} classifiers may serve as suitable choices. Note that for most commonly used loss functions, computing gradients $\nabla \hat\cL$ takes only $\bigO{ND \log L}$ time since data points typically contain only logarithmically many positive labels, {\it i.e.,} $\abs{\cP_i} \approx \bigO{\log L}$ and $\abs{\hat\cN_i} \leq \bigO{\log L}$ by design. Modules I and IV can use distinct loss functions aimed at promoting recall and precision respectively.

\section{The Astec\xspace algorithm}
\label{sec:alg}

\textbf{Module I}: Astec\xspace uses the following feature architecture that can be learnt from limited training data and be computed in 30 $\mu s$ on a CPU thereby meeting the accuracy and latency requirements of short text applications. In particular, Astec\xspace operates with sparse bag-of-words representations for documents, {\it i.e.} $\vx_i \in \bR^V$, and learns $D$ dimensional embeddings for each vocabulary token $\ve_t \in \bR^D: t \in [V]$. The intermediate features used by Astec\xspace are of the form $\hat\vx^0 := \cZ^0(\vx) = \vv + \delta\vv$ where $\vv := \text{ReLU}\br{\sum_{t=1}^{V} x_t\cdot \ve_t}$, {\it i.e.} a ReLU non-linearity over the TF-IDF weighted linear combination of the learnt token embeddings, and $\delta\vv := \text{ReLU}(\vR^0\vv)$ where $\vR^0 \in \bR^{D \times D}$ is a residual matrix. The final features are of the form $\hat\vx := \cZ(\vx) = \vv + \Delta\vv$ where $\Delta\vv := \text{ReLU}(\vR\vv)$ and $\vR \in \bR^{D \times D}$. Note that $\cZ^0(\vx)$ and $\cZ(\vx)$ share the component $\vv$ and only differ in the residual component $\delta\vv, \Delta\vv$. Restricting the spectral norms of $\vR^0$ and $\vR$ encourages a high fidelity transfer in Module III, {\it i.e.,} $\hat\vx \approx \hat\vx^0$. Thus, Astec\xspace's feature architectures $\cZ^0, \cZ$ are parametrized using token embeddings $\vE = [\ve_1,\ldots, \ve_V] \in \bR^{D \times V}$ and $\vR^0, \vR \in \bR^{D \times D}$.

Astec\xspace adopts label clustering for its surrogate task as this was empirically observed to lead to the highest classification accuracies while keeping training time to within a few hours on a single GPU across all datasets in Table~\ref{tab:results_main}. Label centroids, defined as $\vmu^s_l= \frac{\hat\vmu^s_l}{\norm{\hat\vmu^s_l}_2}$, were used to cluster the labels, where, $ \hat\vmu^s_l = \frac1{|\cP_l|} \sum_{i \in \cP_l} \vx_i$, and $\cP_l := \bc{ i : y_{il} = +1}$ is the set of documents for which label $l$ is relevant. The balanced 2-means++ algorithm~\cite{Prabhu18b} was used to recursively cluster the labels into balanced partitions until $\hat{L}$ clusters were obtained. These clusters were treated as meta-labels and new (meta) label vectors $\hat\vy_i \in \bc{-1,+1}^{\hat L}$ were created for each training document as $\hat{y}_{ik} = +1$ for documents $i$ tagged with at least one label in cluster $k$ and $\hat{y}_{ik} = -1$ otherwise.
Using label correlations for improved clustering led to a 2\% improvement in recall as compared to Parabel's clusters which ignored label correlations. Predicting the relevant clusters for a given document was taken as the surrogate task for Module I. 1-vs-All classifiers, parametrized as $\hat\vW = [\hat\vw_1, \ldots, \hat\vw_{\hat{L}}] \in \bR^{D \times \hat{L}}$ and $\cZ^0$ were trained jointly by solving
\begin{align}
\label{equ:surr_opt}
\argmin_{\cZ^0, \hat\vW} & \sum_{i = 1}^N \sum_{k = 1}^{\hat{L}} \log\br{1 + \exp\br{-\hat{y}_{ik}\cdot\hat\vw_k^\top\hat\vx^0_i}
\end{align}
subject to the constraint $\sup_{\vu \neq \vzero} \norm{\vR^0\vu}_2/\norm\vu_2 \leq \lambda$, and with $\hat\vx^0_i = \cZ^0(\vx_i)$. Note that $\hat\vW, \vR^0$ are distinct from their counterparts $\vW, \vR$ used to solve the original task in Module IV. Indeed, $\hat\vW, \vR^0$ were summarily discarded after completion of Module I. However, $\vE$ was shared by $\cZ^0$ and $\cZ$ and hence, retained. The power iteration method~\cite{Miyato18} was used to constrain the spectral norm of $\vR^0$. Adding dropout after each ReLU layer was found to be as effective as using an explicit regularizer $\frac12\sum_l\hat\vw_l^\top\hat\vw_l$. Using $\hat L \leq \bigO{\log L}$, efficient parallelization on a GPU and the Adam optimizer controlled the complexity of the surrogate task. In practice, $\cZ^0$ could be trained in a couple of hours on a single GPU across all datasets in Table~\ref{tab:results_main}.

\textbf{Module II}: The sampling cost of techniques discussed in Section~\ref{sec:phase-2} can be larger than $O(ND\log L)$ if document features $\hat\vx^0_i$ keep changing during training due to $\cZ^0$ getting updated with each mini-batch. Astec\xspace tackles this challenge by freezing the intermediate representations after completion of Module I before training a sub-linear search data structure to discover hard negatives using multiple label representations. For simplicity, Astec\xspace utilizes $\vv$ to sample hard negatives which were found to be just as effective as those sampled using $\hat{\vx}^0$. A four-step procedure was adopted to sample hard negatives. First, multiple representations were computed for each label (refer to section~A.4 in the \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace). Then, two Approximate Nearest Neighbour Structures~(ANNS)~\cite{MalkovY16} were deployed: $\text{ANNS}^x$ over $\bc{\vv_i: i \in [N]}$ and $\text{ANNS}^\mu$ over label representations defined as $\vmu^0_l= \frac{\hat\vmu^0_l}{\norm{\hat\vmu^0_l}_2}$ where $ \hat\vmu^0_l = \frac1{|\cP_l|} \sum_{i \in \cP_l} \vv_i$. Recall that $\vv_i$ is the component shared by $\hat\vx^0_i$ and $\hat\vx_i$. The graphs were queried to generate $\bigO{\log L} \leq 500$ negative labels for each $i \in [N]$ as: (a) $\hat\cN^x_i = \bc{l : j \in \text{ANNS}^x(\vv_i), y_{jl} = +1, y_{il} = -1 }$, and (b) $\hat\cN^\mu_i = \bc{l : \vmu^0_l \in \text{ANNS}^\mu(\vv_i), y_{il} = -1}$. Adding $\bigO{\log L} \leq 50$ random negatives to the shortlist (to account for minor distortions between $\vv_i$ and $\hat\vx_i$) increased prediction accuracy.

Astec\xspace's shortlisting strategy could recall 5\% more relevant labels and the overall predictions could be 13\% more accurate than Slice. Theorem~\ref{thm:sim} indicates that it was justified to generate ANNS shortlists using $\vv_i$ rather than the final features, i.e. $\hat{\vx}_i$. This was also empirically validated as more than 87\% overlap was observed between the label shortlists computed on $\vv_i$ and those computed on $\hat{\vx}_i$. Astec\xspace's negative sampling step incurs a cost of $O(ND \log L)$, assuming $N \leq L^{\bigO1}$ and took only a few minutes on most datasets and at most 2 hours on the AmazonTitles and Q2B datasets with up to 3 million labels. Note that the shortlist could be further extended with the labels selected based on explicit label features, however label features are beyond the scope of Astec\xspace.

\begin{theorem}
\label{thm:sim}
Let $\lambda$ be the spectral norm of $\vR$, $\cP_l = \{i | y_{il} = +1\}$ be the set of positive training points for the label $l$, $\epsilon_l = (1+\lambda\sqrt{|\cP_l|})^2-1$, and $\vmu_l$ be the label representation that could have been computed using final features as $\vmu_l= \frac{\hat\vmu_l}{\norm{\hat\vmu_l}_2}$ where $ \hat\vmu_l = \frac1{|\cP_l|} \sum_{i \in \cP_l} \hat\vx_i$ and $\hat\vx_i = \cZ(\vx)$. Then: (a) $\norm{\hat\vx_i - \vv_i}_2 \leq \lambda\norm{\vv_i}_2 \; \forall i$, i.e. the final and intermediate features for any document can be brought arbitrarily close by restricting $\lambda$; and (b) $\frac{C(\vv_i, \vmu^0_l)}{1+\epsilon_l} \leq C(\hat\vx_i, \vmu_l) \leq C(\vv_i, \vmu^0_l) + \epsilon_l$, i.e. the cosine similarity $C(\va, \vb) := \frac{\va^\top \vb}{\norm\va_2\norm\vb_2}$ used in the ANNS algorithm to determine the similarity between any document $i$ and label $l$ can also be made arbitrarily close between the intermediate and final representations by restricting $\lambda$, the spectral norm of $\vR$. [Please refer to~A.6 in the \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace for a proof.]
\end{theorem}

\textbf{Module III}: Astec\xspace uses a final feature representation of the form $\hat\vx : = \cZ(\vx) = \vv + \Delta\vv$, where $\vv$ is shared with $\cZ^0$ and $\Delta\vv = \text{ReLU}(\vR\vv)$ for a residual matrix $\vR$. Note that $\vR$ is distinct from the residual matrix $\vR^0$  used in $\cZ^0$. Since $D^2 \ll VD$, this introduced a negligible increase in model parameters and computational complexity and $\Delta\vv$ could be computed in under 10 $\mu s$ on a CPU thereby meeting the low-latency constraints. The spectral norm of $\vR$ was restricted to be no greater than $\lambda \leq 1$ to ensure that feature representations offered by $\cZ$ did not lie far away from those offered by $\cZ^0$ (see Theorem~\ref{thm:sim}). Prediction accuracy was observed to degrade if the spectral norm was not restricted or if the residual block was replaced by a fully connected block~\cite{Medini19} since such steps allowed the final features to drift away from the intermediate features.

\textbf{Module IV}: Astec\xspace adopts the high capacity 1-vs-All classifier model and learns a $D$-dimensional classifier per label. Thus, its classifier is parametrized as $\vW = [\vw_1,\ldots, \vw_L] \in \bR^{D \times L}$. The classifier and residual block $\vR$ present in $\cZ$ were jointly trained in $O(ND \log L)$ time by restricting training to the positive and shortlisted negative labels for each data point. Specifically, for any $i \in [N]$, let $\cP_i := \bc{l: y_{il} = +1}$ be the set of positive labels of the data point, $\hat\cN_i = \hat\cN^\mu_i \cup \hat\cN^z_i$ be the negative labels shortlisted in Module II and let $\hat\cS_i := \hat\cN_i \cup \cP_i$. Then the approximate objective $\argmin_{\vR, \vW}\hat\cL(\vR,\vW)$ was solved where
\[
\hat\cL(\vR,\vW) = \sum_{i=1}^N\sum_{l \in \hat\cS_i} \log\br{1+\exp\br{-y_{il}\cdot\vw_l^\top\hat\vx_i}},
\]
subject to the constraint $\sup_{\vu \neq \vzero} \norm{\vR\vu}_2/\norm\vu_2 \leq \lambda$, where $\vW = [\vw_1,\ldots, \vw_L] \in \bR^{D \times L}$ are the 1-vs-All classifiers and $\hat\vx_i = \cZ(\vx_i)$. This formulation could be optimized efficiently using the Adam optimizer on a single P40 GPU in 0.25--5 hours on the datasets considered in this paper.

\textbf{Re-ranking the labels}: Astec\xspace improves its accuracy by learning a novel re-ranker as follows. Astec\xspace's predictions $ \hat{\vy}_i \in \bc{-1,+1}^L$ were obtained for each training point $i$. Excluding the true positives yielded a shortlist of negative labels that Astec\xspace found the most confusing for each point. A re-ranker was then trained to eliminate these mis-predictions by optimizing $\argmin_{\widetilde\cZ, \widetilde\vW}\tilde\cL(\widetilde\cZ, \widetilde\vW)$ where
\[
\tilde\cL(\widetilde\cZ, \widetilde\vW) = \sum_{i=1}^N \sum_{l \in \widetilde\cS_i} \log\br{1 + \exp\br{-y_{il}\cdot\widetilde\vw_l^\top\widetilde\vx_i}},
\]
where $\widetilde\cS_i = \bc{l : y_{il} = +1} \cup \bc{l : \hat{y}_{il} = +1}$, $\widetilde\vx_i = \widetilde\cZ(\vx_i)$ and $\widetilde\cZ$ had an architecture similar to $\cZ$ but with independent parameters $\widetilde\vE, \widetilde\vR$. This increased accuracy by $1$--$3\%$ with comparatively larger gains on larger datasets, with only a 10--20\% increase in training time.

\textbf{Log-time prediction}: Astec\xspace meets the latency requirements of online short text applications on even the largest datasets by making predictions in $O(D^2 + D \log L)$ time. Given a test document $\vx \in \bR^V$, the $O(\log L)$ most relevant labels $\hat\cS :=  \cN^\mu \cup \cN^x$ along with their base similarity scores $s^l(\vv)$ were shortlisted using $\vv := \text{ReLU}\br{\sum_{t=1}^{V} x_t\cdot \ve_t} \in \bR^D$. Next, the final features $\hat\vx = \cZ(\vx)$ and 1-vs-All classifiers were used to yield scores $\hat{y}_l = \alpha\sigma(\vw_l^\top\hat\vx) + (1-\alpha)\sigma(s^l(\vv))$ for shortlisted labels $l \in \hat\cS$ and $\hat{y}_l = 0$ otherwise where $\sigma$ is the sigmoid function and $\alpha \in [0, 1]$ is a hyperparameter. Final predictions were given by linearly combining the re-ranker scores $\tilde{y}_l = \sigma( \widetilde\vw_l^\top\widetilde\vx)$ with the base similarity scores as $\overline{y}_l = \beta \hat{y}_l + (1-\beta)\widetilde{y}_l$ if $l \in \hat\cS$ and $\overline{y}_l = 0$ otherwise, using a hyper-parameter $\beta \in [0,1]$.

\input{Tables/results_main.tex}
\section{Experiments} \label{sec:results}

\textbf{Datasets}: Results are presented on publically available benchmark short text datasets with up to 3 million labels for predicting frequently bought together Amazon items based on just the product title (AmazonTitles-670K and AmazonTitles-3M) as well as using a Wikipedia article's title to predict it's Wikipedia tags (WikiTitles-500K) as well as its related Wikipedia articles (WikiSeeAlsoTitles-350K). All datasets are available on the Extreme Classification Repository~\cite{XMLRepo}. Note that, even though the focus of this paper is on short text applications and Astec\xspace has been designed keeping their specific requirements in mind, results on benchmark long text datasets are presented in the \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace for completeness. Results are also presented on a proprietary Bing dataset with 3 million labels and 21 million training points for matching user queries to advertiser bid phrases (Q2B-3M). The dataset was created by mining Bing's click logs where each user's query was treated as a data point and clicked advertiser bid phrases became its labels. Table~3 in the \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace presents the data set statistics.

 \textbf{XML baselines}: The focus of this paper is on comparing Astec\xspace to MACH~\cite{Medini19} and Slice~\cite{Jain19} as they have been specifically designed for short text applications. Astec\xspace was also compared to other leading deep extreme classifiers including XML-CNN~\cite{Liu17}, XTransformers~\cite{Chang20} and AttentionXML~\cite{You18}. Furthermore, for the sake of completeness, results are presented for non-deep extreme classifiers including XT~\cite{Wydmuch18}, DiSMEC~\cite{Babbar17}, PfastreXML~\cite{Jain16}, Parabel~\cite{Prabhu18b}, Bonsai~\cite{Khandagale19} and AnneXML~\cite{Tagami17}. Implementations of all the baseline algorithms were provided by their authors. The hyper-parameters of these algorithms were set as suggested by their authors wherever applicable and by fine-grained validation otherwise. Results are only presented for those datasets to which an implementation could scale. Results could therefore not be reported for XTransformers  as it could not be trained on any of the datasets on a single GPU in a week. As is customary in extreme classification, results are presented for not only Astec\xspace but also an ensemble with 3 learners referred to as Astec\xspace-3. Astec\xspace's hyper-parameters and their settings on various datasets are discussed in the \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace.

\textbf{Bing baselines}:
The online flight results revealed the gains that Astec\xspace was able to achieve when added to a large ensemble of state-of-the-art techniques currently running in production on Bing including many leading techniques for query synthesis (constrained and unconstrained), graph based techniques (graph neural networks, random walks, sessions based methods, {\it etc.}), embedding methods (Siamese networks and two-tower models), extreme classifiers as well as techniques that leverage additional information (refer to Section~\ref{sec:lit}). Offline results for some of these techniques such as Simrank++~\cite{Ioannis08} and a BERT based sequence-to-sequence constrained synthesis model~\cite{Lian19, Devlin19} are presented on the Q2B-3M dataset for matching queries to bid phrases.

\textbf{Evaluation metrics}: Performance was evaluated using prec-ision@$k$ (P@$k$), and propensity scored precision@$k$ (PSP@$k$) which have been widely used in the extreme classification literature~\cite{Babbar17,Jain16}. Results on additional metrics such as nDCG@$k$ (N@$k$) and propensity scored nDCG@$k$ (PSN@$k$) have been included in \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace which also contains the definitions of all the metrics. All training times have been reported on a 24-core Intel Xeon 2.6 GHz machine with a single Nvidia P40 GPU unless stated otherwise.

\textbf{Table~\ref{tab:results_main} - Offline results}: Astec\xspace's main competitors were Slice and MACH as they are extreme classifiers that have been developed for low-latency short text applications since their features could be extracted in milliseconds on a CPU while all other architectures required a GPU. Slice is not a deep learning method and was therefore trained on the same FastText embeddings that were used to initialize Astec\xspace. Nevertheless, Astec\xspace was 2.29-17.94\% more accurate than Slice and 5.05-12.27\% more accurate  than MACH on the publically available Repository datasets and 20.14\% more accurate than Slice on the Bing Q2B-3M dataset. Unfortunately, MACH could not be trained on the Q2B-3M dataset in a week on a single GPU whereas Astec\xspace could be trained in 13 hours. On the Repository datasets, Astec\xspace was 5-9$\times$ faster to train than MACH. Similarly, none of the other deep extreme classifiers could be trained on the AmazonTitles-3M and Q2B-3M datasets. On the smaller datasets, Astec\xspace was 2.05-4.56\% more accurate and 9-29$\times$ faster to train than AttentionXML and XML-CNN respectively. Furthermore, Astec\xspace was at least 19\% more accurate than all other methods on the Bing Q2B-3M dataset including SimRank++ and a BERT based sequence-to-sequence model that had been specifically designed for matching queries to bid phrases. These results demonstrate that Astec\xspace's features could be learnt accurately from limited training data and that the DeepXML\xspace framework enabled Astec\xspace to be significantly more scalable than all other deep extreme classifiers. Finally, for the sake of completeness, Astec\xspace was compared to non-deep learning extreme classifiers on the Repository datasets where it could be up to 6.12\% more accurate and this increased marginally to 6.77\% for the Astec\xspace-3 ensemble.

\textbf{Online results from Bing flights}: Astec\xspace was able to efficiently train in 20 hours on 4$\times$P40 GPUs on various Bing internal datasets with up to 62 million labels that were far beyond the scaling capabilities of all other deep extreme classifiers. Furthermore, Astec\xspace's features could be extracted in microseconds on a CPU and its overall predictions made in a few milliseconds allowing it to make billions of predictions per day at peak rates of 120,000 queries per second on commodity hardware. This allowed Astec\xspace to be flighted for multiple short text applications on Bing with extremely low-latencies and high-throughputs including text ads, product ads, rich ads, native ads, retail product recommendation, news recommendation, personalized query recommendation, {\it etc}. Unfortunately, due to space constraints, online results when Astec\xspace was added to the ensemble of state-of-the-art techniques currently in production can only be presented for just two of these applications.

\textbf{Personalized ads}: Billions of users were shown personalized ads as they surfed the web based on their browsing history in the first application. Each user's intent was represented by the set of queries that the user could have potentially asked on Bing to reach the last $K$ webpages that they had browsed. Astec\xspace was used to predict the set of Bing queries from each visited webpage's title in milliseconds. A GRU was used to model the user's last $K$ states in near real time. The GRU was then used to select a single predicted query that was passed through the Bing pipeline to show ads to the user. Human expert evaluation revealed that Astec\xspace increased the number of excellent predictions by more than 20\% while reducing the number of fair and poor predictions by more than 10\% as compared to the ensemble in production containing leading extreme classification, deep learning and IR techniques. Astec\xspace was also able to achieve 100\% coverage when deployed online by making predictions for all visited webpages for all users within the latency and throughput constraints. This allowed Astec\xspace to increase the click-through-rate by 6.5\% and revenue by more than 5\%. Table~8 in the \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace shows examples comparing Astec\xspace's predictions to those of traditional approaches and demonstrates Astec\xspace's benefits over the state-of-the-art.

\textbf{Matching queries to bid phrases}: Astec\xspace treated each user's query as input and each advertiser bid phrase as a separate label in order to predict the set of bid phrases that could be matched to the user's query to show ads on Bing. Astec\xspace's offline accuracy was 19\% higher than that of other approaches~(see Table~\ref{tab:results_main}). This translated into an increase of 1.6\% in revenue per thousand queries, 2.9\% in match quality and 8.6\% increase in query coverage over a large ensemble of state-of-the-art techniques when deployed online. Table~\ref{tab:qual_analysis_q2b} lists Astec\xspace's predicted bid phrases for the query ``what is diabetes type''. As can be seen, Astec\xspace's predictions could be more accurate and diverse than those of other methods which could fixate on the phrase ``what is'' and thereby make many mispredictions.

\input{Tables/qualitative_analysis_q2b}

\textbf{Ablations}: Table~6 in the \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace presents the ablation results for each module validating Astec\xspace's design choices. First, much progress has been made in both designing and speeding up feature architectures such as BERT~\cite{Devlin19}, Roberta~\cite{Liu19roberta}, {\it etc.} that can easily be incorporated into the DeepXML\xspace framework, if desired. Unfortunately, doing so did not lead to accuracy gains on the short text extreme datasets and Astec\xspace's features were demonstrated to be at least 2\% more accurate than those based on CNNs~\cite{Liu17}, MLPs~\cite{Medini19} and BERT~\cite{Devlin19}. This indicates that Astec\xspace's features were more suitable for short text extreme classification and that more research is required on how to train and fine-tune transformer and other architectures when many features and labels have very limited training data. Second, training on Astec\xspace's surrogate task was demonstrated to be 1-4\% more accurate than training using self-supervision, label selection or label low-rank projection indicating the suitability of label clustering and the importance of a well designed surrogate task. Third, it was demonstrated that the negative sampling strategy proposed in Astec\xspace could be 13\% more accurate than the strategy in Slice. This further highlights the differences between Astec\xspace and Slice and indicates that, even if Slice could somehow have been trained jointly along with Astec\xspace's features, it would still have been significantly inferior to Astec\xspace. Finally, it was demonstrated that Astec\xspace's performance could not be further improved by replacing its classifier with alternatives such as DiSMEC or Parabel.

\textbf{State-of-the-art algorithms in the DeepXML\xspace framework}: DeepXML\xspace could be used to analyze and improve leading deep extreme classifiers thereby demonstrating its generality and usefulness. For instance, both XML-CNN and MACH could be recast into the DeepXML\xspace framework by replacing Astec\xspace's feature architecture by their CNN or MLP features. This led to accuracy gains of 1.8\% and 2.4\% and training speedups of 10$\times$ and 5$\times$ for XML-CNN and MACH respectively (see Table~4 in the \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace). Note that the original MACH algorithm stops training after the first DeepXML\xspace module, {\it i.e.,} after it has learnt its MLP features on the surrogate task, and therefore compensates by learning a large ensemble. Casting MACH into DeepXML\xspace allowed it to improve accuracy by fine-tuning its features for the extreme task at hand while speeding up training by learning just a single learner rather than an ensemble of 32 learners. Similarly, DeepXML\xspace increased XML-CNN's accuracy by fine-tuning its CNN features through the residual block in the third module and speeded up training by replacing its fully connected output layer with $O(L)$ complexity by a 1-vs-some classifier with $O(\log L)$ complexity in the fourth module. It should be reiterated that Astec\xspace's performance continued to remain superior to that of the improved MACH and XML-CNN.

\textbf{Vision task}: While the primary focus of this paper is on short text applications, DeepXML\xspace was also applied to a vision task to demonstrate its ease-of-use and generality. Table~9 in the \href{http://manikvarma.org/pubs/dahiya21.pdf}{\color{blue}{supplementary material}}\xspace presents results for classifying a retail item's image into its product categories from the AmazonCat-13K dataset available on the Repository. DeepXML\xspace made only a minor modification in Astec\xspace by replacing its feature architecture by ResNet18~\citep{He16} while leaving everything else unchanged. Pre-trained ResNet18 features were also used to train Slice, MACH, Parabel and DiSMEC. Nevertheless, DeepXML\xspace was 4-29\% more accurate than these methods thereby demonstrating its flexibility.

\section{Conclusions}
This paper developed the modular DeepXML\xspace framework that was used to: (a) derive the Astec\xspace algorithm for low-latency short text applications; (b) analyse and improve leading deep extreme classifiers; and (c) provide a convenient and flexible tool for practitioners to plug in components of their choice with minimal effort for tackling diverse applications. Furthermore, Astec\xspace was demonstrated to be significantly more accurate and scalable than leading extreme classifiers for short text documents and could lead to significant gains in various online metrics for multiple applications on Bing.

\begin{acks}
The authors thank Purushottam Kar, Aditya Kusupati, Harsha Vardhan Simhadri, Yash Garg, and Risi Thonangi for helpful feedback.
\end{acks}

\section{Supplementary Material}
\input{Tables/dataset_stats.tex}
 \input{Tables/results_framework}

\subsection{Hyper-parameters}\label{sec:supp:hyper-parameters}
Astec\xspace's hyper-parameters include $\alpha$ for combining the classifier \& shortlist scores, $\hat{L}$ which was the number of labels in the surrogate task which was set to $2^{16}$ across all datasets and the label shortlist size which was set to $500$ in all cases. Note that the different values of $\alpha$ leads to a trade-off in vanilla and propensity scored metrics. The embeddings in $\mathcal{Z}$, residual matrices and classifiers were initialized with FastText~\cite{Joulin17}, the identity matrix and Xavier's method respectively. Astec\xspace was trained using the Adam optimizer with spectral norm constraints~\cite{Miyato18} and its hyper-parameters included the learning rate, the batch size and the number of epochs. Most of these were set to default values across datasets and the most expensive hyper-parameter to tune was the learning rate on the surrogate task.

 Experiments were performed on a P40 GPU card with CUDA 11, and Pytorch 1.8 unless stated otherwise. Dropout with probability 0.5 was used for all datasets. HNSW~\citep{MalkovY16} parameters $M$, $efC$ and $efS$ where set to $100, 300, 300$ and $50, 50, 100$ for $\text{ANNS}^x$ and $\text{ANNS}^{\mu}$ respectively. The surrogate learning task was trained with $|\hat{L}| = 2^{16}$ with learning rate chosen from $\{0.003, 0.005, 0.02 \}$. Increasing $|\hat{L}|$ beyond $2^{16}$ lead to only marginal gains in accuracy but at the cost of increase in training time. The model parameters for the extreme task were learnt with a learning rate chosen from $\{0.002, 0.0005\}$. It should be noted that no hyper-parameter tuning was done for proprietary datasets where Astec lead to significant gains in offline as well as online metrics.

\input{Tables/supp_hyperparameters}
\input{Tables/ablation_framework}

\subsection{Clustering labels}\label{sup:clustering}
Astec\xspace clustered the labels using label centroid representation $\hat{\vmu}_l = \frac{\sum_{i=1}^{N}y_{il}\vx_i}{\norm{\sum_{i=1}^{N}y_{il}\vx_i}_2}$ as the label meta-data was unavailable. 2-means++ algorithm was deployed to solve the following optimization problem, which recursively clusters the labels into two balanced-partitions to finally end up with $\hat{L}$ clusters.

\begin{align*}\label{equ:clustering}
\argmin_{ {\vmu}_{\pm},\valpha \in \{-1, 1\}^{L}} & \sum_{l=1}^{L} \cC_{ll} \left(\frac{1+\valpha_l}{2} \hat{\vu}_l \vmu_{+} + \frac{1-\valpha_l}{2} \hat{\vu}_l \vmu_{-} \right)\\
&+ \sum_{l=1}^{L}\sum_{p=1}^{L} \cC_{lp} \left(\frac{1+\valpha_l}{2} \hat{\vu}_p \vmu_{+} + \frac{1-\valpha_l}{2} \hat{\vmu}_p \vmu_{-}\right) \\
\mbox{s. t.} ~~&~~ \norm{\vmu_{\pm}}_2 = 1,~~-1 \le \sum_{l=1}^{L} \valpha_l \le 1
\end{align*}

where it has been assumed without loss of generality that $L$ labels need to be partitioned at each step, $\alpha_l=\pm1$ means that label $l$ is assigned to cluster with mean $\vmu_{\pm}$ and $\cC$ is the $\ell1$ normalized label correlation matrix~($\vY^{\top} \vY$). In practise, the label correlation matrix was estimated by performing a random walk over a graph with labels as nodes and data-points as edges.

\subsection{Additional surrogate tasks}
\label{subsec:divalgo}
A simple but effective way to train the parameters of chosen feature architecture, {\it i.e.,} $\mathcal{Z}$ could be to select $\hat{L}$ labels based on label frequency in the training set. The hyper-parameter $\hat{L}$ was chosen to balance two constraints. First, $\hat{L}$ should be large enough so that almost all the token embeddings could be learnt in this first phase of training. At the same time, the $|\hat{L}|$ should be small enough so that (\ref{equ:surr_opt}) could be optimized efficiently on a single GPU as a non-extreme problem and without resorting to ANNS shortlisting. It was empirically observed that setting $0.05L \leq \hat{L} \leq 0.2L$ with lower values being preferred for larger problems resulted in accurate intermediate representations. It should be reiterated that the balanced clustering was found to be more accurate and scalable than the alternatives including label selection based techniques.

\subsection{ANNS search and multiple representatives}\label{sup:multiple_rep}
It is worth pointing out that some of the most frequently occurring head labels could be usefully represented by {\em multiple} $k$-means cluster centres while constructing the $\mbox{ANNS}$ small world graph. This allowed for the accurate shortlisting of multi-modal head labels which could not be shortlisted well based on a single label centroid representation, {\it i.e.,} $\vmu_l^0$. For instance, representing the top 4 head labels on the WikiTitles-500K dataset by 300 $k$-means cluster centres rather than just the label mean improved recall@300 by $5\%$ without any noticeable increase in the training or prediction time.

\subsection{Evaluation metrics}
\label{sup:eval}
Performance has been evaluated using propensity scored precision@$k$ and nDCG@$k$, which are unbiased and more suitable metric in the extreme multi-labels setting~\citep{Jain16, Babbar19, Prabhu18, Prabhu18b}. The propensity model and values available on The Extreme Classification Repository~\citep{XMLRepo} were used. Performance has also been evaluated using vanilla precision@$k$ and nDCG@$k$ (with $k$ = 1, 3 and 5) for extreme classification.

For a predicted score vector $\hat{\mathbf{y}} \in R^L$ and ground truth vector $\mathbf{y} \in \{0, 1\}^L$:
		$$P@k = \frac{1}{k} {\sum_{l \in rank_k(\hat{\mathbf{y}})}} y_l$$
		$$PSP@k = \frac{1}{k} {\sum_{l \in rank_k(\hat{\mathbf{y}})}} \frac{y_l}{p_l}$$
		$$DCG@k = \frac{1}{k} {\sum_{l \in rank_k(\hat{\mathbf{y}})}} \frac{y_l}{\log(l+1)}$$
		$$PSDCG@k = \frac{1}{k} {\sum_{l \in rank_k(\hat{\mathbf{y}})}} \frac{y_l}{p_l \log(l+1)}$$
		$$nDCG@k = \frac{DCG@k}{\sum_{l=1}^{\min(k, ||\mathbf{y}||_0)} \frac{1}{\log(l +1) }}$$
		$$PSnDCG@k = \frac{PSDCG@k}{\sum_{l=1}^{k} \frac{1}{\log l +1 }}$$
Here, $p_l$ is propensity score of the label $l$ proposed in~\citep{Jain16}.

\subsection{Theorem proofs}
\begin{proof} \label{sup:proof}
\textbf{(Bound on $||\hat{\vx}_i - \vv_i||_2$)}. For notational convenience, we use $\hat{\vx} := \hat{\vx}_i$ and $\vv := \vv_i$;
\begin{align*}
    \hat{\vx}_i - \vv_i &= ReLU(\vR \vv) \\
    ||\hat{\vx}_i - \vv_i||_2 &=  ||ReLU(\vR \vv)||_2
    \shortintertext{\center Using, $||ReLU(\vu)||_2 \le ||\vu||_2$}
    &\le ||\vR \vv||_2
    \shortintertext{\center Using, $||\vR||_{op} \le \lambda$}
    &\le \lambda||\vv||_2
\end{align*}

\textbf{(Bound on $\hat{\vx}_i$)}.
\begin{align*}
    ||\hat{\vx}_i||_2 &\le ||\vv||_2 + ||\hat{\vx}_i - \vv_i||_2
    \shortintertext{\center Using, $||\hat{\vx}_i - \vv_i||_2 \le \lambda||\vv||_2$}
    &= ||\vv||_2 + \lambda||\vv||_2\\
    &= (1+\lambda)||\vv||_2
\end{align*}

In order to prove bound on cosine similarity, we first prove bound on $||\vmu_l - \vmu^0_l||_2$ and $||\vmu_l||_2$. For notational convenience, we use $\cP := \cP_l$, $\vmu^0 := \vmu^0_l$, and $\vmu := \vmu_l$, \\
\textbf{(Bound on $\vmu_l - \vmu^0_l$)}.
\begin{align*}
    \vmu_l - \vmu^0_l &= \frac{1}{|\cP|} \sum_{i \in \cP}{\hat{\vx}_i - \vv_i} \\
    ||\vmu_l - \vmu^0_l||_2 &= ||\frac{1}{|\cP|} \sum_{i \in \cP}{\hat{\vx}_i - \vv_i}||_2 \\
    &= ||\frac{1}{|\cP|} \sum_{i \in \cP}{ReLU(\vR \vv_i})||_2 \\
    &\le \frac{1}{|\cP|}\sum_{i \in \cP}|| ReLU(\vR \vv_i)||_2
    \shortintertext{\center Using $||ReLU(\vu)||_2 \le ||\vu||_2$}
    &\le \frac{1}{|\cP|}\sum_{i \in \cP}|| \vR \vv_i||_2
    \shortintertext{\center Using, $||\vR||_{op} \le \lambda$}
    &\le  \frac{\lambda}{|\cP|}\sum_{i \in \cP}{|| \vv_i}||_2\\
    \shortintertext{\center Let, $\mathrm{V} := [||\vv_1||_2,..,||\vv_{|\cP|}||_2]$,}
    &= \frac{\lambda}{|\cP|}||\mathrm{V}||_1
    \shortintertext{\center As, $||\mathrm{V}||_1 \le \sqrt{|\cP|}||\mathrm{V}||_2$}
    &\le \frac{\lambda}{\sqrt{|\cP|}}||\mathrm{V}||_2 \\
    &\le \frac{\lambda}{\sqrt{|\cP|}}[\sum_{i \in \cP}||\vv_i||_2^2]^{\frac{1}{2}} \\
    &= \frac{\lambda}{\sqrt{|\cP|}}[\sum_{i \in \cP}{\vv_i}^T \vv_i]^{\frac{1}{2}}
    \shortintertext{\center As, $\sum_{i \in \cP}\vv_i^T \vv_i \le [\sum_{i \in \cP}\vv_i]^T[\sum_{i \in \cP} \vv_i]$}
    &\le \frac{\lambda}{\sqrt{|\cP|}}\{[\sum_{i \in \cP}\vz_i]^T[\sum_{i \in \cP}\vz_i]\}^{\frac{1}{2}}\\
    &\le \frac{\lambda|\cP|}{\sqrt{|\cP|}}[\vmu^{0{\top}} \vmu^0]^{\frac{1}{2}} \\
    &= \lambda\sqrt{|\cP|}||\vmu^0||_2
\end{align*}

\textbf{(Bound on $\vmu_l$)}.
\begin{align*}
\vmu &= \frac{1}{|\cP|} \sum_{i \in \cP}{\hat{\vx}_i} \\
||\vmu||_2 &= ||\frac{1}{|\cP|} \sum_{i \in \vP} \hat{\vx}_i||_2 \\
&= ||\frac{1}{|\cP|} \sum_{i \in \cP}{\vv_i + \frac{1}{|\cP|} \sum_{i \in \cP}{\hat{\vx}_i - \vv_i}||_2} \\
&\le ||\frac{1}{|\cP|} \sum_{i \in \cP}{\vv_i}||_2 + ||\frac{1}{|\cP|} \sum_{i \in \cP}{\hat{\vx}_i - \vv_i}||_2 \\
&= ||\vmu^0||_2 + ||\vmu_l - \vmu^0_l||_2 \\
&\le (1+\lambda\sqrt{|\cP|}) ||\vmu^0||_2
\end{align*}

\textbf{(Lower bound on $C(\hat{\vx}, \vmu_l)$)}.
\begin{align*}
     \frac{C(\hat{\vx}, \vmu)}{C(\vv, \vmu^0)} &= {\frac{\hat{\vx}^{\top}\vmu \cdot ||\vv||_2 ||\vmu^0||_2}{\vv^{\top}\vmu^0 \cdot ||\hat{\vx}||_2 ||\vmu||_2}}
     \shortintertext{\center Using, $||\hat{\vx}||_2 \le (1+\lambda)||\vv||_2$ and, $||\vmu||_2 \le (1+\lambda \sqrt{|\cP|})|| \vmu^0||_2$}
     &\ge \frac{\hat{\vx}^{\top}\vmu}{(1+\lambda)(1+\lambda \sqrt{|\cP|}) \cdot \vv^{\top} \vmu^0}\\
     &\ge \frac{(\hat{\vx})^{\top}(\vmu)}{(1+\lambda \sqrt{|\cP|})^2 \cdot \vv^{\top}\vmu^0}
     \shortintertext{\center Using, $\hat{\vx}^{\top}\vmu \ge \vv^{\top} \vmu^0$, as, $\hat{\vx} - \vv \ge 0$, and $\vmu - \vmu^0 \ge 0$}
     &\ge \frac{1}{(1+\lambda \sqrt{|\cP|})^2}
\end{align*}

\textbf{(Upper bound on $C(\hat{\vx}_i, \vmu_l)$)}
\begin{align*}
     C(\hat{\vx}, \vmu_l) &=
     \frac{\hat{\vx}^{\top} \vmu }{||\hat{\vx}||_2 ||\vmu||_2}
     \shortintertext{\center Using, $||\hat{\vx}||_2 \ge ||\vv||_2$ and, $||\vmu||_2 \ge ||\vmu^0||_2$}
     &\le \frac{\hat{\vx}^{\top} \vmu}{||\vv||_2 ||\vmu^0||_2} \\
     &= \frac{\vv^{\top} \vmu^0 + \vv^{\top} (\vmu - \vmu^0) + (\hat{\vx} - \vv)^{\top}\vmu^0 + (\hat{\vx} - \vv)^{\top} (\vmu - \vmu^0)}{||\vv||_2 ||\vmu^0||_2} \\
     &= C(\vv, \vmu^0) \\
     &+ \frac{  ||\vv^{\top} \vmu||_2 + ||\vv^{\top}\vmu^0||_2 + ||\hat{\vx}^{\top} \vmu||_2 - ||\vv^{\top} \vmu||}{||\vv||_2 ||\vmu^0||_2} \\
     &\le C(\vv, \vmu^0) \\
     &+ \frac{ ||\vv||_2 ||\vmu||_2 + ||\vv||_2 ||\vmu^0||_2 + ||\hat{\vx}||_2 ||\vmu||_2 - ||\vv||_2 ||\vmu||}{||\vv||_2 ||\vmu^0||_2} \\
      \shortintertext{\center Using, $||\hat{\vx}||_2 \le (1+\lambda)||\vv||_2$ and $||\vmu||_2 \le (1+\lambda \sqrt{|\cP|})||\vmu^0||_2$}
      &\le C(\vv, \vmu^0) \\
      &+ \frac{((1+\lambda)(1+\lambda\sqrt{|\cP|})-1)||\vv||_2||\vmu^0||_2}{||\vv||_2 ||\vmu^0||_2} \\
      &\le C(\vv, \vmu^0) + (1+\lambda \sqrt{|\cP|})^2 -1  \qedhere
\end{align*}
\end{proof}

\input{Figs/negative_sampling_ablation}
\input{Tables/qualitative_analysis_wiki}
\input{Tables/qualitative_analysis_pads}
\input{Tables/results_extreme_vision.tex}
\input{Tables/results_full_supp}
\input{Tables/results_titles_supp}

\clearpage

Machine Learning (cs.LG); Information Retrieval (cs.IR)


Using VAEs and Normalizing Flows for One-shot Text-To-Speech Synthesis of Expressive Speech

Abstract: We propose a Text-to-Speech method to create an unseen expressive style using one utterance of expressive speech of around one second. Specifically, we enhance the disentanglement capabilities of a state-of-the-art sequence-to-sequence based system with a Variational AutoEncoder (VAE) and a Householder Flow. The proposed system provides a 22% KL-divergence reduction while jointly improving perceptual metrics over state-of-the-art. At synthesis time we use one example of expressive style as a reference input to the encoder for generating any text in the desired style. Perceptual MUSHRA evaluations show that we can create a voice with a 9% relative naturalness improvement over standard Neural Text-to-Speech, while also improving the perceived emotional intensity (59 compared to the 55 of neutral speech).

\section{Introduction}
\label{sec:intro}
Neural Text-to-Speech models, such as Tacotron and WaveNet \cite{oord2016wavenet, shen2018natural}, are able to produce high quality speech with naturalness close to that of real humans. However, these models require several hours of data to produce high-quality voices.

This work focuses on creating high-quality expressive speech using only one utterance of a target style without requiring retraining. We show that we are able to create high-quality expressive voices at low cost and improve the naturalness of our baseline \textit{state-of-the-art} neutral text-to-speech system. Our approach relies on a architecture combining Householder Flows and Variational AutoEncoders (VAE) in a novel way that provides better KL-Divergence (KLD) between an isotropic Gaussian prior and our VAE latent variables, alongside better reconstruction than existing methods. Reductions in KLD provide lower correlation between our latent variables. Consequently, our approach provides improved manipulation of the latents  \cite{hsu2018hierarchical}, substantial improvements in the disentanglement of the latent space, and raises perceptual metrics of synthesized expressive speech whilst using only one reference utterance.

To the authors knowledge, this work is the first to tackle the problem of one-shot synthesis of an unseen expressive style. Although there is existing work in the related fields of prosody generation \cite{skerry2018towards} and prosody transfer \cite{akuzawa2018expressive}, these approaches rely on either modelling expressive speech or transferring a seen expressive style to an unseen speaker, and are trained on a large corpora of expressive speech. In contrast, we use only one utterance to create an unseen style. Our problem is also distinct from that of speaker-adaptation \cite{nachmani2018fitting, jia2018transfer}, as we wish to improve expressiveness by including emotion for existing speakers, not learn entirely new speakers. Existing work on emotional speech focuses on explicit emotional speech models \cite{lorenzo2018investigating}, model adaptation \cite{wu2018rapid} or voice conversion \cite{gao2018nonparallel,shankar2019multi}, which again require significantly more data than the single utterance required by our proposed system.

\gdef\@thefnmark{}\@footnotetext{Corresponding author email: agvatsal@amazon.com.}

\vspace{-2mm}
\section{Proposed Approach}
\label{sec:approach}
\begin{figure}[t]
	\centering
	\includegraphics[width=1\columnwidth]{architecture4.png}
	\caption{Overview of the system's architecture: a Variational seq-to-seq auto-encoder conditioned on the phonetic sequence. The Universal Neural Vocoder translates the generated spectrogram into the final speech signal.}
	\label{fig:Architecture}
\end{figure}
For our proposed architecture, we extend a \textit{state-of-art} sequence-to-sequence NTTS \cite{shen2018natural,prateek2019other,latorre2019effect} system. This base seq2seq model uses an attention module to map the concatenated outputs of an  RNN-based Phoneme Encoder and a mel-spectrogram conditioned Reference Encoder to mel-spectrograms via an autoregressive  decoder. The reference encoder is trained as the aggregate posterior of a VAE \cite{kingma2013auto}, and is conditioned on mel-spectrograms to predict the mean and variance of a 64-dimensional diagonal Gaussian Distribution (with similar architecture to \cite{skerry2018towards}). Thus, the reference encoder and the decoder form a text-conditioned VAE, with the Phoneme Encoder providing text-conditioning to our network. We consider this as our baseline system, referred to as Vanilla VAE

Our proposed system extends this base model by adding a Householder Normalizing Flow \cite{tomczak2016improving} (as presented in Figure~\ref{fig:Architecture}) to the text-conditioned VAE. Householder Flows consist of multiple steps of easily invertible affine transformations that transform samples from a diagonal Gaussian to samples from a full co-variance Gaussian. We hypothesise that Householder Normalizing Flows allow samples from the reference encoder to better match the diagonal normal Gaussian prior while allowing the samples after the flow to become more co-related. We believe that this allows for both improved disentanglement and better reconstruction by the decoder. To investigate this claim, we tried a variety of architectures differing only in how the Householder vectors used to define each step of the householder flow are predicted. We now summarize our different architectures (see also Figure \ref{fig:vae_plus_flows_architectures}).
\begin{itemize}
    \item Architecture-1: We follow the original implementation of Householder Flows \cite{tomczak2016improving} - predicting the first householder vector as an additional output of the reference encoder, with all other vectors predicted via subsequent affine transformations.
    \vspace{-2.5mm}
    \item Architecture-2: We predict all the Householder vectors as separate outputs of the reference encoder.
    \vspace{-2.5mm}
    \item Architecture-3: We propose a novel architecture where Householder vectors are globally shared across all utterances. This architecture is distinct from all existing implementations of Householder Flows, as the vectors are not conditioned on the input to the VAE encoder in any way (e.g. Architecture-1).
\end{itemize}
\begin{figure}[t]
	\centering
	\includegraphics[width=1\columnwidth]{vae_plus_flos_architectures.png}
	\caption{Diagram of proposed architectures: Architecture-1 (a), Architecture-2 (b), and Architecture-3 (c). $h_i$ refers to the $i^{th}$ householder vector.}
\label{fig:vae_plus_flows_architectures}
\end{figure}
Each architecture is evaluated using 2, 4, 8 and 16 householder vectors as hyperparameters. Across all our architectures, samples from the output of the flow are concatenated with the output of the phoneme encoder instead of samples from the output distribution of the reference encoder.

During training, we use a large training corpus (Section \ref{sec:experiments}) and feed the target spectrogram into the reference encoder. At synthesis time, we instead use a spectrogram from the 1-shot data-set (also described in Section \ref{sec:experiments}) as input to the reference encoder and feed phonemes of the target text into the phoneme encoder. To complete the TTS pipeline, we use the vocoder of  \cite{lorenzo2019towards} to convert the mel-spectrograms to audio waveforms.

\section{CORPORA}
\label{sec:experiments}
\vspace{-1mm}
We train our models on a training data-set combining two corpora: a high-quality proprietary multi-speaker corpus, and a subset of the VCTK \cite{veaux2017cstr} corpus. Both data-sets contained natural speech waveforms from English speakers speaking in non emotional styles. Our internal corpus contains 181 hours of speech from 13 speakers (each contributing between 5 and 35 hours), which we combine with 21 English VCTK speakers (each with 23 minutes of recordings).

Alongside this training corpus, we also have access to a small corpus of expressive recordings for one of our internal speakers. This speaker exhibits the `excited' emotion across three intensity levels: low, medium and high. We randomly picked one utterance for each intensity level to condition the reference encoder in our 1-shot experiments. Furthermore, we selected $50$ text prompts to be synthesized for perceptual evaluations, ensuring that each prompt corresponds to an appropriate emotional and intensity level given its content.

\vspace{-3mm}
\section{RESULTS}
\vspace{-1mm}
\label{sec:experiments}
We evaluated each system using both objective and perceptual metrics. Our objective metrics are the KL-divergence of our reference encoder and the final teacher-forced L2 loss achieved by our the decoder. To choose systems for perceptual evaluations, we compare our system against other \textit{state-of-the-art} methods with HouseHolder Flows (across multiple hyper-parameter configurations),  and choose the best model with and best model without flows. Our perceptual evaluations target speech naturalness, emotional intensity and signal quality.
\begin{figure}[t]
	\centering
	\includegraphics[width=\columnwidth]{kl_and_recon.png}

	\caption{KL-Divergence (blue bars) and Reconstruction Objective (orange bars) Metrics. We plot the average loss for the final epoch when training is stopped for each architecture and hyper-parameter presented in the Proposed Approach section.}
	\label{fig:obj}
\end{figure}

\subsection{Objective Evaluation}
\label{ssec:obj}
We compare the Vanilla VAE system summarized in the Section 2, with the three other architectures predicting householder vectors in different ways. Our proposed architecture, Architecture-3, achieves the lowest combination of KL-divergence and Reconstruction Loss across all systems (Figure \ref{fig:obj}). Furthermore, it is the only architecture which provides lower KL-divergence, reducing it by $22$\%, or improved disentanglement, while also providing better reconstruction when comparing against Vanilla VAE. From informal listening, this model indeed provided better quality in limited scenarios as compared to the other models, with flows, that we tried. We believe this is because the flow in our proposed architecture is now able to learn transformations of factors that are relevant for speech in general rather than having to learn how to transform factors that could change with different inputs. We will refer to this architecture with the best hyper-parameter (16 householder vectors) as VAE+FLOW for the rest of this paper.

In the next section, we verify that the superior disentanglement achieved by our model (without a loss in reconstruction loss), allows for convincing 1-shot generalization of expressive speech.

\vspace{-3mm}
\subsection{Perceptual Evaluations}
\label{ssec:perc}
To evaluate our model with respect to naturalness, emotional strength, and signal quality, we ran three MUltiple Stimuli with Hidden Reference and Anchor (MUSHRA) tests \cite{recommendation2001bs} across 40 listeners from Amazon Mechanical Turk.

For naturalness, we asked our listeners ``Please rate the naturalness of the systems. Make sure that at least 1 system per screen is always rated 100." and provided a scale from 0 (Very poor) to 100 (Completely natural). To measure emotional strength, we asked ``Please rate the strength of the emotion of the systems. Consider the neutral sounding sample (there is one provided for each utterance) as 0 and an extremely emotional sounding sample as 100." and provided a scale from 0 (Neutral) to 100 (Very Emotional). Finally, to measure signal quality, we asked  ``Please rate the systems in terms of audio quality. Try to ignore the content of the speech and the expressivity and instead focus on the quality of the audio signal (e.g glitches, clicks, noise...)." and provided a scale from 0 (Very poor) to 100 (Very good).

Each MUSHRA consisted of 4 systems: recordings of Emotional Speech (hidden higher anchor), Neutral TTS (hidden lower anchor), Vanilla VAE and VAE+FLOW. We show each listener 50 utterances with excited emotions with prompts as described in the Corpus section. To measure statistical significance, we ran t-tests with Bonferroni-Holm corrections at a $95$\% significance level. Our null-hypothesis is that the mean MUSHRA scores of the listening population for the two systems is equal.

\subsubsection{Naturalness}
\label{sssec:natural}
MUSHRA evaluations of naturalness are presented in Figure~\ref{aggregrated2}. When aggregated across all emotional intensities, there is a statistically significant difference between all pairs of systems. Both 1-shot based expressive systems (VAE+FLOW and VAE) obtain higher naturalness ($66.3$ and $67.8$ respectively) than the neutral system ($62.2$). When we break down the results per intensity conveyed in the reference utterance, this improvement in naturalness of both the 1-shot based expressive systems with respect to the neutral system is maintained for low and medium intensities. However, the difference between the two 1-shot expressive systems is not statistically significant. In contrast, when using high intensity utterances as the reference, there is a statistically significant degradation respect to neutral ($66.6$) for both expressive systems: partially reduced by VAE+FLOW ($63.9$) as compared to the vanilla VAE ($60.3$).
\begin{figure}
	\centering
	\includegraphics[width=\columnwidth]{n_e2.png}
	\caption{Perceptual Results of naturalness for low, medium, high and aggregated emotional intensities. We plot MUSHRA responses of listeners as box-plots for each intensity. The number in white and location of the green triangle represents the mean of the listeners response for that system.}
	\label{aggregrated2}

\end{figure}

\subsubsection{Emotional Strength}
\label{sssec:emo-inten}
For emotional strength (Figure \ref{aggregrated}), we found significant differences between all pairs of systems, in the aggregated results as well as results split by intensity. As expected, listeners scored the emotional strength of the neutral system consistently (and lower than rest of the systems) across all the evaluations and both the 1-shot expressive systems were scored with a higher mean score than neutral.
\begin{figure}
	\centering
	\includegraphics[width=\columnwidth]{es_e2.png}
	\caption{Perceptual Results for emotional strength across low, medium, high and aggregated emotional intensities. Legend and metrics are reported as in Figure~\ref{aggregrated2}.}
	\label{aggregrated}
\end{figure}

\subsubsection{Signal Quality}
Figure \ref{aggregrated3} shows the perceptual evaluations of signal quality. When aggregated across all emotional intensities, we see a significant difference between all pairs of system. Furthermore, for low and medium intensities the difference between flow and no-flow is not significant, whereas differences between all other system pairs are significant. In the high intensity case, the difference between all pairs of systems is significant.
\begin{figure}
	\centering
	\includegraphics[width=\columnwidth]{sq_e2.png}
	\caption{Perceptual Results for signal quality for low, medium, high and aggregated emotional intensities. Legend and metrics are reported as in Figure~\ref{aggregrated2}.}
	\label{aggregrated3}
\end{figure}

\vspace{-4mm}
\section{DISCUSSION}
\vspace{-1mm}
Overall, both Vanilla VAE and VAE+FLOW provide a method to perform one-shot adaptation of a new expressive style. Vanilla VAE provided a statistically significantly improvement in emotional strength ($23\%$ relative), and naturalness ($2\%$ relative) over Neutral TTS despite listeners detecting a degradation in signal quality ($20\%$ relative). Our proposed VAE+FLOW architecture, leveraging a simple computationally inexpensive flow, is able to statistically significantly improve further upon Vanilla VAE in emotional strength ($2.5\%$ relative) naturalness ($2.2\%$ relative), and signal quality ($1.7\%$ relative).

After breaking down the perceptual results by emotional intensity, we found that the relative improvement achieved by VAE+FLOW over Vanilla VAE increases with emotional intensity. In this scenario, a higher intensity means a bigger acoustic divergence from the neutral speech. We, therefore, hypothesize that our proposed VAE+FLOW architecture performs better at the extreme cases of high intensity due to improved disentanglement, leading to better generalization capabilities.

This hypothesis is supported by our perceptual tests on emotional strength. All systems suffer a drop in naturalness when synthesizing high excitement (due to the degradation in signal quality caused by trying to synthesize speech very different to the training corpus). However, the improvement in generalization by the VAE+FLOW system provides a smaller naturalness degradation, compared to Neutral TTS, than the Vanilla VAE system in this case ($4\%$ and $9.5\%$ respectively). Finally, we note that while we're consistently increasing emotional strength in the low and medium case, these results are not translating proportionately to listeners perception of naturalness despite natural recordings of those emotions being rated highly. We hypothesize this is due to the limited signal quality that 1-shot systems can currently achieve.
\vspace{-2.5mm}
\section{CONCLUSIONS}
\vspace{-1.5mm}
We proposed a novel approach using householder flows within a text-conditioned VAE that overcomes the common trade-off between disentanglement and reconstruction in a state-of-the-art VAE implementation. The proposed flow-based VAE system significantly reduced KL-divergence ($22\%$ relative reduction) whilst providing a $2\%$ relative improvement in reconstruction. Furthermore, we showed that pre-existing methods can perform 1-shot adaption to a new expressive style with only one utterance. Our VAE+FLOW model improved naturalness and emotional strength over an existing state-of-the-art system particularly when generalizing across more extreme prosody differences (with respect to the neutral training data) using only one utterance.

We have shown that signal quality is the bottleneck preventing joint improvement of naturalness and emotional strength. Future investigations are necessary to create voices with limited resources across a wider variety of expressive styles, seeking to further improve naturalness without harming expressiveness.

\vspace{-1.5mm}
\section{Acknowledgment}
\vspace{-2.5mm}
We would like to thank Pablo Garcia Moreno, Jasha Droppo, Henry Moss, Andrew Breen and Thomas Drugman for insightful discussions relevant to this work.

\vfill\pagebreak
\bibliographystyle{IEEEbib}

Machine Learning (cs.LG); Computation and Language (cs.CL); Sound (cs.SD); Audio and Speech Processing (eess.AS); Machine Learning (stat.ML)