# Multi-stage Retrieve and Re-rank Model for Automatic Medical Coding
Recommendation

Xindi Wang1,2, Robert E. Mercer1, Frank Rudzicz2,3,4
1 Department of Computer Science, University of Western Ontario, Canada
2 Vector Institute for Artificial Intelligence, Canada
3 Faculty of Computer Science, Dalhousie University, Canada
4 Department of Computer Science, University of Toronto, Canada
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

The International Classification of Diseases (ICD) serves as a definitive
medical classification system encompassing a wide range of diseases and
conditions. The primary objective of ICD indexing is to allocate a subset of
ICD codes to a medical record, which facilitates standardized documentation
and management of various health conditions. Most existing approaches have
suffered from selecting the proper label subsets from an extremely large ICD
collection with a heavy long-tailed label distribution. In this paper, we
leverage a multi-stage “retrieve and re-rank” framework as a novel solution to
ICD indexing, via a hybrid discrete retrieval method, and re-rank retrieved
candidates with contrastive learning that allows the model to make more
accurate predictions from a simplified label space. The retrieval model is a
hybrid of auxiliary knowledge of the electronic health records (EHR) and a
discrete retrieval method (BM25), which efficiently collects high-quality
candidates. In the last stage, we propose a label co-occurrence guided
contrastive re-ranking model, which re-ranks the candidate labels by pulling
together the clinical notes with positive ICD codes. Experimental results show
the proposed method achieves state-of-the-art performance on a number of
measures on the MIMIC-III benchmark.

Multi-stage Retrieve and Re-rank Model for Automatic Medical Coding
Recommendation

Xindi Wang1,2, Robert E. Mercer1, Frank Rudzicz2,3,4 1 Department of Computer
Science, University of Western Ontario, Canada 2 Vector Institute for
Artificial Intelligence, Canada 3 Faculty of Computer Science, Dalhousie
University, Canada 4 Department of Computer Science, University of Toronto,
Canada<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>

##  1 Introduction

Electronic health
records111https://www.cms.gov/Medicare/E-Health/EHealthRecords (EHRs) contain
a comprehensive repository of essential administrative and clinical data
pertinent to a person’s care within a specific healthcare provider setting. In
order to conduct meaningful statistical analysis, these EHR data are annotated
with structured codes in a classification system known as medical codes. The
International Classification of
Diseases222https://www.who.int/standards/classifications/classification-of-
diseases (ICD) is one of the most widely-used coding systems, and it provides
a taxonomy of classes, each uniquely identified by a code assigned to an
episode of patient care.

The task of medical coding associates ICD codes with EHR documents. The status
quo of assigning medical codes is a manual process, which is labour-intensive,
time-consuming, and error-prone Xie and Xing (2018). To reduce coding errors
and cost, the demand for automated medical coding has become imperative.
Previous deep learning approaches regarded medical coding as an extreme multi-
label text classification problem Shi et al. (2017); Mullenbach et al. (2018);
Baumel et al. (2018); Xie et al. (2019); Yuan et al. (2022), where an encoder
is typically employed to learn the representations of the clinical notes and a
label-specific binary classifier is subsequently constructed on top of the
encoder for label predictions. However, some remaining difficulties have still
posed immense challenges. First, clinical documents are lengthy (containing on
average 1596 words in the MIMIC-III dataset) and noisy (including terse
abbreviations, symbols, and misspellings). Second, the label set is extremely
large and complex; for instance, in the $10^{th}$ ICD edition, there are over
130,000 codes333https://www.cdc.gov/nchs/icd/icd10cm_pcs.htm. Third, the
distribution of ICD codes is extremely long-tailed; while some ICD codes occur
frequently, many others seldom appear, if at all, because of the rarity of the
diseases. For instance, among the 942 unique 3-digit ICD codes in the MIMIC-
III dataset Johnson et al. (2016), the ten most common codes account for 26%
of all code occurrences and the 437 least common codes account for only 1% of
occurrences Bai and Vucetic (2019).

Figure 1: An example of a medical record from the MIMIC-III dataset which
includes the discharge summary, assigned ICD codes and auxiliary knowledge. We
colour each code and its corresponding mentions in the discharge summary and
auxiliary knowledge. We use the auxiliary knowledge of the notes to retrieve
the candidate subset of the label space.

To address the aforementioned challenges, we propose a novel multi-stage
retrieve and re-rank framework, where the goal is to first generate a curated
ICD list and then provide suggested ICD codes for a given medical record. In
contrast to prior approaches, for instance, CAML Mullenbach et al. (2018),
MultiResCNN Li and Yu (2020) and KEPTLongformer Yang et al. (2022), that
primarily consider ICD indexing as a multi-label text classification task, we
introduce a new perspective that conceptualizes the task as a recommendation
problem. More precisely, we first conduct a two-stage retrieval process
leveraging auxiliary knowledge and BM25 to obtain a small subset of candidate
ICD codes from the large number of labels to alleviate issues caused by the
label set and imbalanced label distribution. EHR auxiliary knowledge holds
significant potential, but it has often been underutilized in prior studies.
In addition to clinical texts, our focus centers on two code terminologies:
Diagnosis-Related Group codes444https://www.cms.gov/Medicare/Medicare-Fee-for-
Service-Payment/AcuteInpatientPPS/MS-DRG-Classifications-and-Software (DRG)
and Current Procedural Terminology codes555https://www.ama-
assn.org/amaone/cpt-current-procedural-terminology (CPT), as well as patient
prescribed medications. These external sources can serve as robust indicators
for predicting ICD codes. For instance, within a drug prescription, the
presence of a medication like “Namenda” can strongly imply a likelihood of
Alzheimer’s disease, as depicted in Figure 1. Subsequently, we design a re-
ranking model via co-occurrence guided contrastive learning to refine the
candidate set, which can deal with lengthy clinical notes and generate
semantically meaningful representations via the pre-trained language model and
leverage code co-occurrence to generate co-occurrence-aware label
representations. The co-occurrence of codes in clinical texts yields valuable
insights into the interconnections among different diseases or conditions. As
illustrated in Figure 1, the code for “Dementia in conditions classified
elsewhere without behavioral disturbance” (294.10) can be easily found in the
text; however, inferring the code “Alzheimer’s disease” (331.0) presents a
more intricate challenge with less explicit clues. Fortunately, a robust
association exists between these two diseases, with “Alzheimer’s disease”
serving as a prevalent cause of “dementia”. This linkage can be effectively
captured as these two diseases frequently co-occur within the clinical notes.
This empowers us to gain a deeper understanding of the contexts, which could
mitigate the limitation of long-tailed label distributions as rare labels
might be suggested based on these relationships. We train the re-ranking model
via contrastive learning as it has strong discriminative power that can
extract features uniquely associated with each class, which empowers the model
to make more accurate recommendations.

To summarize, the major contributions of this paper are:

  * •

We formalize the medical coding task as a recommendation problem and present a
novel multi-stage retrieve and re-rank framework to make more accurate
predictions by ruling out the irrelevant codes before ranking, rather than
making direct predictions on the entire large label set.

  * •

To address the large label set and long-tailed distribution issues, in the
two-stage retrieval process we use external knowledge and BM25 to retrieve a
subset of candidate labels from the large label space. We further leverage the
code co-occurrence in the re-ranking stage to capture the internal connections
among the codes.

  * •

We apply contrastive learning in the re-ranking stage. It effectively pulls
together the representations of a clinical note and its corresponding golden
truth labels, which allows the model to make more accurate predictions.

##  2 Related Work

The automatic ICD indexing task is well established in the healthcare domain.
Extensive research using deep learning has been dedicated to ICD indexing,
including recurrent-based neural networks (RNNs), convolution-based neural
networks (CNNs), and their variations Mullenbach et al. (2018); Li and Yu
(2020); Shi et al. (2017); Xie and Xing (2018). These architectures are able
to extract and categorize semantic features, reducing the need for medical
domain expertise during the traditional feature selection stage seen in
conventional algorithms Teng et al. (2023). The ICD indexing task is
formulated as a multi-label classification problem in these approaches.
Mullenbach et al. (2018) introduced a combination of CNN with an attention
mechanism to effectively capture pertinent information within clinical texts
for each ICD code. Building on this foundation, Xie et al. (2019) enhanced the
CNN attention model by integrating a multi-scale feature attention technique.
Many CNN variants were subsequently introduced to address the challenges posed
by lengthy and noisy clinical texts, including MultiResCNN Li and Yu (2020),
DCAN Ji et al. (2020), and EffectiveCAN Liu et al. (2021). RNN-based models,
renowned for their capacity to capture contextual information across input
texts, have also been widely used for ICD indexing. Shi et al. (2017) proposed
a character-aware Long Short-Term Memory (LSTM) recurrent network to learn the
underlying representations of clinical texts. Xie and Xing (2018) introduced a
tree-of-sequences LSTM architecture alongside adversarial learning to capture
hierarchical relationships among ICD codes. Additionally, Baumel et al. (2018)
presented a Hierarchical Attention-Bidirectional Gated Recurrent Unit (HA-GRU)
model, facilitating document labeling by identifying sentences relevant to
each ICD code. LAAT Vu et al. (2020) used a bidirectional Long-Short Term
Memory (BiLSTM) encoder and a customized label-wise attention mechanism to
cultivate label-specific vectors across distinct clinical text fragments.

To address the hierarchical relationships intrinsic to ICD codes, Graph
Convolutional Neural Networks (GCNNs) Kipf and Welling (2017) have emerged as
a powerful tool. Rios and Kavuluru (2018) and Xie et al. (2019) used GCNNs to
capture both the hierarchical interplay among ICD codes and the semantic
information specific to each code. HyperCore Cao et al. (2020) took a
comprehensive approach by considering both code hierarchy and code co-
occurrence, employing GCNNs to learn code representations within the co-graph.

Incorporating external knowledge beyond ICD code information has also gained
traction. Bai and Vucetic (2019) introduced a Knowledge Source Integration
(KSI) model that integrates external knowledge from Wikipedia. This
integration calculated matching scores between clinical notes and disease-
related Wikipedia documents, in order to enrich the available information for
ICD predictions. Additionally, Yuan et al. (2022) proposed a Multiple Synonym
Matching Network (MSMN) to use synonyms of ICD codes, enhancing the quality of
code representation learning. Expanding on this, Yang et al. (2022) integrated
a pre-trained language model with three domain-specific knowledge sources:
code hierarchy, synonyms, and abbreviations. This fusion of knowledge sources
contributes significantly to the performance of ICD classification.

##  3 Method

###  3.1 A Multi-stage Framework

Figure 2: Overview of the proposed multi-stage retrieve and re-rank framework.
The model first leverages auxiliary knowledge and BM25 to retrieve a candidate
list from the full label space, then uses a re-rank model that leverages the
code co-occurrence guided contrastive learning to generate the final relevant
labels.

We formulate the medical coding task as a recommendation task given medical
records $\mathcal{D}=\\{d_{1},d_{2},...,d_{N}\\}$ and a set of ICD codes
$\mathcal{Y}=\\{y_{1},y_{2},...,y_{L}\\}$ with associated external auxiliary
knowledge $\mathcal{K}$. We construct the label information as a graph
structure $\mathcal{G}$, using code co-occurrence relations, and we train a
multi-stage recommender system $\mathcal{R}$, based on the text information
$\mathcal{D}$, constructed label information $\mathcal{G}$, and the external
auxiliary knowledge $\mathcal{K}$. The system $\mathcal{R}$ needs to predict
the relevant labels given a document $d\notin\mathcal{D}$.

In this section, we present a multi-stage retrieve and re-rank framework for
ICD indexing, which is shown in Figure 2. Our model is composed of a two-stage
retrieval process that uses auxiliary knowledge of the EHR and BM25 to obtain
a shortened candidate list, and a re-ranking process that conducts code co-
occurrence guided contrastive learning to further improve the recommended ICD
list.

###  3.2 The Retrieval Stage

#### Using Auxiliary Knowledge

To retrieve the candidate list using auxiliary knowledge, we incorporate
insights from three external sources of knowledge: diagnosis-related group
(DRG) codes, current procedural terminology (CPT) codes, and medications
prescribed to patients. DRG codes are used by hospitals and healthcare
providers to classify patients into groups based on their diagnosis,
treatment, and length of stay. These codes are used for reimbursement
purposes, and they help determine the amount that healthcare providers are
remunerated for their services. DRG codes are further classified into medical
DRGs (which exclude operating room procedures) and surgical DRGs. CPT codes
are used to describe medical procedures and services provided by healthcare
providers. They provide a standardized way of documenting and billing for
medical services. CPT codes are used by insurance companies to determine
reimbursement rates for healthcare providers. Such code terminologies
significantly contribute to the refinement of ICD indexing. Moreover, the
medications prescribed to patients offer a wealth of predictive information
for ICD codes. These prescriptions often mark the conclusion of a patient’s
care episode. As patients approach the conclusion of their treatment, the
prescribed medications serve a critical role in managing their conditions.
Consequently, these medications emerge as potent indicators of underlying
health conditions or diagnoses. Their inclusion in the retrieval process
greatly enhances the accuracy and relevance of the corresponding ICD code
recommendations. The aforementioned auxiliary knowledge, such as DRG codes,
CPT codes, and drug prescriptions, typically appears in the EHR data and is
readily accessible.

Given a clinical note $d$, we retrieve the candidate ICD list by calculating
the auxiliary knowledge and label co-occurrence matrix using conditional
probabilities, i.e., $P(y_{i}\,|\,k_{j})$, which denote the probabilities of
occurrence of ICD $y_{i}$ when auxiliary knowledge $k_{j}$ appears.

$P(y_{i}\,|\,k_{j})=\frac{C_{y_{i}\cap k_{j}}}{C_{k_{j}}},$ (1)

where $C_{y_{i}\cap k_{j}}$ denotes the number of co-occurrences of $y_{i}$
and $k_{j}$, and $C_{k_{j}}$ is the number of occurrences of $k_{j}$ in the
training set. To avoid the noise of rare co-occurrences, a threshold $\eta$
filters noisy correlations. $\tilde{K}_{j}$ denotes the selected ICD set for
auxiliary knowledge $j$.

$\tilde{K}_{j}=\\{y_{i}|P(y_{i}|k_{j})>\eta,\;i=1,...,L\\},$ (2)

where $L$ is the total number of ICD codes in the label set, and $\eta=0.005$.
We then join the ICD codes retrieved from the auxiliary knowledge co-
occurrences for the DRG codes, CPT codes and prescribed drugs to form the
candidate ICD subset $\mathcal{C}_{\mathrm{auxiliary}}$:

$\mathcal{C}_{\mathrm{auxiliary}}(d)=\\\
\tilde{K}_{\mathrm{DRG}}(d)\cup\tilde{K}_{\mathrm{CPT}}(d)\cup\tilde{K}_{\mathrm{drug}}(d),$
(3)

where $\mathcal{C}_{\mathrm{auxiliary}}\subseteq\mathcal{Y}$.

#### Using BM25

The retrieval stage using auxiliary knowledge incorporates the co-relations
between ICD codes and external knowledge, but ignores the relationship between
clinical texts and labels. To increase the recall of the retrieval stage, we
adopt BM25 Robertson and Walker (1994) to allow lexical matching between the
medical documents and labels on the retrieved candidate list
$\mathcal{C}_{\mathrm{auxiliary}}$. Given a medical record $d$ and an ICD code
$y$, the score between $d$ and $y$ is calculated as:

$\mathrm{BM25}(d,y)=\\\ \sum_{w\in d\cap
t_{y}}\mathrm{IDF}(w)\frac{\scriptstyle\mathrm{TF}(w,t_{y})\cdot(k_{1}+1)}{\scriptstyle\mathrm{TF}(w,t_{y})\cdot
k_{1}(1-b+b\frac{|\mathcal{Y}|}{\textit{avgdl}})},$ (4)

and

$\textit{avgdl}=\frac{1}{|\mathcal{Y}|}\sum_{y\in\mathcal{Y}}|t_{y}|,$ (5)

where $t_{y}$ represents the words in the label descriptors, $|\mathcal{Y}|$
is the length of the label descriptors in words, avgdl is the average length
of text information in the label.

When the BM25 score between $d$ and $y_{i}$ exceeds a certain threshold
$\theta$, we add $y_{i}$ as a candidate of $d$:

$\mathcal{C}_{\mathrm{BM25}}(d)=\\\
\\{{y_{i}}|\mathrm{BM25}(d,y_{i})>\theta,y_{i}\in\mathcal{C}_{\mathrm{auxiliary}}\\},$
(6)

where $\theta=200$. Given a clinical note $d$, its candidate ICD set is first
generated by using the auxiliary knowledge in the retrieval stage and then
reduced by using BM25, where
$\mathcal{C}_{\mathrm{BM25}}\subseteq\mathcal{C}_{\mathrm{auxiliary}}$ and
$\mathcal{C}_{\mathrm{auxiliary}}\subseteq\mathcal{Y}$.

###  3.3 The Re-ranking Stage

#### Clinical Text Encoder

Encouraged by the success of the pre-trained language model Longformer Beltagy
et al. (2020) in dealing with longer texts, we use Clinical-Longformer Li et
al. (2023), specifically pre-trained in the medical domain, as a text encoder.
Given a medical document $d$ as input that consists of a sequence of tokens:

$d=\\{[\texttt{CLS}],x_{1},x_{2},...,x_{n-2},[\texttt{SEP}]\\},$ (7)

where $[\texttt{CLS}]$ and $[\texttt{SEP}]$ are two special tokens that
indicate the beginning and end of the sequence, and $n$ is the sequence
length, the Clinical-Longformer encodes the tokens and outputs the hidden
representations for each token:

$H_{\mathrm{hidden}}=\mathrm{ClinicalLongformer}(d),$ (8)

where $H_{\mathrm{hidden}}\in\mathbb{R}^{n\times h_{e}}$, and $h_{e}$ is the
hidden size. Following previous work Wang et al. (2022); Yang et al. (2022),
we use the hidden state of the $\mathrm{[CLS]}$ token to represent the
document, which is the first token of $H_{\mathrm{hidden}}$, denoted as
$H_{\mathrm{T}}$.

#### Label Encoder

The occurrence of two ICD codes together in clinical texts frequently
indicates a simultaneous presence or a causal connection between specific
diseases. This implies that the codes representing these interconnected
diseases often manifest together within clinical notes. We employ a Graphormer
Ying et al. (2021) to incorporate the co-occurrence relationships among ICD
codes. Unlike the original GNN, Graphormer models graphs using Transformer
layers Vaswani et al. (2017) with spatial encoding and edge encoding, which
could effectively encode the structural information (i.e., code co-occurrence)
of a graph into the model. We create a directed code co-occurrence graph
$\mathcal{G}=(\mathcal{Y},\mathcal{E})$, where node set $\mathcal{Y}$ is the
labels and edge set $\mathcal{E}$ denotes the co-occurrence relations. This
graph is constructed using the code co-occurrence matrix, which has been used
as the edge matrix for the graph. We create the code co-occurrence matrix by
using the correlated relationship between labels based on conditional
probabilities. This approach encapsulates the interdependence between various
ICD codes in a quantifiable manner, offering valuable insights into the
underlying connections among disease codes within the clinical texts. To be
more specific, we calculate the probability of occurrence of label $y_{j}$
when label $y_{i}$ appears as follows:

$P(y_{j}\,|\,y_{i})=\frac{C_{y_{i}\cap y_{j}}}{C_{y_{i}}}$ (9)

where $C_{y_{j}\cap y_{i}}$ denotes the number of co-occurrences of $y_{i}$
and $y_{j}$, and $C_{y_{i}}$ is the number of occurrences of $y_{i}$ in the
training set. To facilitate graph construction, we binarize the correlation
probability $P(y_{j}\,|\,y_{i})$. This entails converting the probability
values into binary values which indicates whether a correlation exists (or
not) between two labels. The operation can be written as:

$\mathcal{E}_{ij}=\begin{cases}0,\text{if }P(y_{j}\,|\,y_{i})<\lambda\\\
1,\text{if }P(y_{j}\,|\,y_{i})\geq\lambda,\end{cases}$ (10)

where $\mathcal{E}$ is the binary correlation matrix that is used to form the
edge set, and $\lambda$ is the hyper-parameter threshold to filter the noise
edges. In our experiment, $\lambda=1$, which means that a edge is formed when
the two labels in each pair always appear together.

To encode the graph $\mathcal{G}$, we first generate the initial node features
using the ICD full descriptors for each code $y$ via Clinical-Longformer:

$\begin{gathered}\mathrm{y}=\\{[\texttt{CLS}],x_{1},x_{2},...,x_{n-2},[\texttt{SEP}]\\},\\\
H_{v}=\mathrm{ClinicalLongformer}(\mathrm{y}),\end{gathered}$ (11)

where $\mathrm{y}$ represents a sequence of words in the label descriptors of
label $y$, $H_{v}\in\mathbb{R}^{n\times h_{e}}$, and $h_{e}$ is the hidden
size. We use the hidden state of the first token ($\mathrm{[CLS]}$) to
represent the initial node feature denoted as $H_{\mathrm{node}}^{i}$ for the
$i^{th}$ label.

With all initial node features stacked as a matrix
$V=\\{H_{\mathrm{node}}^{1},H_{\mathrm{node}}^{2},...,H_{\mathrm{node}}^{L}\\}$,
where $V\in\mathbb{R}^{h_{e}\times L}$, a standard self-attention layer is
then used for feature migration. To leverage the structural information, a
novel spatial encoding method is used to modify the Query-Key product matrix
$A^{\mathcal{G}}$ in the self-attention layer:

$A^{\mathcal{G}}_{\textit{ij}}=\frac{(H_{\mathrm{node}}^{i}W^{\mathcal{G}}_{Q})(H_{\mathrm{node}}^{j}W^{\mathcal{G}}_{K})^{\intercal}}{\sqrt{h_{e}}}+b_{\phi(y_{i},y_{j})},$
(12)

where $W^{\mathcal{G}}_{Q}$ and $W^{\mathcal{G}}_{K}$ are layer-specific
weight matrices, and $\phi(y_{i},y_{j})$ is the spatial relation between
$y_{i}$ and $y_{j}$ in graph $\mathcal{G}$, and the function $\phi(\cdot)$ is
defined as the connectivity between the nodes in $\mathcal{G}$, which is the
co-occurrence relation among labels. $b_{\phi(y_{i},y_{j})}$ is a learnable
scalar indexed by $\phi(y_{i},y_{j})$, and shared across all layers. The
attention score $A^{\mathcal{G}}_{\textit{ij}}$, then, has been used to
aggregate the multi-head attention for the final output:

$h^{l+1}=\mathrm{MHA}(\mathrm{LN}(h^{l}))+h^{l},$ (13)

where $\mathrm{LN}$ denotes the layer normalization, $\mathrm{MHA}$ denotes
the multi-head self-attention, $h^{l}$ and $h^{l+1}\in\mathbb{R}^{L\times
h_{e}}$ indicate the node representation of the $l^{th}$ and $(l+1)^{th}$
layers. We use the last layer to represent the label feature denoted as
$H_{L}$. For more details on the full structure of Graphormer, please refer to
the original paper Ying et al. (2021).

#### Contrastive Learning for Re-ranking

Now, we construct a code co-occurrence guided contrastive learning framework.
Unlike supervised learning that aims to understand “what is what”, contrastive
learning adopts a different perspective by learning “what is similar or
dissimilar to what”. In our problem setting, we focus on the distances between
a clinical document and its associated ICD codes, rather than solely between
samples themselves. We consider the ground truth labels as positive samples,
while the negative samples comprise all the other labels within the label
space. Given $H_{\mathrm{T}}$, the representation for a clinical note $d$, and
the set of representations of its corresponding ICD codes denoted as
$H^{+}_{\mathrm{L}}$, we denote the representations of $N$ negative ICD codes
randomly chosen from the ICD codes of the documents in the batch (batch size
is $N$), which are not ICD codes of document $d$, as $H^{-}_{\mathrm{L}}$.
Contrastive learning aims to learn the effective representations by pulling
$d$ and $H^{+}_{\mathrm{L}}$ together while pushing apart $d$ and
$H^{-}_{\mathrm{L}}$, represented as $S$ and $D$, respectively, in the
equation below. The contrastive loss can be defined as:

$\mathcal{L}=-\mathrm{log}\frac{S/\tau}{S/\tau+D/\tau},$ (14)

where $S=\exp(\sum_{c\in
L^{+}_{\mathrm{L}}}\cos(H_{\mathrm{T}},c)/|H^{+}_{\mathrm{L}}|)$,
$D=\exp(\sum_{c^{\prime}\in
L^{-}_{\mathrm{L}}}\cos(H_{\mathrm{T}},c^{\prime})/N)$, and $\tau$ is the
temperature hyper-parameter. During inference, a comparison is conducted by
measuring the distance between the query clinical note and ICD codes in the
embedding space, which ranks the ICD codes and then provides recommendations
of the potential ICD candidates.

##  4 Experiments

###  4.1 Dataset and Pre-processing

We conduct our experiments on the publicly available benchmark MIMIC-III
Johnson et al. (2016) dataset that contains a variety of patient data types,
including discharge summaries, demographic details, interventions, laboratory
results, physiologic measures, and medication information. Following previous
work, we are interested in the de-identified discharge summaries with
annotated ICD-9 codes. There are 52,722 discharge summaries and 8,922 unique
ICD-9 codes in the dataset. We mainly use three major data resources from the
dataset: (1) de-identified discharge summaries (from the NOTEEVENTS table);
(2) ICD-9 codes (from DIAGNOSES_ICD and PROCEDURES_ICD tables); and (3)
auxiliary knowledge including DRG codes, CPT codes and drug prescriptions
(from DRGCODES, CPTEVENTS, and PRESCRIPTIONS tables).

To preprocess the clinical notes, we first remove all de-identified
information, then replace punctuation and atypical alphanumerical character
combinations (e.g., ‘3a’, ‘4kg’) with white space, and lowercase every token.
We truncate the discharge summaries at a maximum length of 4000 tokens. We
follow Mullenbach et al. (2018) to form two settings: full codes (MIMIC-III-
full) and top-50 frequent codes (MIMIC-III-top 50). In MIMIC-III-full, there
are 47,719 discharge summaries for training, with 1,632 for validation, and
with 3,372 for testing.

###  4.2 Implementation and Evaluation

We implement our model in PyTorch Paszke et al. (2019) on a single NVIDIA A100
40G GPU. We use the Adam optimizer and early stopping strategies using the
Micro-F1 score over the validation set as the stopping criterion to avoid
over-fitting. We set the initial learning rate as 5e-5 with batch size 16. We
choose a learning rate scheduler which is warmed up with cosine decay, and the
warm up ratio is set to 0.1. Our code is available at
https://github.com/xdwang0726/ICD-contrastive-curriculum.

For evaluating the performance of our proposed model, we employ three commonly
used metrics: F1-score (Micro and Macro), AUC (Micro and Macro), and precision
at $K$ ($\mathrm{P@K}$).

##  5 Results and Discussion

Models | MIMIC-III-full | MIMIC-III-top 50
---|---|---
AUC | F1 | P@K | AUC | F1 | P@5
Macro | Micro | Macro | Micro | P@8 | P@15 | Macro | Micro | Macro | Micro
CAML Mullenbach et al. (2018) | 0.895 | 0.986 | 0.088 | 0.539 | 0.709 | 0.561 | 0.875 | 0.909 | 0.532 | 0.614 | 0.609
DR-CAML Mullenbach et al. (2018) | 0.897 | 0.985 | 0.086 | 0.529 | 0.690 | 0.548 | 0.884 | 0.916 | 0.576 | 0.633 | 0.618
MultiResCNN Li and Yu (2020) | 0.910 | 0.986 | 0.085 | 0.552 | 0.734 | 0.584 | 0.899 | 0.928 | 0.606 | 0.670 | 0.641
LAAT Vu et al. (2020) | 0.919 | 0.988 | 0.099 | 0.575 | 0.738 | 0.591 | 0.925 | 0.946 | 0.666 | 0.715 | 0.675
Joint-LAAT Vu et al. (2020) | 0.921 | 0.988 | 0.107 | 0.575 | 0.735 | 0.590 | 0.925 | 0.946 | 0.661 | 0.716 | 0.671
EffectiveCAN Liu et al. (2021) | 0.915 | 0.988 | 0.106 | 0.589 | 0.758 | 0.606 | 0.915 | 0.938 | 0.644 | 0.702 | 0.656
MSMN Yuan et al. (2022) | 0.950 | 0.992 | 0.103 | 0.584 | 0.752 | 0.599 | 0.928 | 0.947 | 0.683 | 0.725 | 0.680
KEPTLongformer Yang et al. (2022) | - | - | 0.118 | 0.599 | 0.771 | 0.615 | 0.926 | 0.947 | 0.689 | 0.728 | 0.672
Ours | 0.949 | 0.995 | 0.114 | 0.603 | 0.775 | 0.623 | 0.927 | 0.947 | 0.687 | 0.732 | 0.685

Table 1: Comparison to previous methods across three main evaluation metrics
MIMIC-III dataset. Bold: the optimal values.

In order to asses the efficacy of our proposed framework, we compare with the
existing state-of-the-art (SOTA) models, as outlined in Table 1. The top score
for each metric is denoted in bold. As shown, our model outperforms in the
majority of evaluation metrics, with the exception of Macro-AUC and Macro-F1
on MIMIC-III-full and MIMIC-III-top 50. Notably, our model achieves comparable
performance on Micro-F1 and Micro-AUC, and improves precision at $K$ on both
MIMIC-III-full and MIMIC-III-top 50. These results provide solid evidence to
validate the efficacy of integrating auxiliary knowledge in the retrieval
stage and leveraging code co-occurrence guided contrastive learning in the re-
ranking stage.

As the occurrence frequencies of the ICD codes are imbalanced, our focus lies
in assessing the efficacy of our model, specifically on the infrequently
appearing ICD codes. We categorize the ICD codes into four groups based on
their occurrences in the training set: [0, 10), [10, 50), [50, 500), and
[500,). Figure 3 illustrates the distribution of ICD codes and their
occurrence percentages across the four categorized groups in the training set,
which show that the distribution of ICD frequency is highly biased, conforming
to a long-tail distribution. Figures 3b and 3c present the performance of our
model on MIMIC-III-full in comparison to the CAML baseline Mullenbach et al.
(2018) across the four ICD groups on Macro-AUC and Micro-F1, respectively. Our
model demonstrates significant improvements for both frequent and infrequent
labels on both metrics.

Figure 3: (a) ICD code distribution. (b) Macro-AUC performance comparison of
our model and CAML on ICD codes at different frequency. (c) Micro-F1
performance comparison of our model and CAML on ICD codes at different
frequency.

To confirm the specific contributions of these modules in terms of enhancing
both the effectiveness and robustness of the model, we conduct ablation
studies with three different settings: (a) we examine the effectiveness of
using auxiliary knowledge in the retrieval stage by removing the retrieval
stage and rank the ICD codes on the whole label set; (b) we examine the
influence of different embedding methods by replacing the Clinical-Longformer
with Clinical-BERT; and (c) we test the effectiveness of label embedding by
replacing the encoding of the label with the average of word embeddings in the
label descriptors. The experimental results are shown in Table 2. We also
conduct case studies to qualitatively understand the effects of incorporating
the label co-occurrence and the auxiliary knowledge. Two case studies have
been presented in Appendix A.

Methods | F1 | P@K
---|---|---
Macro | Micro | P@8 | P@15
Full Model | 0.114 | 0.603 | 0.775 | 0.623
w/o auxiliary knowledge | 0.097 | 0.579 | 0.748 | 0.587
embedded w/ Clinical-BERT | 0.083 | 0.548 | 0.711 | 0.546
w/o Graphormer | 0.102 | 0.583 | 0.753 | 0.591

Table 2: Ablation experiment results on the MIMIC-III-full. Bold: the optimal
values.

#### Effectiveness of Using Auxiliary Knowledge for Retrieval

We employ three distinct types of auxiliary knowledge in the retrieval stage:
DRG codes, CPT codes, and drug prescriptions. As shown in Table 2, removing
auxiliary knowledge leads to a decline in performance, indicating the pivotal
role of the retrieval stage. This outcome further provides evidence that
external knowledge effectively addresses the challenge presented by a large
pool of potential ICD codes. Through integrating external knowledge, the
retrieval stage attains the capability to refine the candidate list using the
co-occurrence relationships between ICD codes and the auxiliary knowledge,
thereby amplifying both the efficiency and accuracy of the re-ranking stage.
The selection of an appropriate candidate list for a given medical record
hinges upon a hyper-parameter, specifically the threshold $\eta$ governing the
co-occurrence between auxiliary knowledge and ICD codes. The choice of $\eta$
determines the candidate numbers that implicitly affect the overall
performance of the model. Setting $\eta=0.005$, the candidate list guarantees
inclusion of 99.22% of the gold-standard ICD codes, resulting in an average of
1,460 codes in the subset. Notably, this accounts for approximately one-sixth
of the complete code set. A further reduction using BM25 limits the candidate
list to 1,299 on average.

#### Comparison of Clinical-Longformer and Clinical-BERT

Increasing the maximum token limit is important in the context of clinical
notes analysis as clinical texts are lengthy. Specially, in the MIMIC-III
dataset, the average length of the discharge summaries is 1,596. Given this
substantial token volume in the clinical notes, encoding a maximum number of
tokens prior to downstream analysis becomes a pivotal requirement, which
facilitates robust and meaningful subsequent analysis. To test the
effectiveness of using longer sequences, we compare the model performance of
Clinical-Longformer and a BERT-based pre-trained language model (i.e.,
Clinical-BERT) which can encode a maximum of 512 tokens. As shown in Table 2,
Clinical-Longformer substantially outperforms Clinical-BERT, indicating the
importance of the maximum token limit on language models in the automatic
medical coding task.

#### Effectiveness of Learning Label Features Using Code Co-occurrence

The graph structure has been shown to be effective in modeling code
correlations and Graphormer efficiently learns code representations. The
findings presented in Table 2 highlight the affirmative impact of integrating
code co-occurrence into label representations. By using Graphormer, the model
effectively captures and exploits the intricate connections and
interdependencies among the labels, thereby improving the overall performance.
This indicates that incorporating code co-occurrence information with
Graphormer empowers the model to gain insights from the collaborative
behaviours of the labels, consequently facilitating a more holistic
comprehension of the underlying label co-relations.

##  6 Conclusion

In this paper, we regard the medical coding task as a recommendation problem
and present a novel multi-stage retrieve and re-rank framework. The primary
objective of the proposed framework is twofold: to construct a curated list of
ICD codes and, subsequently, to further refine the candidate list for a given
medical record. Specifically, we first conduct a two-step retrieval process,
incorporating auxiliary knowledge and the BM25 algorithm. This approach
retrieves a concise subset of the candidate list, mitigating the challenges of
a very large and imbalanced label distribution. We then use a re-ranking model
to refine the previously obtained candidate list, employing code co-occurrence
guided contrastive learning. Experimental results demonstrate that our
proposed framework outperforms the previous SOTA, which suggests that it
provides more precise and contextually grounded ICD recommendations for the
given medical records. In the future, our proposed framework may be extended
with more external knowledge such as the Unified Medical Language System
(UMLS) and code synonymy.

## Limitations

Our usage of auxiliary knowledge is limited to external knowledge that
includes DRG codes, CPT codes, and drug prescriptions, only. Other knowledge
including disease-symptom, disease-lab relations, Unified Medical Language
System (UMLS), and others, could also be potentially useful for the auto ICD
coding task. We also acknowledge that the auxiliary knowledge we used is
labeled by human annotators, which may require some extra effort. We are not
quite sure about the workload for annotating different code terminologies, but
we believe linking different code terminologies is important.

Our study is constrained by its evaluation limited to MIMIC-III-full and
MIMIC-III-top 50 datasets, primarily concentrated on common disease. To
comprehensively assess the model’s performance on rare diseases, future work
could benefit from a curated list of rare diseases validated by domain
experts.

## Ethics Statement

We are using the publicly available clinical dataset MIMIC-III, which contains
de-identified patient information. We do not see any ethics issue here in this
paper.

## Acknowledgements

We would like to thank all reviewers for their comments, which helped improve
this paper considerably. Computational resources used in preparing this
research were provided, in part, by Compute
Ontario666https://www.computeontario.ca, Digital Research Alliance of
Canada777https://ccdb.alliancecan.ca, the Province of Ontario, the Government
of Canada through CIFAR, and companies sponsoring the Vector
Institute888https://www.vectorinstitute.ai/partners. This research is
partially funded by The Natural Sciences and Engineering Research Council of
Canada (NSERC) through a Discovery Grant to R. E. Mercer. F. Rudzicz is
supported by a CIFAR Chair in AI.

## References

  * Bai and Vucetic (2019) Tian Bai and Slobodan Vucetic. 2019.  Improving medical code prediction from clinical text via incorporating online knowledge sources.  _The World Wide Web Conference_ , page 72–82.
  * Baumel et al. (2018) Tal Baumel, Jumana Nassour-Kassis, Michael Elhadad, and Noémie Elhadad. 2018.  Multi-label classification of patient notes a case study on ICD code assignment.  In _The Workshops of the Thirty-Second AAAI Conf. on Art. Intell._ , pages 409–416.
  * Beltagy et al. (2020) Iz Beltagy, Matthew E. Peters, and Arman Cohan. 2020.  Longformer: The long-document transformer.  _arXiv:2004.05150_.
  * Cao et al. (2020) Pengfei Cao, Yubo Chen, Kang Liu, Jun Zhao, Shengping Liu, and Weifeng Chong. 2020.  Hypercore: Hyperbolic and co-graph representation for automatic icd coding.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , page 3105–3114.
  * Ji et al. (2020) Shaoxiong Ji, Erik Cambria, and Pekka Marttinen. 2020.  Dilated convolutional attention network for medical code assignment from clinical text.  In _Proceedings of the 3rd Clinical Natural Language Processing Workshop_ , pages 73–78.
  * Johnson et al. (2016) Alistair E. W. Johnson, Tom J. Pollard, Lu Shen, Li wei H. Lehman, Mengling Feng, Mohammad Mahdi Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G. Mark. 2016.  MIMIC-III, a freely accessible critical care database.  _Scientific Data_ , 3.
  * Kipf and Welling (2017) Thomas N. Kipf and Max Welling. 2017.  Semi-Supervised Classification with Graph Convolutional Networks.  In _Proceedings of the 5th Int. Conference on Learning Representations_.
  * Li and Yu (2020) Fei Li and Hong Yu. 2020.  ICD coding from clinical text using multi-filter residual convolutional neural network.  In _Proceedings of the Thirty-Fourth AAAI Conference on Artificial Intelligence_ , pages 8180–8187.
  * Li et al. (2023) Yikuan Li, Ramsey M Wehbe, Faraz S Ahmad, Hanyin Wang, and Yuan Luo. 2023.  A comparative study of pretrained language models for long clinical text.  _Journal of the American Medical Informatics Association_ , 30(2):340–347.
  * Liu et al. (2021) Yang Liu, Hua Cheng, Russell Klopfer, Matthew R. Gormley, and Thomas Schaaf. 2021.  Effective convolutional attention network for multi-label clinical document classification.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pages 5941–5953.
  * Mullenbach et al. (2018) James Mullenbach, Sarah Wiegreffe, Jon Duke, Jimeng Sun, and Jacob Eisenstein. 2018.  Explainable prediction of medical codes from clinical text.  In _Proc. of the 2018 Conf. of the North American Chapter of the Ass. for Comp. Ling.: Human Language Technologies, Volume 1 (Long Papers)_ , pages 1101–1111.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019.  PyTorch: An imperative style, high-performance deep learning library.  In _Advances in Neural Information Processing Systems 32_ , pages 8024–8035.
  * Rios and Kavuluru (2018) Anthony Rios and Ramakanth Kavuluru. 2018.  Few-shot and zero-shot multi-label learning for structured label spaces.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pages 3132–3142.
  * Robertson and Walker (1994) S. E. Robertson and S. Walker. 1994.  Some simple effective approximations to the 2-poisson model for probabilistic weighted retrieval.  In _Proceedings of the 17th Annual International ACM SIGIR Conference on Research and Development in Information Retrieval_ , page 232–241.
  * Shi et al. (2017) Haoran Shi, Pengtao Xie, Zhiting Hu, Ming Zhang, and Eric P. Xing. 2017.  Towards automated ICD coding using deep learning.  _ArXiv_ , abs/1711.04075.
  * Teng et al. (2023) Fei Teng, Yiming Liu, Tianrui Li, Yi Zhang, Shuangqing Li, and Yue Zhao. 2023.  A review on deep neural networks for icd coding.  _IEEE Trans. on Knowledge and Data Eng._ , 35(5):4357–4375.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In _Proc. of the 31st Int. Conference on Neural Information Processing Systems_ , page 6000–6010.
  * Vu et al. (2020) Thanh Vu, Dat Quoc Nguyen, and Anthony Nguyen. 2020.  A label attention model for ICD coding from clinical text.  In _Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI-20_ , pages 3335–3341.  Main track.
  * Wang et al. (2022) Zihan Wang, Peiyi Wang, Lianzhe Huang, Xin Sun, and Houfeng Wang. 2022.  Incorporating hierarchy into text encoder: a contrastive learning approach for hierarchical text classification.  In _Proc. of the 60th Annual Meeting of the Assoc. for Computational Linguistics (Volume 1: Long Papers)_ , pages 7109–7119.
  * Xie and Xing (2018) Pengtao Xie and Eric Xing. 2018.  A neural architecture for automated ICD coding.  In _Proc. of the 56th Annual Meeting of the Ass. for Comp. Ling. (Vol. 1: Long Papers)_ , pages 1066–1076.
  * Xie et al. (2019) Xiancheng Xie, Yun Xiong, Philip S. Yu, and Yangyong Zhu. 2019.  EHR coding with multi-scale feature attention and structured knowledge graph propagation.  In _Proc. of the 28th ACM Int. Conf. on Information and Knowledge Management_ , page 649–658.
  * Yang et al. (2022) Zhichao Yang, Shufan Wang, Bhanu Pratap Singh Rawat, Avijit Mitra, and Hong Yu. 2022.  Knowledge injected prompt based fine-tuning for multi-label few-shot ICD coding.  In _Findings of the Assoc. for Computational Linguistics: EMNLP 2022_ , pages 1767–1781.
  * Ying et al. (2021) Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu. 2021.  Do transformers really perform badly for graph representation? In _Thirty-Fifth Conference on Neural Information Processing Systems_.
  * Yuan et al. (2022) Zheng Yuan, Chuanqi Tan, and Songfang Huang. 2022.  Code synonyms do matter: Multiple synonyms matching network for automatic ICD coding.  In _Proc. of the 60th Annual Meeting of the Assoc. for Comp. Ling. (Vol. 2: Short Papers)_ , pages 808–814.

##  Appendix A Case Studies

Figure 4: Case study on the effectiveness of incorporating label co-
occurrence. Correctly predicted labels are marked in green and the incorrect
ones are marked in red.

We conducted case studies to qualitatively explore the impacts of integrating
label co-occurrence (illustrated in Figure 4) and auxiliary knowledge
(depicted in Figure 5). We compared the full model with models that did not
integrate the label co-occurence and the external knowledge on the predictions
of two patient records. For each patient, we present the discharge summary,
ground truth ICD codes, label co-occurrence information, and auxiliary
knowledge information, along with the predicted ICD codes from the full model
and ablated models.

In Case 1, the ground truth ICD codes ‘785.51 Cardiogenic shock’ and ‘V49.86
Do not resuscitate status’ are not explicitly mentioned in the discharge
summary. The observed label co-occurrence between ‘427.5 Cardiac arrest’ and
‘785.51 Cardiogenic shock’, as well as co-relation between ‘96.71 Continuous
invasive mechanical ventilation for less than 96 consecutive hours’ and
‘V49.86 Do not resuscitate status’ provide strong indicators suggesting the
presence of the codes ‘785.51’ and ‘V49.86’. Without the label co-occurrence
signals, the ablated model missed the predictions of codes ‘785.51’ and
‘V49.86’, indicating a failure to leverage latent label information.

In Case 2, the patient has been diagnosed with ‘244.9 Unspecified acquired
hypothyroidism’ with less explicit information in the discharge summary.
Notably, the presence of the medication ‘Levothyroxine’ in the drug
prescription, an element of auxiliary knowledge, suggests that the patient is
likely to have acquired hypothyroidism. The ablated model, lacking the
auxiliary knowledge, misses the prediction of code ‘244.9’. The aforementioned
Cases 1 and 2 highlight the benefits of incorporating label co-occurrence and
auxiliary knowledge, respectively.

Figure 5: Case study on the effectiveness of incorporating auxiliary
knowledge. Correctly predicted labels are marked in green and the incorrect
ones are marked in red.


# VOICE CLONING: A MULTI-SPEAKER TEXT-TO-SPEECH SYNTHESIS APPROACH BASED ON
TRANSFER LEARNING

###### Abstract

Deep learning models are becoming predominant in many fields of machine
learning. Text-to-Speech (TTS), the process of synthesizing artificial speech
from text, is no exception. To this end, a deep neural network is usually
trained using a corpus of several hours of recorded speech from a single
speaker. Trying to produce the voice of a speaker other than the one learned
is expensive and requires large effort since it is necessary to record a new
dataset and retrain the model. This is the main reason why the TTS models are
usually single speaker. The proposed approach has the goal to overcome these
limitations trying to obtain a system which is able to model a multi-speaker
acoustic space. This allows the generation of speech audio similar to the
voice of different target speakers, even if they were not observed during the
training phase.

Index Terms—  text-to-speech, deep learning, multi-speaker speech synthesis,
speaker embedding, transfer learning

##  1 Introduction

Text-to-Speech (TTS) synthesis, the process of generating natural speech from
text, remains a challenging task despite decades of investigation. Nowadays
there are several TTS systems able to get impressive results in terms of
synthesis of natural voices very close to human ones. Unfortunately, many of
these systems learn to synthesize text only with a single voice. The goal of
this work is to build a TTS system which can generate in a data efficient
manner natural speech for a wide variety of speakers, not necessarily seen
during the training phase. The activity that allows the creation of this type
of models is called Voice Cloning and has many applications, such as restoring
the ability to communicate naturally to users who have lost their voice or
customizing digital assistants such as Siri.

Over time, there has been a significant interest in end-to-end TTS models
trained directly from text-audio pairs; Tacotron 2 [1] used WaveNet [2] as a
vocoder to invert spectrograms generated by sequence-to-sequence with
attention [3] model architecture that encodes text and decodes spectrograms,
obtaining a naturalness close to the human one. It only supported a single
speaker. Gibiansky et al. [4] proposed a multi-speaker variation of Tacotron
able to learn a low-dimensional speaker embedding for each training speaker.
Deep Voice 3 [5] introduced a fully convolutional encoder-decoder architecture
which supports thousands of speakers from LibriSpeech [6]. However, these
systems only support synthesis of voices seen during training since they learn
a fixed set of speaker embeddings. Voiceloop [7] proposed a novel architecture
which can generate speech from voices unseen during training but requires tens
of minutes of speech and transcripts of the target speaker. In recent
extensions, only a few seconds of speech per speaker can be used to generate
new speech in that speaker’s voice. Nachmani et al. [8] for example, extended
Voiceloop to utilize a target speaker encoding network to predict speaker
embedding directly from a spectrogram. This network is jointly trained with
the synthesis network to ensure that embeddings predicted from utterances by
the same speaker are closer than embeddings computed from different speakers.
Jia et al. [9] proposed a speaker encoder model similar to [8], except that
they used a network independently-trained exploring transfer learning from a
pre-trained speaker verification model towards the synthesis model.

This work is similar to [9] however introduces different architectures and
uses a new transfer learning technique still based on a pre-trained speaker
verification model but exploiting utterance embeddings rather than speaker
embeddings. In addition, we use a different strategy to condition the speech
synthesis with the voice of speakers not observed before and compared several
neural architectures for the speaker encoder model. The paper is organized as
follows: Section 2 describes the model architecture and its formal definition;
Section 3 reports experiments and results done to evaluate the proposed
solution; finally conclusions are reported in Section 4.

##  2 Model Architecture

Following [9], the proposed system consists of three components: a _speaker
encoder_ , which computes a fixed-dimensional embedding vector from a few
seconds of reference speech of a target speaker; a _synthesizer_ , which
predicts a mel spetrogram from an input text and an embedding vector; a
_neural vocoder_ , which infers time-domain waveforms from the mel
spectrograms generated by the synthesizer. At inference time, the speaker
encoder takes as input a short reference utterance of the target speaker and
generates, according to its internal learned speaker characteristics space, an
embedding vector. The synthesizer takes as input a phoneme (or grapheme)
sequence and generates a mel spectrogram, conditioned by the speaker encoder
embedding vector. Finally the vocoder takes the output of the synthesizer and
generates the speech waveform. This is illustrated in Figure 1.

###  2.1 Problem Definition

Fig. 1: High level overview of the three components of the system.

Consider a dataset of N speakers each of which has M utterances in the time-
domain. Let’s denote the j-th utterance of the i-th speaker as uij while the
feature extraced from the j-th utterance of the i-th speaker as xij ($1\leq
i\leq N$ and $1\leq j\leq M$). We chose as feature vector xij the mel
spectrogram.

The speaker encoder $\mathcal{E}$ has the task to produce meaningful embedding
vectors that should characterize the voices of the speakers. It computes the
embedding vector eij corresponding to the utterance uij as:

$\mathbf{e}_{ij}=\mathcal{E}\left(\mathbf{x}_{ij};\mathbf{w}_{\mathcal{E}}\right)$
(1)

where wE represents the encoder model parameters. Let’s define it _utterance
embedding_. In addition to defining embedding at the utterance level, we can
also define the _speaker embedding_ :

$\mathbf{c}_{i}=\frac{1}{n}\sum_{j=1}^{n}\mathbf{e}_{ij}$ (2)

In [9], the synthesizer $S$ predicts xij given cij and tij, the transcript of
the utterance uij:

$\hat{\mathbf{x}}_{ij}=\mathcal{S}\left(\mathbf{c}_{i},\mathbf{t}_{ij};\mathbf{w}_{\mathcal{S}}\right)$
(3)

where wS represents the synthesizer model parameters. In our approach, we
propose to use the utterance embedding rather than the speaker embedding:

$\hat{\mathbf{x}}_{ij}=\mathcal{S}\left(\mathbf{e}_{ij},\mathbf{t}_{ij};\mathbf{w}_{\mathcal{S}}\right)$
(4)

We will motivate this choice in Paragraph 2.4.

Finally, the vocoder $\mathcal{V}$ generates uij given
$\hat{\mathbf{\textbf{x}}}_{ij}$. So we have:

$\hat{\mathbf{u}}_{ij}=\mathcal{V}\left(\hat{\mathbf{x}}_{ij};\mathbf{w}_{\mathcal{V}}\right)$
(5)

where wV represents the vocoder model parameters.

This system could be trained in an end-to-end mode trying to optimize the
following objective function:

$\min_{\mathbf{w}_{\mathcal{E}},\mathbf{w}_{\mathcal{S}},\mathbf{w}_{\mathcal{V}}}L_{\mathcal{V}}\left(\mathbf{u}_{ij},\mathcal{V}\left(\mathcal{S}\left(\mathcal{E}\left(\mathbf{x}_{ij};\mathbf{w}_{\mathcal{E}}\right),\mathbf{t}_{ij};\mathbf{w}_{\mathcal{S}}\right);\mathbf{w}_{\mathcal{V}}\right)\right)$
(6)

where LV is a loss function in the time-domain. However, it requires to train
the three models using the same dataset, moreover, the convergence of the
combined model could be hard to reach. To overcome this drawback, the
synthesizer can be trained independently to directly predict the mel
spectrogram xij of a target utterance uij trying to optimize the following
objective function:

$\min_{\mathbf{w}_{\mathcal{S}}}L_{\mathcal{S}}\left(\mathbf{x}_{ij},\mathcal{S}\left(\mathbf{e}_{ij},\mathbf{t}_{ij};\mathbf{w}_{\mathcal{S}}\right)\right)$
(7)

where LS is a loss function in the time-frequency domain. It is necessary to
have a pre-trained speaker encoder model available to compute the utterance
embedding eij.

The vocoder can be trained either directly on the mel spectrograms predicted
by the synthesizer or on the groundtruth mel spectrograms:

$\min_{\mathbf{w}_{\mathcal{V}}}L_{\mathcal{V}}\left(\mathbf{u}_{ij},\mathcal{V}\left(\mathbf{x}_{ij};\mathbf{w}_{\mathcal{V}}\right)\right)\text{
or
}\min_{\mathbf{w}_{\mathcal{V}}}L_{\mathcal{V}}\left(\mathbf{u}_{ij},\mathcal{V}\left(\hat{\mathbf{x}}_{ij};\mathbf{w}_{\mathcal{V}}\right)\right)$
(8)

where LV is a loss function in the time-domain. In the second case, a pre-
trained synthesizer model is needed.

If the definition of the objective function was quite simple for both the
synthesizer and the vocoder, unfortunately this is not the case for the
speaker encoder. The encoder does not have labels to be trained on because its
task is only to create the space of characteristics necessary to create the
embedding vectors. The Generalized End-to-End (GE2E) [10] loss brings a
solution to this problem and it allows the training of the speaker encoder
independently. Consequently, we can define the following objective function:

$\min_{\mathbf{w}_{\mathcal{E}}}L_{\mathcal{G}}(\mathbf{S};w_{\mathcal{E}})=\sum_{j,i}L\left(\mathbf{e}_{ji}\right)$
(9)

where S represents a similarity matrix and LG is the GE2E loss function.

###  2.2 Speaker Encoder

Fig. 2: Speaker encoder model architecture. Input is composed of a time
sequence of dimension 40. The last linear layer takes the hidden state of the
last GRU layer as input.

The speaker encoder must be able to produce an embedding vector that
meaningfully represents speaker characteristics in the transformed space
starting from a target speaker’s utterance. Furthermore, the model should
identify these characteristics using a short speech signal, regardless of its
phonetic content and background noise. This can be achieved by training a
neural network model on a text-independent speaker verification task that
tries to optimize the GE2E loss so that embeddings of utterances from the same
speaker have high cosine similarity, while those of utterances from different
speakers are far apart in the embedding space.

The network maps a sequence of mel spectrogram frames to a fixed-dimensional
embedding vector, known as d-vector [11, 12]. Input mel spectrograms are fed
to a network consisting of one Conv1D [13] layer of 512 units followed by a
stack of 3 GRU [14] layers of 512 units, each followed by a linear projection
of 256 dimension. Following [9], the final embedding dimension is 256 and it
is created by L2-normalizing the output of the top layer at the final frame.
This is shown in Figure 2. We noticed that this architecture was the best
among the various tried and tested, as we will see in Section 3.

During the training phase, all the utterances are split into partial
utterances that are 1.6 seconds long (160 frames). Also at inference time, the
input utterance is split into segments of 1.6 seconds with 50% overlap and the
model processes each segment individually. Following [9, 10], the final
utterance-wise d-vector is generated by L2 normalizing the window-wise
d-vectors and taking the element-wise average.

###  2.3 Synthesizer and Vocoder

The synthesizer component of the system is a sequence-to-sequence model with
attention [1, 3] which is trained on pairs of text derived token sequences and
audio derived mel spectrogram sequences. Furthermore, the network is trained
in a transfer learning configuration (see Paragraph 2.4), using an
independently-trained speaker encoder to extract embedding vectors useful to
condition the outcomes of this component. In view of reproducibility, the
adopted vocoder component of the system is a Pytorch github
implementation111https://github.com/fatchord/WaveRNN of the neural vocoder
WaveRNN [15]. This model is not directly conditioned on the output of the
speaker encoder but just on the input mel spectrogram. The multi-speaker
vocoder is simply trained by using data from many speakers (see Section 3).

###  2.4 Transfer Learning Modality

The conditioning of the synthesizer via speaker encoder is the fundamental
part that makes the system multi-speaker: the embedding vectors computed by
the speaker encoder allow the conditioning of the mel spectrograms generated
by the synthesizer so that they can incorporate the new speaker voice. In [9],
the embedding vectors are speaker embeddings obtained by Equation 2. We used
the utterance embeddings computed by Equation 1. In fact, at inference time
only one utterance of the target speaker is fed to the speaker encoder which
therefore produces a single utterance-level d-vector. Thus, in this case, it
is not possible to create an embedding at the speaker level since the average
operation cannot be applied. This implies that only utterance embeddings can
be used during the inference phase. In addition, an average mechanism could
cause some loss in terms of accuracy. This is due to larger variations in
pitch and voice quality often occurring in utterances of the same speaker
while utterances have lower intra-variation. Following [9], the embedding
vectors computed by the speaker encoder are concatenated only with the
synthesizer encoder output in order to condition the synthesis. However, we
experimented with a new concatenation technique: first we passed the embedding
through a single linear layer and then we applied the concatenation between
the output of this layer and the synthesizer encoder one. The goal was to
exploit the weights of the linear layer to make the embedding vector more
meaningful, since the layer was trained together with the synthesizer. We
noticed that this method achieved good convergence of training and was about
75% times faster than the former vector concatenation.

##  3 Experiments and Results

We used different publicly available datasets to train and evaluate the
components of the system. For the speaker encoder, different neural network
architectures were tested. Each of them was trained using a combination of
three public sets: LibriTTS [16] train-other and dev-other; VoxCeleb [17] dev
and VoxCeleb2 [18] dev. In this way, we obtained a number of speakers equal to
8,381 and a number of utterances equal to 1,419,192, not necessarily all clean
and noiseless. Furthermore, transcripts were not required. The models were
trained using Adam [19] as optimizer with an initial learning rate equal to
0.001. Moreover, we experimented with different learning rate decay
strategies.

During the evaluation phase, we used a combination of the corresponding test
sets of the training ones, obtaining a number of speakers equal to 191 and a
number of utterances equal to 45,132. Both training and test sets have been
sampled at 16 kHz and input mel spectrograms were computed from 25ms STFT
analysis windows with a 10ms step and passed through a 40-channel mel-scale
filterbank.

We separately trained the synthesizer and the vocoder using the same training
set given by the combination of the two “clean” sets of LibriTTS, obtaining a
number of speakers equal to 1,151, a number of utterances equal to 149,736 and
a total number of hours equal to 245,14 of 22.05 kHz audio. We trained the
synthesizer using the L1 loss [20] and Adam as optimizer. Moreover, the input
texts were converted into phoneme sequences and target mel spectrogram
features are computed on 50 ms signal windows, shifted by 12.5 ms and passed
through an 80-channel mel-scale filterbank. The vocoder was trained using
groundtruth waveforms rather than the synthesizer outputs.

###  3.1 Baseline System

We choose as baseline for our work the Corentin Jemine’s real-time voice
cloning system [21], a public re-implementation of the Google system [9]
available on github222https://github.com/CorentinJ/Real-Time-Voice-Cloning.
This system is composed out of three components: a recurrent speaker encoder
consisting of 3 LSTM [22] layers and a final linear layer, each of which has
256 units; a sequence-to-sequence with attention synthesizer based on [1] and
WaveRNN [15] as vocoder.

###  3.2 Speaker Encoder: Proposed System

To evaluate all the speaker encoder models and choose the best one, the
Speaker Verification Equal Error Rate (SV-EER) was estimated by pairing each
test utterance with each enrollment speaker. The models implemented are:

  * •

_rec_conv network_ : 5 Conv1D layers, 1 GRU layer and a final linear layer;

  * •

_rec_conv_2 network_ : 3 Conv1D layers, 2 GRU layers each followed by a linear
projection layer;

  * •

_gru network_ : 3 GRU layers each followed by a linear projection layer;

  * •

_advanced_gru network_ : 1 Conv1D layer and 3 GRU layers each followed by a
linear projection layer (Figure 2);

  * •

_lstm network_ : 1 Conv1D layer and 3 LSTM [22] layers each followed by a
linear projection layer.

All layers have 512 units except the linear ones which have 256. Moreover,
dropout rate of 0.2 was used after all the layers except before the first and
after the last. All the models were trained using a batch size of 64 speakers
and 10 utterances for each speaker. The results obtained are shown in Table 1.

Table 1: Speaker Verification Equal Error Rates.

Name Step Time Train Loss SV-EER LR Decay rec_conv 0.33s 0.36 0.073 Reduce on
Plateau rec_conv_2 0.45s 0.49 0.075 Reduce on Plateau gru 1,45s 0.33 0.054
Every 100,000 step advanced_gru 0.86s 0.14 0.040 Exponential lstm 1.08s 0.17
0.052 Exponential

We designed the advanced gru network trying to combine the advantages of
convolutional and gru networks. In fact, looking at the table, this
architecture was much faster than the gru network during training, and
obtained the best SV-EER on the test set. Figure 3 illustrates the projection
in a two-dimensional space of the utterance embeddings computed by the
advanced gru network on the basis of 6 utterances extracted from 12 speakers
of the test set. In Figure 4, the 12 speakers are 6 men and 6 women. The
projections were made using UMAP [23]. Both the figures show that the model
has created a space of internal features that is robust regarding the
speakers, creating well-formed clusters of speakers based on their utterances
and nicely separating male speakers from female ones.

The SV-EER obtained on the test set from the speaker encoder model of the
proposed system is 0.040 vs the baseline one which is 0.049.

Fig. 3: Advanced Gru Network test utterance embeddings projection. Fig. 4:
Advanced Gru Network six utterances for six male vs six utterances for six
female taken from the test set.

###  3.3 Similarity Evaluation

To assess how similar the waveforms generated by the system were from the
original ones, we transformed the audio signals produced into utterance
embeddings (using the speaker encoder advanced gru network) and then projected
them in a two-dimensional space together with the utterance embeddings
computed on the basis of the groundtruth audio. As test speakers, we randomly
choose eight target speakers: four speakers (two male and two female) were
extracted from the test-set-clean of LibriTTS [16], three (two male and one
female) from VCTK [24] and finally a female proprietary voice. For each
speaker we randomly extracted 10 utterances and compared them with the
utterances generated by the system calculating the cosine similarity. The
speakers averaged values of cosine similarity between the generated and
groundtruth utterance embeddings range from 0.56 to 0.76. Figure 5 shows that
synthesized utterances tend to lie close to real speech from the same speaker
in the embedding space.

Fig. 5: Groundtruth utterance embeddings vs the corresponding generated ones
of the 8 speakers chosen for testing.

###  3.4 Subjective Evaluation

Finally, we evaluated how the generated utterances were, subjectively
speaking, similar in terms of speech timbre to the original ones. To do this,
we gathered Mean Similarity Scores (MSS) based on a 5 points mean opinion
score scale, where 1 stands for “very different” and 5 for “very similar”. Ten
utterances of the proprietary female voice were cloned using both the proposed
and the baseline system and then 12 subjects, most of them TTS experts, were
asked to listen to the 20 samples, randomly mixed, and rate them. Participants
were also provided with an original utterance as reference. The question asked
was: “How do you rate the similarity of these samples with respect to the
reference audio? Try to focus on vocal timbre and not on content, intonation
or acoustic quality of the audio”. The results obtained are shown in Table 2.
Although not conclusive, this experiment highlights a subjective evidence of
the goodness of the proposed approach, despite the significant variance of
both systems: this is largely due to the low number of test participants.

Table 2: MSS of the baseline and the proposed systems.

System MSS baseline $2.59\pm 1.03$ proposed $3.17\pm 0.97$

##  4 Conclusions

In this work, our goal was to build a Voice Cloning system which could
generate natural speech for a variety of target speakers in a data efficient
manner. Our system combines an independently trained speaker encoder network
with a sequence-to-sequence with attention architecture and a neural vocoder
model. Using a transfer learning technique from a speaker-discriminative
encoder model based on utterance embeddings rather than speaker embeddings,
the synthesizer and the vocoder are able to generate good quality speech also
for speakers not observed before. Despite the experiments showed a reasonable
similarity with real speech and improvements over the baseline, the proposed
system does not fully reach human-level naturalness in contrast to the single
speaker results from [1]. Additionally, the system is not able to reproduce
the speaker prosody of the target audio. These are consequences of the
additional difficulty of generating speech for a variety of speakers given
significantly less data per speaker unlike when training a model on a single
speaker.

##  5 Acknowledgements

The authors thank Roberto Esposito, Corentin Jemine, Quan Wang, Ignacio Lopez
Moreno, Skjalg Lepsøy, Alessandro Garbo and Jürgen Van de Walle for their
helpful discussions and feedback.

## References

  * [1] J. Shen, R. Pang, R. J. Weiss, M. Schuster, N. Jaitly, Z. Yang, Z. Chen, Y. Zhang, Y. Wang, R. Skerrv-Ryan, R. A. Saurous, Y. Agiomvrgiannakis, and Y. Wu,  “Natural tts synthesis by conditioning wavenet on mel spectrogram predictions,”  in 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2018, pp. 4779–4783.
  * [2] Aäron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew W. Senior, and Koray Kavukcuoglu,  “Wavenet: A generative model for raw audio,”  CoRR, vol. abs/1609.03499, 2016.
  * [3] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio,  “Neural machine translation by jointly learning to align and translate,”  CoRR, vol. abs/1409.0473, 2015.
  * [4] Andrew Gibiansky, Sercan Arik, Gregory Diamos, John Miller, Kainan Peng, Wei Ping, Jonathan Raiman, and Yanqi Zhou,  “Deep voice 2: Multi-speaker neural text-to-speech,”  in Advances in Neural Information Processing Systems 30, I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, Eds., pp. 2962–2970. Curran Associates, Inc., 2017.
  * [5] Wei Ping, Kainan Peng, Andrew Gibiansky, Sercan O. Arik, Ajay Kannan, Sharan Narang, Jonathan Raiman, and John Miller,  “Deep voice 3: 2000-speaker neural text-to-speech,”  in International Conference on Learning Representations, 2018.
  * [6] V. Panayotov, G. Chen, D. Povey, and S. Khudanpur,  “Librispeech: An asr corpus based on public domain audio books,”  in 2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2015, pp. 5206–5210.
  * [7] Yaniv Taigman, Lior Wolf, Adam Polyak, and Eliya Nachmani,  “Voiceloop: Voice fitting and synthesis via a phonological loop,”  in International Conference on Learning Representations, 2018.
  * [8] Eliya Nachmani, Adam Polyak, Yaniv Taigman, and Lior Wolf,  “Fitting new speakers based on a short untranscribed sample,”  CoRR, vol. abs/1802.06984, 2018.
  * [9] Ye Jia, Yu Zhang, Ron J. Weiss, Quan Wang, Jonathan Shen, Fei Ren, Zhifeng Chen, Patrick Nguyen, Ruoming Pang, Ignacio Lopez-Moreno, and Yonghui Wu,  “Transfer learning from speaker verification to multispeaker text-to-speech synthesis,”  CoRR, vol. abs/1806.04558, 2018.
  * [10] Lipeng Wan, Qi shan Wang, Alan Papir, and Ignacio Lopez-Moreno,  “Generalized end-to-end loss for speaker verification,”  2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 4879–4883, 2018.
  * [11] Georg Heigold, Ignacio Moreno, Samy Bengio, and Noam Shazeer,  “End-to-end text-dependent speaker verification,”  2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 5115–5119, 2016.
  * [12] Ehsan Variani, Xin Lei, Erik McDermott, Ignacio Lopez Moreno, and Javier Gonzalez-Dominguez,  “Deep neural networks for small footprint text-dependent speaker verification,”  in Proc. ICASSP, 2014.
  * [13] Serkan Kiranyaz, Onur Avci, Osama Abdeljaber, Turker Ince, Moncef Gabbouj, and Daniel J. Inman,  “1d convolutional neural networks and applications: A survey,”  ArXiv, vol. abs/1905.03554, 2019.
  * [14] Junyoung Chung, Caglar Gulcehre, Kyunghyun Cho, and Yoshua Bengio,  “Empirical evaluation of gated recurrent neural networks on sequence modeling,”  in NIPS 2014 Workshop on Deep Learning, December 2014, 2014.
  * [15] Nal Kalchbrenner, Erich Elsen, Karen Simonyan, Seb Noury, Norman Casagrande, Edward Lockhart, Florian Stimberg, Aäron van den Oord, Sander Dieleman, and Koray Kavukcuoglu,  “Efficient neural audio synthesis,”  in ICML, 2018.
  * [16] Heiga Zen, Viet Dang, Rob Clark, Yu Zhang, Ron J. Weiss, Ye Jia, Zhifeng Chen, and Yonghui Wu,  “Libritts: A corpus derived from librispeech for text-to-speech,”  in INTERSPEECH, 2019.
  * [17] Arsha Nagrani, Joon Son Chung, and Andrew Zisserman,  “Voxceleb: A large-scale speaker identification dataset,”  in INTERSPEECH, 2017.
  * [18] Joon Son Chung, Arsha Nagrani, and Andrew Zisserman,  “Voxceleb2: Deep speaker recognition,”  in INTERSPEECH, 2018.
  * [19] Diederik Kingma and Jimmy Ba,  “Adam: A method for stochastic optimization,”  International Conference on Learning Representations, 12 2014.
  * [20] Katarzyna Janocha and Wojciech Czarnecki,  “On loss functions for deep neural networks in classification,”  ArXiv, vol. abs/1702.05659, 2017.
  * [21] Corentin Jemine,  “Master thesis: Automatic multispeaker voice cloning,” 2019,  Unpublished master’s thesis, Université de Liège, Liège, Belgique.
  * [22] Klaus Greff, Rupesh K. Srivastava, Jan Koutnik, Bas R. Steunebrink, and Jurgen Schmidhuber,  “Lstm: A search space odyssey,”  IEEE Transactions on Neural Networks and Learning Systems, vol. 28, no. 10, pp. 2222–2232, Oct 2017.
  * [23] Leland McInnes and John Healy,  “Umap: Uniform manifold approximation and projection for dimension reduction,”  ArXiv, vol. abs/1802.03426, 2018.
  * [24] Christophe Veaux, Junichi Yamagishi, and Kirsten MacDonald,  “Cstr vctk corpus: English multi-speaker corpus for cstr voice cloning toolkit,” 2018.