# Auto-Encoding Knowledge Graph for Unsupervised
Medical Report Generation

Fenglin Liu1,2, Chenyu You4, Xian Wu5, Shen Ge5, Sheng Wang3, Xu
Sun111footnotemark: 1
1MOE Key Laboratory of Computational Linguistics, School of EECS, Peking
University
2School of ECE, Peking University
3Paul G. Allen School of Computer Science and Engineering, University of
Washington
4Department of Electrical Engineering, Yale University 5Tencent
{fenglinliu98<EMAIL_ADDRESS><EMAIL_ADDRESS>
{kevinxwu<EMAIL_ADDRESS><EMAIL_ADDRESS>
Corresponding authors.

###### Abstract

Medical report generation, which aims to automatically generate a long and
coherent report of a given medical image, has been receiving growing research
interests. Existing approaches mainly adopt a supervised manner and heavily
rely on coupled image-report pairs. However, in the medical domain, building a
large-scale image-report paired dataset is both time-consuming and expensive.
To relax the dependency on paired data, we propose an unsupervised model
Knowledge Graph Auto-Encoder (KGAE) which accepts independent sets of images
and reports in training. KGAE consists of a pre-constructed knowledge graph, a
knowledge-driven encoder and a knowledge-driven decoder. The knowledge graph
works as the shared latent space to bridge the visual and textual domains; The
knowledge-driven encoder projects medical images and reports to the
corresponding coordinates in this latent space and the knowledge-driven
decoder generates a medical report given a coordinate in this space. Since the
knowledge-driven encoder and decoder can be trained with independent sets of
images and reports, KGAE is unsupervised. The experiments show that the
unsupervised KGAE generates desirable medical reports without using any image-
report training pairs. Moreover, KGAE can also work in both semi-supervised
and supervised settings, and accept paired images and reports in training. By
further fine-tuning with image-report pairs, KGAE consistently outperforms the
current state-of-the-art models on two datasets.

##  1 Introduction

Medical images, such as radiology and pathology images, and their
corresponding reports are widely used for clinical diagnosis and treatment [8,
11]. A medical report is usually a paragraph of multiple sentences which
describes both the normal and abnormal findings in the medical image. In
clinical practice, writing a report can be time-consuming and tedious for
experienced radiologists, and error-prone for inexperienced radiologists [4].
Therefore, given the large volume of medical images, automatically generating
reports can improve current clinical practice in diagnostic radiology and
assist radiologists in clinical decision-making [15, 25]. Specifically, it can
relieve radiologists from such heavy workload and alert radiologists of the
abnormalities to avoid misdiagnosis and missed diagnosis. Therefore, automatic
medical report generation attracts remarkable attention in both artificial
intelligence and clinical medicine.

Figure 1: Illustration of our Knowledge Graph Auto-Encoder, which consists of
a pre-constructed knowledge graph, a knowledge-driven encoder and a knowledge-
driven decoder. The Green and Red lines denote the data flow in the training
process and testing process of report generation, respectively.

Recently, inspired by the great success of neural machine translation [2, 37,
48, 49, 47], image captioning [43, 38, 45] and medical imaging analysis [46,
50, 51], the data-driven deep neural models, particularly those based on the
encoder-decoder frameworks [15, 16, 25, 22, 40, 44, 52, 53, 6], have achieved
great success in advancing the state-of-the-art of medical report generation.
However, these models are trained in a supervised learning manner and heavily
rely on labeled paired image-report datasets [9, 17], which are not easy to
acquire in the real world. Specifically, the medical-related data can only be
manually labeled by professional radiologists, and also involves privacy
issues. Therefore, the medical report generation datasets are particularly
labor-intensive and expensive to obtain. As a result, the scales of existing
widely-used datasets for medical report generation models [16, 25, 22, 23],
i.e., MIMIC-CXR (0.22M samples) [17] and IU X-ray (4K samples) [9], are
relatively small compared to image recognition datasets, e.g., ImageNet (14M
samples) [10], and image captioning datasets, e.g., Conceptual Captions (3.3M
samples) [35]. In addition, the MIMIC-CXR and IU X-Ray datasets only include
Chest X-Ray images, for other types of medical images (MRI, Dermoscopy,
Retinal, etc.) of other body parts (brain, skin, eye, etc.), the image-report
pairs could be much less or even unavailable. Therefore, to relax the reliance
on the paired data sets, making use of all available data, like independent
image or report sets, is becoming increasingly important.

In this paper, we propose an unsupervised model Knowledge Graph Auto-Encoder
(KGAE), which utilizes independent sets of images and reports in training (the
image and report set are separate and have no overlap). KGAE consists of a
pre-constructed knowledge graph, a knowledge-driven encoder and a knowledge-
driven decoder. As shown in Figure 1, the knowledge graph works as the shared
latent space of images and reports. The knowledge-driven encoder can take
either the image $I$ or the report $R$ as queries and project them to
corresponding coordinates $\mathcal{G}_{I}$ and $\mathcal{G}_{R}$ in the
latent space. In this manner, since $\mathcal{G}_{I}$ and $\mathcal{G}_{R}$
share the same latent space, we can use positions in latent space to measure
the relationship between images and reports which narrows the gap between
visual and textual domains. In brief, to bridge the gap between vision and
language domains without training on the pairs of images and reports, we adopt
the knowledge graph to create a latent space and propose a knowledge-driven
encoder, which includes a common mapping function to project images and
reports to the same latent space. As a result, our encoder can extract the
image and report knowledge representations, i.e., the knowledge related to the
image and report, they (image, report knowledge) share the common latent
space, which allows our model to bridge the gap between vision and language
domains without the training on the pairs of image and report. Next, we
introduce the knowledge-driven decoder to exploit $\mathcal{G}_{I}$ and
$\mathcal{G}_{R}$ to generate the report. In the training stage, we estimate
the parameters of the decoder by reconstructing the input report $R$ based on
$\mathcal{G}_{R}$, i.e., $R\to\mathcal{G}_{R}\to R$ auto-encoding pipeline; In
the prediction stage, we directly input $\mathcal{G}_{I}$ into the trained
decoder to generate the report. In this way, our approach can produce
desirable reports without any labeled image-report pairs.

Overall, the contributions of this paper are as follows:

  * •

In this paper, we make the first attempt to conduct unsupervised medical
report generation where the image-report pairs are not available. To this end,
we propose the Knowledge Graph Auto-Encoder (KGAE). By leveraging a pre-
constructed knowledge graph, we introduce the knowledge-driven encoder and
decoder which are trained with independent sets of images and reports.
According to the experimental results, the unsupervised KGAE can even
outperform several supervised approaches.

  * •

In addition to the unsupervised mode, KGAE can also be applied in a semi-
supervised or supervised manner. Under the semi-supervised setting, by using
only 60% of paired dataset, KGAE is able to achieve competitive results with
current state-of-art models; Under the supervised setting, by training on
fully paired datasets as in existing works, KGAE can set new state-of-the-art
performances on the IU X-ray and MIMIC-CXR, respectively.

  * •

The analysis, both quantitative and qualitative, as well as a human evaluation
conducted by professional radiologists, further proves the effectiveness of
our approach.

##  2 Related Works

Medical report generation aims to generate a relatively long paragraph to
describe a given medical image. It is similar to the image captioning task [5,
28], which aims to generate a sentence to describe a given image. In image
captioning, the encoder-decoder framework [38], where the encoder [20, 12]
computes visual representations for the image and the decoder [13, 37]
generates a target sentence based on the visual representations, has achieved
great success [1, 30, 34, 43]. However, instead of only generating one single
sentence, medical report generation aims to generate a long paragraph
including multiple structured sentences that describe both the normal and
abnormal parts [25, 15]. To this end, given the success of encoder-decoder
framework on image captioning, most existing medical report generation models
attempt to exploit the hierarchical LSTM (HLSTM) [15, 16, 19] or Transformer
[6] to generate an accurate, long and coherent report. However, existing
models require the paired image-report datasets, which are time-consuming and
expensive to collect. In this paper, we propose the unsupervised model
Knowledge Graph Auto-Encoder (KGAE) which doesn’t need paired images and
reports.

Note that although the knowledge graph has been integrated in existing medical
report generation models [53, 22], these approaches are supervised and require
paired images and reports. Thus, their objectives and motivations of using
knowledge graph are different from our work. In detail, existing knowledge-
graph based medical report generation methods aim to adopt the knowledge graph
to boost the performance of supervised models. However, in our work, we aim to
generate a medical report without using any coupled image-report training
pairs, i.e., unsupervised medical report generation. A key challenge of
unsupervised medical report generation is to bridge the gap between vision and
language domains. To this end, we adopt the knowledge graph to create a latent
space and propose a knowledge-driven encoder to project image and report to
the same latent space.

##  3 Approach

We first formulate the conventional supervised medical report generation
problems; Then, we describe the proposed Knowledge Graph Auto-Encoder for
unsupervised medical report generation in detail.

###  3.1 Conventional Supervised Medical Report Generation Models

Given a medical image $I$, the goal is to generate a descriptive report $R$.
Most models [15, 16, 6] normally include an image encoder and a report
decoder, which can be formulated as:

$\text{Image Encoder}:I\to I^{\prime};\ \ \text{Report Decoder}:I^{\prime}\to
R,$ (1)

where $I^{\prime}\in\mathbb{R}^{N_{I}\times d}$ denotes the image embeddings
extracted by the image encoder, e.g., ResNet-50 [12]. Then, $I^{\prime}$ is
used to guide the generation of the target report $R$ in the report decoder,
e.g., LSTM [13] and Transformer [37]. During training, given the ground truth
report for the input image, we can train the encoder-decoder model by
minimizing a supervised training loss, e.g., cross-entropy loss. However,
paired image-report ($I$-$R$) datasets are particularly labor-intensive and
expensive to obtain. Therefore, in this paper, we aim to relax the reliance on
the paired datasets and make use of the independent sets of images and reports
instead.

###  3.2 Knowledge Graph Auto-Encoder

As shown in Figure 1, the proposed KGAE includes a knowledge graph, a
knowledge-driven encoder and a knowledge-driven decoder, which will be
described in detail in the following sections.

Knowledge Graph The motivation of using the knowledge graph is that writing
medical reports usually requires particular domain knowledge [11, 22].
Therefore, we can utilize the knowledge graph, which models the domain-
specific knowledge structure, to serve as a bridge to correlate the knowledge
representations of the visual and textual domains. In particular, we construct
an off-the-shelf global medical knowledge graph $\mathcal{G}=(V,E)$ covering
the common abnormalities and normalities, where
$V=\\{v_{i}\\}_{i=1:N_{\text{KG}}}\in\mathbb{R}^{N_{\text{KG}}\times d}$ is a
set of nodes and $E=\\{e_{i,j}\\}_{i,j=1:N_{\text{KG}}}$ is a set of edges. In
detail, based on the report corpus, i.e., MIMIC-CXR [17], we consider the
$N_{\text{KG}}$ frequent clinical abnormalities (e.g., "enlarged heart size",
"pleural effusion", and "bibasilar consolidation") and normalities (e.g.,
"heart size is normal" and "lungs are clear") as nodes. The edge weights are
calculated by the normalized co-occurrence of different nodes computed from
report corpus. Please note that this knowledge graph is built purely from the
training set of MIMIC-CXR, thus there is no label leakage. After that, the
knowledge graph is embedded by a graph embedding module, i.e., graph
convolution network [31, 24], which is implemented as follows:

$\footnotesize{v}_{i}^{\prime}={v}_{i}+\text{ReLU}(\sum_{j=1}^{N_{\text{KG}}}{e}_{i,j}{W}_{v}{v}_{j}),$
(2)

where $\text{ReLU}(\cdot)$ represents the ReLU activation function and
$W_{v}\in\mathbb{R}^{d\times d}$ is the learnable matrix. As a result, we can
acquire a set of node embeddings
$V^{\prime}=\\{v^{\prime}_{1},v^{\prime}_{2},\dots,v^{\prime}_{N_{\text{KG}}}\\}\in\mathbb{R}^{N_{\text{KG}}\times
d}$. It is worth noting that more complex graph structures could be
constructed by using more large-scale medical textbooks. Therefore, our model
is not limited to the currently constructed graph and could provide a good
basis for the research of unsupervised medical report generation for other
domains.

Knowledge-driven Encoder The knowledge-driven encoder (KE) is designed to
utilize the knowledge graph $\mathcal{G}$ to extract knowledge representations
($\mathcal{G}_{I}\in\mathbb{R}^{N_{I}\times d}$ and
$\mathcal{G}_{R}\in\mathbb{R}^{N_{R}\times d}$) of both image $I$ and report
$R$ to bridge the vision and the language domains, which can be formulated as:

$\displaystyle\mathcal{G}_{I}=\text{KE}_{I}(I,\mathcal{G});\ \
\mathcal{G}_{R}=\text{KE}_{R}(R,\mathcal{G}).$ (3)

In implementations, as shown in Figure 1, we first adopt the ResNet-50 [12]
and the Transformer [37] as the image embedding module and report embedding
module to embed the image $I$ and report $R$, acquiring the image embeddings
$I^{\prime}\in\mathbb{R}^{N_{I}\times d}$ and report embeddings
$R^{\prime}\in\mathbb{R}^{N_{R}\times d}$, respectively. Next, to extract the
knowledge representations from knowledge graph $\mathcal{G}$, we adopt the
attention mechanism [37] to implement the KE. The motivation stems from that
the attention mechanism can compute the association weights between different
features and allows probabilistic many-to-many relations instead of monotonic
relations, as in [37, 43, 29]. As a result, we take $I^{\prime}$ and
$R^{\prime}$ as the queries, and take the knowledge graph $\mathcal{G}$ (i.e.,
$V^{\prime}=\\{v^{\prime}_{1},v^{\prime}_{2},\dots,v^{\prime}_{N_{\text{KG}}}\\}\in\mathbb{R}^{N_{\text{KG}}\times
d}$) as the lookup matrix:

$\displaystyle\mathcal{G}_{I}=\text{KE}_{I}(I,\mathcal{G})=$
$\displaystyle\mathcal{F}(\text{Attention}_{I}(I^{\prime},V^{\prime}));\ \
\mathcal{G}_{R}=\text{KE}_{R}(R,\mathcal{G})=\mathcal{F}(\text{Attention}_{R}(R^{\prime},V^{\prime}))$
(4)
$\displaystyle\text{Attention}(x,y)=\text{softmax}\left(\frac{xW_{\text{q}}\left(yW_{\text{k}}\right)^{\top}}{\sqrt{d}}\right)yW_{\text{v}}.$

where $x\in\mathbb{R}^{N_{x}\times d_{x}},y\in\mathbb{R}^{N_{y}\times d_{y}}$;
$W_{\text{q}}\in\mathbb{R}^{d_{x}\times d_{x}}$,
$W_{\text{k}}\in\mathbb{R}^{d_{y}\times d_{y}}$ and
$W_{\text{v}}\in\mathbb{R}^{d_{y}\times d_{y}}$ ($d_{x}=d_{y}$) are learnable
parameters. $\mathcal{F}$ is implemented as two fully-connected (FC) layers
with a ReLU in between, i.e., FC-ReLU-FC. The resulted
$\mathcal{G}_{I}\in\mathbb{R}^{N_{I}\times d}$ and
$\mathcal{G}_{R}\in\mathbb{R}^{N_{R}\times d}$ turn out to be a set of
attended (i.e., extracted) knowledge related to the image and report. The
knowledge representations $\mathcal{G}_{I}$ and $\mathcal{G}_{R}$ share the
common latent space which bridges the vision and the language domains. Note
that the $\mathcal{F}$ used in $\text{KE}_{I}$ and $\text{KE}_{R}$ shares the
same parameters and is thus introduced to further boost the bridging
capabilities of different modalities. It means that we adopt two KEs, i.e.,
$\text{KE}_{I}$ and $\text{KE}_{R}$, and each KE includes an attention model
as well as a common mapping function $\mathcal{F}$. The only shared weights of
$\text{KE}_{I}$ and $\text{KE}_{R}$ are the parameters of $\mathcal{F}$, and
the parameters of attention models are independent.

Knowledge-driven Decoder The decoder is designed to generate the reports based
on the graph representations $\mathcal{G}_{I}$ or $\mathcal{G}_{R}$. For
clarity, we use $\mathcal{G}_{k}$ to represent the $\mathcal{G}_{R}$ and
$\mathcal{G}_{I}$ during the training and testing stages, respectively. In
implementations, since medical report generation requires generating a long
paragraph, we choose the (three-layer) Transformer [37] as the basic module of
our decoder and incorporate the proposed Knowledge-driven Attention (KA) to
effectively model the long sequences.

At each decoding step $t$, the decoder takes the embedding of current input
word $x_{t}=w_{t}+e_{t}\in\mathbb{R}^{2d}$ as input, where $w_{t}$ and $e_{t}$
denote the word embedding and fixed position embedding, respectively, and
generate each word $r_{t}$ in report $R=\\{r_{1},r_{2},\dots,r_{T}\\}$, which
can be defined as follows:

$\footnotesize\begin{array}[]{ll}&h_{t}=\text{Attention}(x_{t},x_{1:t})\\\
&h^{\prime}_{t}=\text{KA}(h_{t},\mathcal{G}_{k},B)\end{array},\ \text{where\
}\mathcal{G}_{k}=\left\\{\begin{array}[]{ll}\mathcal{G}_{R},&\footnotesize\text{Training}\\\
\mathcal{G}_{I},&\footnotesize\text{Testing}\end{array}\right.;\ r_{t}\sim
p_{t}=\text{softmax}(\text{FFN}(h^{\prime}_{t})\text{W}_{p}),$ (5)

where FFN stands for Feed-Forward Network in the original Transformer [37];
$\text{W}_{p}\in\mathbb{R}^{2d\times|D|}$ is the learnable parameter ($|D|$
denotes the vocabulary size); $B$ represents the knowledge bank in the KA.

In particular, the knowledge-driven attention (KA) is inspired by the success
of many works that attempt to augment a working memory into the network, e.g.,
memory network. Memory network preserves a dynamic knowledge base for
subsequent inference [36, 21, 41, 42, 39, 45, 32]. For example, [36, 21] and
[41, 42] proposed to preserve the context of a document and visual knowledge
to solve language modeling and visual question answering, respectively. In
this paper, since the knowledge representations $\mathcal{G}_{k}$ are
extracted from the off-the-shelf global knowledge graph $\mathcal{G}$, we
incorporate a knowledge memory mechanism to distill and preserve the fine-
grained medical knowledge related to report generation task, which encourages
our decoder to generate more accurate and desirable reports. Specifically, we
introduce the knowledge bank
$B=\\{b_{1},b_{2},\dots,b_{N_{\text{B}}}\\}\in\mathbb{R}^{N_{\text{B}}\times
d}$, where $N_{\text{B}}$ stands for the total number of the knowledge
corresponding to report generation. As a result, given the graph
representations $\mathcal{G}_{k}\in\mathbb{R}^{N_{k}\times d}$, the knowledge
memory mechanism is defined as follows:

$\displaystyle{B}_{k}$
$\displaystyle=\text{softmax}\left(\mathcal{G}_{k}B^{\top}\right)B.$ (6)

Through above operation, each feature/vector in the
$\mathcal{G}_{k}\in\mathbb{R}^{N_{k}\times d}$ can distill the fine-grained
knowledge preserved in $B$ to acquire ${B}_{k}\in\mathbb{R}^{N_{k}\times d}$
for accurate report generation.

Based on the above mechanism, the knowledge-driven attention in Eq. (5) is
defined as:

$h^{\prime}_{t}=\text{KA}(h_{t},\mathcal{G}_{k},B)=\text{Attention}(h_{t},[\mathcal{G}_{k};B_{k}])=\text{Attention}\left(h_{t},\left[\mathcal{G}_{k};\text{softmax}\left(\mathcal{G}_{k}B^{\top}\right)B\right]\right).$
(7)

where $[\cdot;\cdot]$ denotes the concatenation operation. It is worth noting
that these operations are all differentiable, thus the bank $B$ can be learned
in an end-to-end fashion.

In our subsequent analysis, we will show that the introduced knowledge memory
mechanism indeed distills and preserves the desired medical knowledge, and
thus boost the generation of reports.

###  3.3 Implementation Details

Unsupervised Training Details To train our KGAE in an unsupervised manner,
instead of using the paired image-report dataset in the conventional
supervised model, we only require an image set CheXpert [14], which includes
224,316 X-ray images, and a separate report corpus MIMIC-CXR [17] \+ IU X-ray
[9], which includes 222,758 + 2,770 = 225,528 reports111There are no paired
image-report samples between CheXpert and MIMIC-CXR+IU X-ray..

In detail, to train our knowledge-driven encoder $\text{KE}_{I}$ and
$\text{KE}_{R}$ (see Eq. (4)), we feed the $\mathcal{G}_{I}$ and
$\mathcal{G}_{R}$ into a common multi-label classification network [53, 15]
trained with binary cross entropy loss for 14 common radiographic observations
classification222Atelectasis, Cardiomegaly, Consolidation, Edema, Enlarged
Cardiomediastinum, Fracture, Lung Lesion, Lung Opacity, No Finding, Pleural
Effusion, Pleural Other, Pneumonia, Pneumothorax, Support Devices.. In this
way, our encoder can extract the knowledge representations $\mathcal{G}_{I}$
and $\mathcal{G}_{R}$ of both image and report in a common latent space,
effectively bridging the vision and the language domains. To train the
knowledge-driven decoder, as well as the knowledge bank $B$ (see Eq. (5)),
since there are no coupled image-report pairs, we propose to reconstruct the
report $R$ based on the $\mathcal{G}_{R}$. Therefore, through Eq. (5), taking
the input report $R=\\{r_{1},r_{2},\dots,r_{T}\\}$ as the ground truth report,
we can train our approach by minimizing the cross-entropy loss:

$\footnotesize L_{\text{CE}}=-\sum_{t=1}^{T}\log\left(p\left(r_{t}\mid
r_{1:t-1}\right)\right).$ (8)

In this way, we can train our decoder in the $R\to\mathcal{G}_{R}\to R$ auto-
encoding pipeline.

During testing, we first adopt the knowledge-driven encoder to extract the
knowledge representations $\mathcal{G}_{I}$ of the test image (see Eq. (4)).
Then, we directly feed $\mathcal{G}_{I}$ into the decoder to generate final
report in the $I\to\mathcal{G}_{I}\to R$ pipeline (see Eq. (5)). In this way,
our approach can relax the reliance on the image-report pairs. In our
following experiments, we validate the effectiveness of our approach, which
even outperforms some supervised approaches.

Semi-Supervised and Supervised Training Details To further validate the
effectiveness of our approach, we fine-tune the unsupervised KGAE using
partial and full image-report pairs to acquire the KGAE-Semi(-Supervised) and
KGAE-Supervised, respectively, where the former can evaluate the performance
of our approach under limited labeled pairs for training and the latter can
compare the performance of KGAE with state-of-the-art supervised approaches.
In the (semi-)supervised setting, given the image-report pairs, i.e., $I$-$R$,
we first incorporate the original visual information into the knowledge
representation $\mathcal{G}_{I}$, and then train our KGAE by generating the
ground truth report in the $I\to\mathcal{G}_{I}\to R$ pipeline and minimizing
the cross-entropy loss in Eq. (8). During testing, we also follow the
unsupervised setting to generate the final report in the
$I\to\mathcal{G}_{I}\to R$ pipeline.

##  4 Experiments

We first introduce the datasets, metrics and detailed settings used for
evaluation. Then, we present the evaluation of our approach under the
unsupervised, semi-supervised and supervised training settings.

###  4.1 Datasets, Metrics and Settings

Datasets In this paper, we adopt the test sets of IU X-ray [9] and MIMIC-CXR
[17] for evaluation. All protected health information (e.g., patient name and
date of birth) was de-identified. In particular, the IU X-ray [9] is a widely-
used public benchmark dataset for medical report generation and contains 7,470
chest X-ray images associated with 3,955 fully de-identified medical reports.
Each report is composed of impression, findings and indication sections, etc.
[22]. Following [25], our method also focuses on the findings section as it is
the most important component of reports. Then, following [16, 22, 25, 6], we
randomly select 70%-10%-20% image-report pairs of dataset to form the
training-validation-testing sets. The MIMIC-CXR [17] includes 377,110 chest
x-ray images associated with 227,835 reports. The dataset is officially split
into 368,960 images (222,758 reports) for training, 2,991 images (1,808
reports) for validation and 5,159 images (3,269 reports) for testing. It is
worth noting that we focus on the unsupervised medical report generation,
where the image-report pairs are not available, thus the image-report training
pairs of both the IU X-ray and MIMIC-CXR datasets are discarded and are not
used in our unsupervised training stage. Only the training reports of MIMIC-
CXR and IU X-ray, i.e., 222,758 + 2,770 = 225,528 reports, are used as the
independent report corpus to train our unsupervised model. Only under the
(semi-)supervised training setting, we will adopt the image-report training
pairs to train our approach.

Metrics To fairly compare with existing models [25, 6], we adopt the
evaluation toolkit [5] to calculate the widely-used natural language
generation metrics, i.e., BLEU [33], METEOR [3] and ROUGE-L [26], which
measure the match between the generated reports and ground truth reports, but
are not specialized for the abnormalities in the reports. Therefore, to
measure the accuracy of descriptions for clinical abnormalities, we further
report clinical efficacy metrics following the work of Chen et al. [6]. The
clinical efficacy metrics are calculated by comparing the generated reports
with ground truth reports in 14 different categories related to thoracic
diseases and support devices, producing the Precision, Recall and F1 scores.

Settings The size $d$ is set to 256. For the attention mechanism in Eq. (4)
and Eq. (5), we adopt the multi-head attention [37], the number of heads $n$
in multi-head attention is set to 8. The intermediate dimension in
$\mathcal{F}$, i.e., Eq. (4), is set to 1024. Based on the average performance
on the validation set, the $N_{\text{B}}$ in knowledge bank, i.e., Eq. (6), is
set to 10,000. The $N_{\text{KG}}$ in the knowledge graph is set to 200. In
our knowledge-driven encoder, the image embedding module adopts the ResNet-50
[12] pre-trained on ImageNet [10] and fine-tuned on CheXpert dataset [14] to
extract the image embeddings in the shape of 7 $\times$ 7 $\times$ 2048, which
will be projected to $d=256$, acquiring $I^{\prime}\in\mathbb{R}^{49\times
256}$, i.e., $N_{I}=49$; The report embedding module is implemented by the
Transformer [37] equipping with the self-attention mechanism provided in Lin
et al. [27]; $N_{R}=N_{I}=49$. In both unsupervised and (semi-)supervised
settings (see Section 3.3), the batch size is set to 16 and Adam optimizer
[18] with a learning rate of 1e-4 is used for parameter optimization. We also
use momentum of 0.8 and weight decay of 0.999. All re-implementations and our
experiments were run on 8 V100 GPUs.

###  4.2 Automatic Evaluation

We evaluate the performance of our approach under unsupervised, semi-
supervised and supervised settings. The results are shown in Table 1, Table 2
and Figure 2. We select several supervised methods, including a recently
state-of-the-art model R2Gen [6], for comparison. These models follow the
encoder-decoder architecture, trained on the full pairs of images and reports.

Unsupervised Setting As shown in Table 1 and Table 2, our unsupervised model
KGAE achieves competitive results with some supervised models in both IU X-ray
and MIMIC-CXR datasets, and even outperforms several supervised models.
Specifically, on the IU X-ray dataset, Table 1 shows that KGAE surpasses the
NIC [43], AdaAtt [30] and Att2in [34] in terms of all metrics, and the
Transformer [6] in terms of BLEU-1,2,3. On the MIMIC-CXR dataset, Table 2
shows that KGAE outperforms the $\mathcal{M}^{2}$ Trans. [7] in terms of
Precision, Recall and F1. The competitive results prove the effectiveness of
our approach in addressing the unsupervised medical report generation, and
thus can generate desirable medical reports without the training on the pairs
of image and report.

Semi-Supervised Setting To further prove the effectiveness of our approach, we
fine-tune the unsupervised KGAE using partial downstream paired image-report
datasets (see Section 3.3), resulting in the KGAE-Semi. To this end, in Figure
2, we evaluate the performance of our approach on both IU X-ray and MIMIC-CXR
datasets with respect to the increasing amount of paired data. For a fair
comparison, we also re-train the state-of-the-art model R2Gen [6] using the
same amount of pairs. As we can see, our model outperforms the R2Gen under all
ratios of paired dataset used for training. It is worth noting that the fewer
the image-report pairs, the larger the margins, e.g., under the very limited
pairs setting (20% of paired datasets), our approach significantly surpasses
the R2Gen by 8.5% absolute BLEU-4 score on IU X-ray and 9.8% absolute F1 score
on MIMIC-CXR. Intuitively, since our approach can relax the reliance on the
paired datasets, we can make use of available unpaired image and report data
as a solid bias for medical report generation task. Table 1 and Table 2
further prove the effectiveness of our approach, which achieves results
competitive with current state-of-the-art models by using only 60% of paired
dataset.

Table 1: Performance in terms of natural language generation metrics on the IU X-ray and MIMIC-CXR. B-n, M and R-L are short for BLEU-n, METEOR and ROUGE-L, respectively. Higher is better in all columns. It is worth noting that the KGAE is trained in an unsupervised manner, KGAE-Semi is trained 60% paired data, KGAE-Supervised and existing methods are trained with full paired data. Methods | Year |  Ratio of Pairs |  IU X-ray [9] |  MIMIC-CXR [17]
---|---|---|---|---
B-1 | B-2 | B-3 | B-4 | M | R-L | B-1 | B-2 | B-3 | B-4 | M | R-L
NIC [38] | 2015 | 100% | 0.216 | 0.124 | 0.087 | 0.066 | - | 0.306 | 0.299 | 0.184 | 0.121 | 0.084 | 0.124 | 0.263
AdaAtt [30] | 2017 | 100% | 0.220 | 0.127 | 0.089 | 0.068 | - | 0.308 | 0.299 | 0.185 | 0.124 | 0.088 | 0.118 | 0.266
Att2in [34] | 2017 | 100% | 0.224 | 0.129 | 0.089 | 0.068 | - | 0.308 | 0.325 | 0.203 | 0.136 | 0.096 | 0.134 | 0.276
Transformer [6] | 2020 | 100% | 0.396 | 0.254 | 0.179 | 0.135 | 0.164 | 0.342 | 0.314 | 0.192 | 0.127 | 0.090 | 0.125 | 0.265
$\mathcal{M}^{2}$ Trans. [7] | 2020 | 100% | 0.437 | 0.290 | 0.205 | 0.152 | 0.176 | 0.353 | 0.238 | 0.151 | 0.102 | 0.067 | 0.110 | 0.249
R2Gen [6] | 2020 | 100% | 0.470 | 0.304 | 0.219 | 0.165 | 0.187 | 0.371 | 0.353 | 0.218 | 0.145 | 0.103 | 0.142 | 0.277
KGAE | Ours | 0% | 0.417 | 0.263 | 0.181 | 0.126 | 0.149 | 0.318 | 0.221 | 0.144 | 0.096 | 0.062 | 0.097 | 0.208
KGAE-Semi | 60% | 0.497 | 0.320 | 0.232 | 0.171 | 0.189 | 0.379 | 0.352 | 0.219 | 0.149 | 0.108 | 0.147 | 0.290
KGAE-Supervised | 100% | 0.512 | 0.327 | 0.240 | 0.179 | 0.195 | 0.383 | 0.369 | 0.231 | 0.156 | 0.118 | 0.153 | 0.295

Table 2: Performance of automatic evaluation in terms of clinical efficacy metrics, which measure the accuracy of descriptions for clinical abnormalities, on the test set of MIMIC-CXR dataset. Higher is better in all columns. Methods | Year |  Ratio of Pairs |  MIMIC-CXR [17]
---|---|---|---
Precision | Recall | F1
NIC [38] | 2015 | 100% | 0.249 | 0.203 | 0.204
AdaAtt [30] | 2017 | 100% | 0.268 | 0.186 | 0.181
Att2in [34] | 2017 | 100% | 0.322 | 0.239 | 0.249
Up-Down [1] | 2018 | 100% | 0.320 | 0.231 | 0.238
$\mathcal{M}^{2}$ Trans. [7] | 2020 | 100% | 0.197 | 0.145 | 0.133
Transformer [6] | 2020 | 100% | 0.331 | 0.224 | 0.228
R2Gen [6] | 2020 | 100% | 0.333 | 0.273 | 0.276
KGAE | Ours | 0% | 0.214 | 0.158 | 0.156
KGAE-Semi | 60% | 0.360 | 0.302 | 0.307
KGAE-Supervised | 100% | 0.389 | 0.362 | 0.355

Figure 2: Performance of re-implemented state-of-the-art model R2Gen [6] and
our approach on the test sets with respect to various amount of paired data
used for training. The margins in different ratios are shown with the polyline
and the right y-axis. As we can see, our approach consistently outperforms the
R2Gen.

Supervised Setting We fine-tune the unsupervised KGAE using full image-report
pairs, acquiring the KGAE-Supervised model (see Section 3.3). Table 1 and
Table 2 show that KGAE-Supervised sets the new state-of-the-art results on the
two datasets in all metrics. Moreover, in terms of the clinical efficacy
metrics, our approach achieves 0.389 precision score, 0.362 recall score and
0.355 F1 score, outperforming the state-of-the-art model R2Gen [6]. The
superior clinical efficacy scores demonstrate the capability of our approach
to produce higher quality descriptions for clinical abnormalities than
existing models.

Overall Combining the results of unsupervised, semi-supervised and supervised
settings, the proposed KGAE can relax the dependency on the paired datasets,
and thus makes the medical report generation model use the available separate
image and report data to boost the performance. The advantages under the
scenarios with limited labeled pairs (i.e., semi-supervised setting) show that
KGAE might be applied to other medical images (MRI, Dermoscopy, Retinal,
etc.), where the coupled images and reports pairs could be much less or even
unavailable.

Table 3: We invite three professional clinicians to conduct human evaluation for comparing our approach with $\mathcal{M}^{2}$ Trans. [7] and R2Gen [6] under various amount of pairs for training in terms of the comprehensiveness of the generated true abnormalities and the faithfulness of the generated normalities and abnormalities to the ground truth reports. All values are reported in percentage (%). Metrics |  | KGAE (0%)
---
vs. $\mathcal{M}^{2}$ Trans. (100%)
| KGAE (0%)
---
vs. R2Gen (100%)
| KGAE-Semi (20%)
---
vs. R2Gen (20%)
| KGAE-Supervised (100%)
---
vs. R2Gen (100%)
Loss | Tie | Win | Loss | Tie | Win | Loss | Tie | Win | Loss | Tie | Win
Faithfulness | 34 | 18 | 48 | 62 | 15 | 23 | 21 | 10 | 69 | 25 | 22 | 53
Comprehensiveness | 27 | 20 | 53 | 57 | 17 | 26 | 24 | 11 | 65 | 32 | 21 | 47

###  4.3 Human Evaluation

We conduct human evaluations to verify the effectiveness of KGAE in clinical
practice. Specifically, to assist radiologists in clinical decision-making and
reduce their workload, it is important to generate accurate reports
(faithfulness), i.e., the model does not generate normalities and
abnormalities that does not exist according to doctors, with comprehensive
abnormalities (comprehensiveness), the model does not leave out the
abnormalities. Therefore, we randomly select 100 samples from the MIMIC-CXR
and invite three professional clinicians to compare our approach and baselines
independently. The clinicians are unaware of which model generates these
reports. The results are shown in Table 3. As we can see, under the
unsupervised setting, our approach achieves competitive results with the
supervised model, outperforming the $\mathcal{M}^{2}$ Trans with winning pick-
up percentages. Under the (semi-)supervised setting, our method is better than
state-of-the-art model R2Gen [6] in all metrics, especially for the semi-
supervised setting (20% of paired dataset), our method substantially surpasses
the R2Gen, which is in accordance with the automatic evaluation, by $69-21=48$
and $65-24=41$ points in terms of the faithfulness and comprehensiveness
metrics, respectively.

##  5 Analysis

In this section, we conduct several analysis to better understand our proposed
approach.

###  5.1 Knowledge Graph Sensitivity

In this section, to evaluate the knowledge graph sensitivity, we evaluate the
performances using different knowledge graphs defined on IU X-Ray only, MIMIC-
CXR only, both MIMIC-CXR and IU X-Ray. Table 4 shows the results of KGAE (0%)
and KGAE-Supervised (100%) on the IU X-ray dataset. As we can see, our KGAE
using different knowledge graphs can consistently outperform several existing
supervised models, i.e., NIC, AdaAtt, Att2in, across all metrics (Table 1).
Similarly, our KGAE-Supervised with different knowledge graphs can also
consistently outperform existing state-of-the-art model, i.e., R2Gen (Table
1). The results prove the robustness of our proposed model to the pre-defined
knowledge graph. Therefore, this work could provide a good basis or starting
point for the research of unsupervised medical report generation in other
clinical domains such as MRI and Dermoscopy.

Table 4: Analysis of the knowledge graph sensitivity. Performance of our approach using different knowledge graphs defined on IU X-Ray only, MIMIC-CXR only, both MIMIC-CXR and IU X-Ray. Methods |  Ratio of Pairs | Knowledge Graphs |  IU X-ray [9]
---|---|---|---
B-1 | B-2 | B-3 | B-4 | M | R-L
KGAE | 0% | IU X-Ray | 0.425 | 0.271 | 0.185 | 0.114 | 0.123 | 0.310
MIMIC-CXR | 0.417 | 0.263 | 0.181 | 0.126 | 0.149 | 0.318
MIMIC-CXR + IU X-Ray | 0.419 | 0.260 | 0.180 | 0.124 | 0.143 | 0.321
KGAE-Supervised | 100% | IU X-Ray | 0.519 | 0.331 | 0.235 | 0.174 | 0.191 | 0.376
MIMIC-CXR | 0.512 | 0.327 | 0.240 | 0.179 | 0.195 | 0.383
MIMIC-CXR + IU X-Ray | 0.505 | 0.323 | 0.239 | 0.181 | 0.198 | 0.385
Table 5: We conduct the quantitative analysis under both the unsupervised setting (KGAE) and supervised setting (KGAE-Supervised). We analysis the effect of the shared $\mathcal{F}$ in the encoder (Eq. (4)), and the effect of the number of the knowledge in the decoder’s knowledge bank $B$ (Eq. (6)). Methods | Settings | $\mathcal{F}$ | $B$ |  IU X-ray [9] |  MIMIC-CXR [17]
---|---|---|---|---|---
B-1 | B-2 | B-3 | B-4 | M | R-L | B-1 | B-2 | B-3 | B-4 | M | R-L
KGAE (0%) | (a) | - | - | 0.291 | 0.170 | 0.127 | 0.086 | 0.120 | 0.301 | 0.166 | 0.095 | 0.064 | 0.032 | 0.070 | 0.174
(b) | $\surd$ | - | 0.352 | 0.227 | 0.154 | 0.109 | 0.133 | 0.313 | 0.192 | 0.118 | 0.079 | 0.045 | 0.085 | 0.187
(c) | $\surd$ | 5,000 | 0.403 | 0.252 | 0.170 | 0.117 | 0.144 | 0.316 | 0.211 | 0.137 | 0.093 | 0.058 | 0.093 | 0.202
(d) | $\surd$ | 15,000 | 0.412 | 0.260 | 0.178 | 0.121 | 0.148 | 0.320 | 0.217 | 0.140 | 0.092 | 0.059 | 0.101 | 0.205
Full Model | $\surd$ | 10,000 | 0.417 | 0.263 | 0.181 | 0.126 | 0.149 | 0.318 | 0.221 | 0.144 | 0.096 | 0.062 | 0.097 | 0.208
KGAE-Supervised (100%) | (e) | - | - | 0.482 | 0.310 | 0.219 | 0.166 | 0.184 | 0.373 | 0.323 | 0.204 | 0.128 | 0.104 | 0.115 | 0.266
(f) | $\surd$ | - | 0.470 | 0.304 | 0.217 | 0.162 | 0.181 | 0.368 | 0.320 | 0.206 | 0.129 | 0.101 | 0.118 | 0.267
(g) | $\surd$ | 5,000 | 0.505 | 0.323 | 0.232 | 0.178 | 0.190 | 0.379 | 0.363 | 0.228 | 0.152 | 0.115 | 0.147 | 0.290
(h) | $\surd$ | 15,000 | 0.498 | 0.315 | 0.221 | 0.170 | 0.186 | 0.374 | 0.368 | 0.235 | 0.158 | 0.114 | 0.150 | 0.293
Full Model | $\surd$ | 10,000 | 0.512 | 0.327 | 0.240 | 0.179 | 0.195 | 0.383 | 0.369 | 0.231 | 0.156 | 0.118 | 0.153 | 0.295

###  5.2 Ablation Study

In Table 5, we conduct quantitative analysis to better understand our approach
under both the unsupervised and supervised training settings. For different
ablation settings, the KGAE-Supervised is acquired by further fine-tuning the
KGAE using image-report pairs.

As we can see, under the unsupervised setting, both the introduced shared
$\mathcal{F}$ and knowledge bank $B$ can significantly boost the performance,
which proves our arguments and verifies the effectiveness of our approach in
performing the unsupervised medical report generation.

Under the supervised setting, as shown in settings (e,f), applying shared
$\mathcal{F}$ generates unchanged and impaired performance on the MIMIC-CXR
and IU X-ray datasets, respectively. We speculate the reason is that the
supervised model no longer requires the $\mathcal{F}$ to bridge the vision and
the language domains. Therefore, the performance is unchanged on the large
dataset MIMIC-CXR. However, the increased parameters introduced by
$\mathcal{F}$ might bring overfitting or increase the difficulty in
optimization on the small dataset IU X-ray, which somewhat hinders the
performance. For the knowledge bank $B$, we can find that the results of
setting (h) outperforms (g) on the MIMIC-CXR dataset, but underperforms (g) on
the IU X-ray dataset. We speculate the reason is that the large MIMIC-CXR
dataset contains more knowledge than the small IU X-ray dataset, so a larger
knowledge bank may learn more knowledge of MIMIC-CXR to boost the performance
while introducing more noisy knowledge into the IU X-ray dataset to degrade
the performance.

Figure 3: Reports generated by our approach and a state-of-the-art model R2Gen
[6]. Under the unsupervised setting, we show four nodes with top-4 attention
weights in Eq. (4) to visualize $\mathcal{G}_{I}$ and $\mathcal{G}_{R}$. The
Purple and Blue colored text denote the generated wrong sentences
(Faithfulness) and true abnormalities (Comprehensiveness), respectively;
Underlined text denotes the repeated sentences.

###  5.3 Qualitative Analysis

In Figure 3, we conduct the qualitative analysis to better understand our
approach. As we can see, the visualization verifies the effectiveness of our
knowledge-driven encoder in extracting the knowledge representations of both
image and report. For the generated reports, our unsupervised KGAE generates a
desirable report, which correctly describes “innumerable nodules are present”
and “heart size is normal”. When removing the bank $B$, the model tends to
generate plausible general reports with no prominent abnormal narratives and
some repeated reports, which shows that the knowledge memory mechanism can
indeed distill and preserve the desired medical knowledge to boost the
generation of reports. Under the semi-supervised setting, the R2Gen can not
well handle the medical report generation task and generates some repeated
sentences of normalities (Underlined text) and fails to depict some rare but
important abnormalities, i.e., “nodules” and “scoliosis”, while our approach
can generate fluent report supported by accurate abnormalities. Under the
supervised setting, our approach can generate an accurate report showing
significant alignment with the ground truth report. It further prove our
arguments and the effectiveness of our proposed approach.

##  6 Conclusions and Discussions

In this paper, we propose the Knowledge Graph Auto-Encoder (KGAE). Without any
image-report pairs, KGAE can extract the knowledge representations of both
image and report from the knowledge graph to bridge the visual and textual
domains, and generate desirable reports by being trained in the auto-encoding
pipeline. The experiments verify the effectiveness of our approach, which even
exceeds several supervised models. Moreover, by further fine-tuning KGAE using
paired datasets, we achieve the state-of-the-art results on two public
datasets with the best human preference.

In the future, 1) since we can relax the dependency on paired data, it can be
interesting to apply the KGAE to other types of medical images of other body
parts, where the image-report pairs could be much less or even unavailable, to
assist radiologists in clinical decision-making and reduce their workload; 2)
We can replace the explicit pre-defined knowledge graph with an implicit large
matrix (e.g., knowledge bank in our decoder) to improve the generalization
ability of our approach.

Societal Impacts: In this paper, we target the problem of medical report
generation. Although the proposed model outperforms state-of-the-art
approaches, it aims to assist the radiologists instead of replacing them. For
experienced radiologists, given a large amount of medical images, our model
can automatically generate medical reports, the radiologists only need to make
revisions rather than write a new report from scratch. However, it is possible
that some radiologists direct copy the generated report as the final report.
Also for less experienced radiologists, they may not be able to correct the
errors in machine-generated reports. In order to apply the proposed model in
clinical practice, it is required to add process control to avoid unintended
use.

Limitations: Although the proposed KGAE can work in an unsupervised manner, we
still need the independent sets of medical images and medical reports which
may still be difficult to collect for some types of medical images. In
addition, our approach introduces the knowledge graph to bridge visual and
textual domain, in the paper, we collect the frequent clinical findings as
nodes and build the knowledge graph from the set of reports automatically.
When applying to new domains, we need to collect a new set of clinical
findings.

## References

  * Anderson et al. [2018] P. Anderson, X. He, C. Buehler, D. Teney, M. Johnson, S. Gould, and L. Zhang.  Bottom-up and top-down attention for image captioning and VQA.  In _CVPR_ , 2018.
  * Bahdanau et al. [2015] D. Bahdanau, K. Cho, and Y. Bengio.  Neural machine translation by jointly learning to align and translate.  In _ICLR_ , 2015.
  * Banerjee and Lavie [2005] S. Banerjee and A. Lavie.  METEOR: an automatic metric for MT evaluation with improved correlation with human judgments.  In _IEEvaluation@ACL_ , 2005.
  * Brady et al. [2012] A. Brady, R. Ó. Laoide, P. Mccarthy, and R. Mcdermott.  Discrepancy and error in radiology: Concepts, causes and consequences.  _The Ulster Medical Journal_ , 81:3 – 9, 2012.
  * Chen et al. [2015] X. Chen, H. Fang, T. Lin, R. Vedantam, S. Gupta, P. Dollár, and C. L. Zitnick.  Microsoft COCO captions: Data collection and evaluation server.  _arXiv preprint arXiv:1504.00325_ , 2015.
  * Chen et al. [2020] Z. Chen, Y. Song, T. Chang, and X. Wan.  Generating radiology reports via memory-driven transformer.  In _EMNLP_ , 2020.
  * Cornia et al. [2020] M. Cornia, M. Stefanini, L. Baraldi, and R. Cucchiara.  Meshed-memory transformer for image captioning.  In _CVPR_ , 2020.
  * Delrue et al. [2011] L. Delrue, R. Gosselin, B. Ilsen, A. Van Landeghem, J. de Mey, and P. Duyck.  Difficulties in the interpretation of chest radiography.  _Comparative Interpretation of CT and Standard Radiography of the Chest_ , pages 27–49, 2011.
  * Demner-Fushman et al. [2016] D. Demner-Fushman, M. D. Kohli, M. B. Rosenman, S. E. Shooshan, L. Rodriguez, S. K. Antani, G. R. Thoma, and C. J. McDonald.  Preparing a collection of radiology examinations for distribution and retrieval.  _J. Am. Medical Informatics Assoc._ , 23(2):304–310, 2016.
  * Deng et al. [2009] J. Deng, W. Dong, R. Socher, L. Li, K. Li, and F. Li.  Imagenet: A large-scale hierarchical image database.  In _CVPR_ , 2009.
  * Goergen et al. [2013] S. K. Goergen, F. J. Pool, T. J. Turner, J. E. Grimm, M. N. Appleyard, C. Crock, M. C. Fahey, M. F. Fay, N. J. Ferris, S. M. Liew, et al.  Evidence-based guideline for the written radiology report: Methods, recommendations and implementation challenges.  _Journal of medical imaging and radiation oncology_ , 57(1):1–7, 2013.
  * He et al. [2016] K. He, X. Zhang, S. Ren, and J. Sun.  Deep residual learning for image recognition.  In _CVPR_ , 2016.
  * Hochreiter and Schmidhuber [1997] S. Hochreiter and J. Schmidhuber.  Long short-term memory.  _Neural Comput._ , 9(8):1735–1780, 1997.
  * Irvin et al. [2019] J. Irvin, P. Rajpurkar, M. Ko, Y. Yu, S. Ciurea-Ilcus, C. Chute, H. Marklund, B. Haghgoo, R. L. Ball, K. S. Shpanskaya, J. Seekins, D. A. Mong, S. S. Halabi, J. K. Sandberg, R. Jones, D. B. Larson, C. P. Langlotz, B. N. Patel, M. P. Lungren, and A. Y. Ng.  Chexpert: A large chest radiograph dataset with uncertainty labels and expert comparison.  In _AAAI_ , 2019.
  * Jing et al. [2018] B. Jing, P. Xie, and E. P. Xing.  On the automatic generation of medical imaging reports.  In _ACL_ , 2018.
  * Jing et al. [2019] B. Jing, Z. Wang, and E. P. Xing.  Show, describe and conclude: On exploiting the structure information of chest x-ray reports.  In _ACL_ , 2019.
  * Johnson et al. [2019] A. E. W. Johnson, T. J. Pollard, S. J. Berkowitz, N. R. Greenbaum, M. P. Lungren, C. Deng, R. G. Mark, and S. Horng.  MIMIC-CXR: A large publicly available database of labeled chest radiographs.  _arXiv preprint arXiv:1901.07042_ , 2019.
  * Kingma and Ba [2014] D. P. Kingma and J. Ba.  Adam: A method for stochastic optimization.  In _ICLR_ , 2014.
  * Krause et al. [2017] J. Krause, J. Johnson, R. Krishna, and L. Fei-Fei.  A hierarchical approach for generating descriptive image paragraphs.  In _CVPR_ , 2017.
  * Krizhevsky et al. [2012] A. Krizhevsky, I. Sutskever, and G. E. Hinton.  Imagenet classification with deep convolutional neural networks.  In _NIPS_ , 2012.
  * Kumar et al. [2016] A. Kumar, O. Irsoy, P. Ondruska, M. Iyyer, J. Bradbury, I. Gulrajani, V. Zhong, R. Paulus, and R. Socher.  Ask me anything: Dynamic memory networks for natural language processing.  In _ICML_ , 2016.
  * Li et al. [2019] C. Y. Li, X. Liang, Z. Hu, and E. P. Xing.  Knowledge-driven encode, retrieve, paraphrase for medical image report generation.  In _AAAI_ , 2019.
  * Li et al. [2020] M. Li, F. Wang, X. Chang, and X. Liang.  Auxiliary signal-guided knowledge encoder-decoder for medical report generation.  _arXiv preprint arXiv:2006.03744_ , 2020.
  * Li et al. [2016] Y. Li, D. Tarlow, M. Brockschmidt, and R. S. Zemel.  Gated graph sequence neural networks.  In _ICLR_ , 2016.
  * Li et al. [2018] Y. Li, X. Liang, Z. Hu, and E. P. Xing.  Hybrid retrieval-generation reinforced agent for medical image report generation.  In _NeurIPS_ , 2018.
  * Lin [2004] C.-Y. Lin.  ROUGE: A package for automatic evaluation of summaries.  In _ACL_ , 2004.
  * Lin et al. [2017] Z. Lin, M. Feng, C. N. dos Santos, M. Yu, B. Xiang, B. Zhou, and Y. Bengio.  A structured self-attentive sentence embedding.  In _ICLR_ , 2017.
  * Liu et al. [2018] F. Liu, X. Ren, Y. Liu, H. Wang, and X. Sun.  simnet: Stepwise image-topic merging network for generating detailed and comprehensive image captions.  In _EMNLP_ , 2018.
  * Liu et al. [2019] F. Liu, Y. Liu, X. Ren, X. He, and X. Sun.  Aligning visual regions and textual concepts for semantic-grounded image representations.  In _NeurIPS_ , 2019.
  * Lu et al. [2017] J. Lu, C. Xiong, D. Parikh, and R. Socher.  Knowing when to look: Adaptive attention via a visual sentinel for image captioning.  In _CVPR_ , 2017.
  * Marcheggiani and Titov [2017] D. Marcheggiani and I. Titov.  Encoding sentences with graph convolutional networks for semantic role labeling.  In _EMNLP_ , 2017.
  * Miller et al. [2016] A. H. Miller, A. Fisch, J. Dodge, A. Karimi, A. Bordes, and J. Weston.  Key-value memory networks for directly reading documents.  In _EMNLP_ , 2016.
  * Papineni et al. [2002] K. Papineni, S. Roukos, T. Ward, and W. Zhu.  BLEU: a Method for automatic evaluation of machine translation.  In _ACL_ , 2002.
  * Rennie et al. [2017] S. J. Rennie, E. Marcheret, Y. Mroueh, J. Ross, and V. Goel.  Self-critical sequence training for image captioning.  In _CVPR_ , 2017.
  * Soricut et al. [2018] R. Soricut, N. Ding, P. Sharma, and S. Goodman.  Conceptual captions: A cleaned, hypernymed, image alt-text dataset for automatic image captioning.  In _ACL_ , 2018.
  * Sukhbaatar et al. [2015] S. Sukhbaatar, A. Szlam, J. Weston, and R. Fergus.  End-to-end memory networks.  In _NIPS_ , 2015.
  * Vaswani et al. [2017] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin.  Attention is all you need.  In _NIPS_ , 2017.
  * Vinyals et al. [2015] O. Vinyals, A. Toshev, S. Bengio, and D. Erhan.  Show and tell: A neural image caption generator.  In _CVPR_ , 2015.
  * Vinyals et al. [2016] O. Vinyals, C. Blundell, T. Lillicrap, K. Kavukcuoglu, and D. Wierstra.  Matching networks for one shot learning.  In _NIPS_ , 2016.
  * Wang et al. [2018] X. Wang, Y. Peng, L. Lu, Z. Lu, and R. M. Summers.  Tienet: Text-image embedding network for common thorax disease classification and reporting in chest x-rays.  In _CVPR_ , 2018.
  * Xiong et al. [2016] C. Xiong, S. Merity, and R. Socher.  Dynamic memory networks for visual and textual question answering.  In _ICML_ , 2016.
  * Xu and Saenko [2016] H. Xu and K. Saenko.  Ask, attend and answer: Exploring question-guided spatial attention for visual question answering.  In _ECCV_ , 2016.
  * Xu et al. [2015] K. Xu, J. Ba, R. Kiros, K. Cho, A. C. Courville, R. Salakhutdinov, R. S. Zemel, and Y. Bengio.  Show, attend and tell: Neural image caption generation with visual attention.  In _ICML_ , 2015.
  * Xue et al. [2018] Y. Xue, T. Xu, L. R. Long, Z. Xue, S. K. Antani, G. R. Thoma, and X. Huang.  Multimodal recurrent model with attention for automated radiology report generation.  In _MICCAI_ , 2018.
  * Yang et al. [2019] X. Yang, K. Tang, H. Zhang, and J. Cai.  Auto-encoding scene graphs for image captioning.  In _CVPR_ , 2019.
  * You et al. [2019] C. You, G. Li, Y. Zhang, X. Zhang, H. Shan, M. Li, S. Ju, Z. Zhao, Z. Zhang, W. Cong, et al.  CT super-resolution GAN constrained by the identical, residual, and cycle learning ensemble (gan-circle).  _IEEE Transactions on Medical Imaging_ , 2019.
  * You et al. [2021a] C. You, N. Chen, and Y. Zou.  Knowledge distillation for improved accuracy in spoken question answering.  In _ICASSP_ , 2021a.
  * You et al. [2021b] C. You, N. Chen, and Y. Zou.  MRD-Net: Multi-Modal Residual Knowledge Distillation for Spoken Question Answering.  In _IJCAI_ , 2021b.
  * You et al. [2021c] C. You, N. Chen, and Y. Zou.  Self-supervised contrastive cross-modality representation learning for spoken question answering.  _arXiv preprint arXiv:2109.03381_ , 2021c.
  * You et al. [2021d] C. You, R. Zhao, L. Staib, and J. S. Duncan.  Momentum contrastive voxel-wise representation learning for semi-supervised volumetric medical image segmentation.  _arXiv preprint arXiv:2105.07059_ , 2021d.
  * You et al. [2021e] C. You, Y. Zhou, R. Zhao, L. Staib, and J. S. Duncan.  Simcvd: Simple contrastive voxel-wise representation distillation for semi-supervised medical image segmentation.  _arXiv preprint arXiv:2108.06227_ , 2021e.
  * Yuan et al. [2019] J. Yuan, H. Liao, R. Luo, and J. Luo.  Automatic radiology report generation based on multi-view image fusion and medical concept enrichment.  In _MICCAI_ , 2019.
  * Zhang et al. [2020] Y. Zhang, X. Wang, Z. Xu, Q. Yu, A. L. Yuille, and D. Xu.  When radiology report generation meets knowledge graph.  In _AAAI_ , 2020.


# KerasCV and KerasNLP: Vision and Language Power-Ups

Lead Authors
Matthew Watson, Divyashree Shivakumar Sreepathihalli, François Chollet
Martin Görner, Kiranbir Sodhia, Ramesh Sampath, Tirth Patel
Haifeng Jin, Neel Kovelamudi, Gabriel Rasskin, Samaneh Saadat
Luke Wood, Chen Qian, Jonathan Bischof, Ian Stenbit {mattdangerw,
divyasreepat, fchollet, mgorner<EMAIL_ADDRESS>
{rameshsampath, tirthp, haifengj, nkovela, grasskin<EMAIL_ADDRESS>
{lukewoodcs, qianchen94era<EMAIL_ADDRESS><EMAIL_ADDRESS>Keras Team,
Google LLC, USA Community Contributors
Abheesht Sharma, Anshuman Mishra

###### Abstract

We present the Keras domain packages KerasCV and KerasNLP, extensions of the
Keras API for Computer Vision and Natural Language Processing workflows,
capable of running on either JAX, TensorFlow, or PyTorch. These domain
packages are designed to enable fast experimentation, with a focus on ease-of-
use and performance. We adopt a modular, layered design: at the library’s
lowest level of abstraction, we provide building blocks for creating models
and data preprocessing pipelines, and at the library’s highest level of
abstraction, we provide pretrained “task” models for popular architectures
such as Stable Diffusion, YOLOv8, GPT2, BERT, Mistral, CLIP, Gemma, T5, etc.
Task models have built-in preprocessing, pretrained weights, and can be fine-
tuned on raw inputs. To enable efficient training, we support XLA compilation
for all models, and run all preprocessing via a compiled graph of TensorFlow
operations using the `tf.data` API. The libraries are fully open-source
(Apache 2.0 license) and available on GitHub.

Keywords: KerasCV, KerasNLP, Keras multibackend, Deep learning, Generative AI

##  1 Introduction

Keras (Chollet et al., 2015) is among the most widely used tools for machine
learning today111https://survey.stackoverflow.co/2022/. The Keras library acts
as a high-level abstraction for machine learning models and layers, and seeks
to be accessible to a broad group of machine learning researchers and
practitioners by focusing on rapid experimentation and progressive disclosure
of complexity.

Notably, recent developments in Computer Vision (CV) and Natural Language
Processing (NLP) have created new challenges for practitioners. The most
obvious is the shift towards larger and larger models trained on self-
supervised tasks. Pretraining a state of the art model is now cost-prohibitive
for many researchers and practitioners, in particular in NLP. Access to open-
source model architectures with pretrained weights is imperative for a large
amount of CV and NLP work.

Additionally, pairing efficient preprocessing and metrics computation for
modern models has become more difficult, with a proliferation of disparate
techniques, backends, and licenses. An ML researcher or practitioner today
must select among a range of auto-differentiation frameworks such as JAX,
TensorFlow, and PyTorch, and even within each framework, they are often forced
to stay within a specific modeling library for cross-compatibility of
components.

Further, improving the train-time performance of models on NLP problems
presents additional hurdles. The XLA compiler (Sabne, 2020) offers dramatic
speedups for many model architectures, but adds complex restrictions on the
shape and flow of tensor operations. The TensorFlow-based `tf.data` (Murray et
al., 2021) and `tf.text` APIs provide a scalable, dynamic, and multi-process
approach for preprocessing, but many common text operations do not easily
compile to a TensorFlow graph.

Aiming to reduce these framework barriers for CV and NLP practitioners and
researchers, we present KerasCV and KerasNLP, extensions of the Keras API for
computer vision and natural language processing workflows. These packages
expand upon the modular approach of Keras, adding pretrained backbone models,
easy-to-use domain-specific losses and metrics, out-of-the-box support for XLA
(Sabne, 2020) compilation and data and model parallelism. Because these domain
packages are written on top of Keras 3, all of their modeling components
natively support JAX (Bradbury et al., 2018), TensorFlow (Abadi et al., 2015),
and PyTorch (Paszke et al., 2019), and can be freely used in framework-native
workflows that do not otherwise involve any Keras components.

##  2 The Keras Domain Packages API

We adopt a layered approach to API design. Our library has three primary
levels of abstraction:

  * •

Foundational Components: A collection of composable modules for building and
training preprocessing pipelines, models, and evaluation logic. These are pure
Keras 3 components which can be used outside of the Keras Domain Packages
ecosystem.

  * •

Pretrained Backbones: A collection of pretrained model backbones for fine-
tuning. For NLP models, matching tokenizers can be created alongside
backbones.

  * •

Task Models: A collection of end-to-end models specialized for a specific
task, e.g. text generation in NLP or object detection in CV. These task models
combine the preprocessing and modeling modules from the lower API levels to
create a unified training and inference interface that can operate directly on
plain text or image input. Task models aim to allow fine-tuning with zero
configuration for common use cases.

Each additional API layer is built on top of the previous one. Modules from
each level can be mixed and matched in usage, for example, extending a
pretrained backbone with foundational preprocessing modules to pack input
sequences or perform data augmentation.

Any KerasCV and KerasNLP model can be instantiated as a PyTorch
`torch.nn.Module`, a TensorFlow `tf.Module`, or as a stateless JAX function.
This means that the models can be used with PyTorch ecosystem packages, with
the full range of TensorFlow deployment and production tools (such as TF-
Serving, TF.js and TFLite), and with JAX large-scale TPU training
infrastructure.

##  3 Training, Serving, and Deployment

KerasCV and KerasNLP offer large vision and language models. State of the art
models are expected to continually increase in size in the future. To address
these problems, KerasCV and KerasNLP are compatible with the Keras Unified
Distribution API. This API enables both model parallelism and data parallelism
across all Keras backends. The API maintains a clear separation between the
model definition, training logic, and sharding configuration. As a result,
models within KerasCV and KerasNLP can be written as if they were intended to
run on a single device. Later, specific sharding configurations can be added
to these models when it’s time to train
them.222https://keras.io/api/distribution/

##  4 Pretrained models on Kaggle Models

All pretrained models of KerasCV and KerasNLP are published on Kaggle Models
333https://www.kaggle.com/organizations/keras/models. Importantly, these
models are also available on Kaggle competition notebooks in Internet-off
mode.

| SAM | Gemma | BERT | Mistral
---|---|---|---|---
| train | predict | train | predict | train | predict | train | predict
Batch Size | 1 | 7 | 8 | 32 | 54 | 531 | 8 | 32
Keras 2 (TF) | 386.93 | 3,187.09 | NA | NA | 841.84 | 965.21 | NA | NA
Keras 3 (TF) | 355.25 | 762.67 | 232.52 | 1,134.91 | 404.17 | 962.11 | 185.92 | 966.06
Keras 3 (JAX) | 361.69 | 660.16 | 273.67 | 1,128.21 | 414.26 | 865.29 | 213.22 | 957.25
Keras 3 (PT) | 1,388.87 | 2,973.64 | 525.15 | $7,952.67^{*}$ | 1320.441 | 3869.72 | 452.12 | $10932.59^{*}$
Keras 3 (best) | 355.25 | 660.16 | 232.52 | 1,128.21 | 404.17 | 865.29 | 185.92 | 957.25
Table 1: Average time taken (in ms/step) per training or inference step across
different models, namely SAM (Kirillov et al., 2023), Gemma (Team et al.,
2024), BERT (Devlin et al., 2019) and Mistral (Jiang et al., 2023).
* LLM inference with the PyTorch backend is abnormally slow at this time because KerasNLP uses static sequence padding. This will be addressed soon.

##  5 Performance

Framework performance depends on the specific model. Keras 3 offers
flexibility by letting users select the fastest framework for their task.
Picking the fastest backend for a given model consistently outperforms Keras 2
as seen in Table 1. All benchmarks are done with a single NVIDIA A100 GPU with
40GB of GPU memory on a Google Cloud Compute Engine of machine type
a2-highgpu-1g with 12 vCPUs and 85GB host memory.

For fair comparison, we use the same batch size across frameworks if it is the
same model and task (fit or predict). However, for different models and tasks,
due to their different sizes and architectures, we use different batch sizes
to avoid either running out of memory (too large) or under GPU utilization
(too small). We also used the same batch size for Gemma and Mistral since they
are the same model type with similar number of parameters. (see Table 1).

##  6 Related Work

A library with clear parallels to KerasNLP and KerasCV is the HuggingFace
Transformers library (Wolf et al., 2020). Both libraries offer access to
pretrained model checkpoints for a number of widely-used transformer
architectures.

The Transformers library is built with a “repeat yourself”
approach444https://huggingface.co/blog/transformers-design-philosophy.
KerasNLP, in contrast, is built with a layered approach, with an explicit goal
of allowing the re-implementation of any large language model in a relatively
small amount of code. We believe there are strengths and weaknesses to both of
these approaches.

##  7 Future Work

Future work will focus on extending the project’s capabilities. We plan to
expand our multimodal model offerings, supporting a wider range of
applications. Additionally, we will optimize integrations with backend-
specific large model serving solutions, ensuring seamless deployment and
scalability.

##  8 Conclusions

KerasCV and KerasNLP are new toolboxes offering both modular components for
rapid prototyping of new models, as well as standard pretrained backbones and
task models for many computer vision and natural language processing
workflows. They can be leveraged by users of either JAX (Bradbury et al.,
2018), TensorFlow (Abadi et al., 2015), or PyTorch (Paszke et al., 2019).
Thanks to backend optionality and XLA (Sabne, 2020) compilation, KerasCV and
KerasNLP deliver state-of-the-art training and inference performance.
KerasCV555https://keras.io/guides/keras_cv/ and
KerasNLP666https://keras.io/guides/keras_nlp/ offer extensive user guides,
available at Keras.io.

Acknowledgments

We thank all contributors to Keras (Chollet et al., 2015), KerasCV (Wood et
al., 2024), KerasNLP (Watson et al., 2024), TensorFlow (Abadi et al., 2015),
TensorFlow Text, TensorFlow Data, and the XLA (Sabne, 2020) compiler, all of
which are crucial to the functionality provided in KerasCV and KerasNLP.

##  Appendix A Preprocessing Layers

KerasCV provides a comprehensive suite of preprocessing layers that empower
users to construct state-of-the-art, industry-grade data augmentation
pipelines for image classification, object detection, image segmentation and
image generation tasks. These layers implement a wide range of commonly used
data augmentation techniques, enabling users to effortlessly enhance the
robustness and generalizability of their models. By using preprocessing
layers, users can ensure that their models are trained on data that is
representative of the data that they will encounter at inference time. KerasCV
offers 38 data augmentation layers. These layers implement a wide range of
commonly used data augmentation techniques, enabling users to effortlessly
manipulate image data in a variety of ways and handle all types of labels out-
of-the-box (e.g. class labels, box labels, mask labels).

TF Data is a TensorFlow API for building input pipelines. Input pipelines are
responsible for loading data from disk, preprocessing it, and batching it. TF
Data provides a number of features that make it a powerful tool for
preprocessing data for machine learning, such as:

  * •

Dataset APIs: for loading data from a variety of sources, such as CSV files,
`TFRecords`, and images.

  * •

Preprocessing functions: for performing common preprocessing tasks, such as
decoding images, resizing images, and normalizing images.

  * •

Batching functions: for grouping data into batches.

  * •

Prefetching and caching: for improving the performance of input pipelines.

⬇

# Apply grayscale preprocessing to input

(images, labels), _ = keras.datasets.cifar10.load_data()

to_grayscale = keras_cv.layers.preprocessing.Grayscale()

augmented_images = to_grayscale(images)

##  Appendix B Preset API

The presets API provides a convenient way to create state-of-the-art CV and
NLP models. Presets are pre-configured models that have been trained on a
specific dataset and can be used for a specific task.

To use the presets API, one simply needs to import the `keras_cv.models` or
`keras_nlp.models` module and then call the `from_preset()` method on the
desired model class. The presets API provides a number of advantages over
creating models from scratch. First, presets are pretrained on a large
dataset, which means that they can achieve high accuracy on a variety of
tasks. Second, presets are pre-configured, which means that users do not need
to worry about setting hyperparameters. Third, presets are easy to use, which
means that users can get started with them quickly.

⬇

# Load architecture and weights from preset

model = keras_cv.models.RetinaNet.from_preset(

"resnet50_imagenet",

)

# Load randomly initialized model from preset architecture with weights

model = keras_cv.models.RetinaNet.from_preset(

"resnet50_imagenet",

load_weights=False,

)

##  Appendix C Backbone API

Both KerasCV and KerasNLP offer a Backbone API. Backbones can be thought of as
the central architecture of a model, without the final output layer. This
allows users to leverage powerful pretrained backbones (often trained on vast
datasets) as the starting point for their own customized models. The
pretrained backbones within KerasCV and KerasNLP offer more than just a
starting point, they are also fine-tunable. Several examples of how to do this
can be seen on the Keras.io webpage (Chollet et al., 2015).

##  Appendix D Task Models

KerasCV and KerasNLP provide a number of task models that are designed for
specific tasks. These task models are built on top of the KerasCV and KerasNLP
modeling layers and provide a high-level of performance. These models are
ready for use in applications, but can be further fine-tuned if desired. Some
examples of available task models include image classification, object
detection, semantic segmentation, image generation, text generation, text
classification, and question answering.

### Pretrained Task Models

Pretrained task models can be used by using presets trained on different
datasets. This allows users to quickly and easily get started with deep
learning without having to train a model from scratch. For example, KerasCV
provides a number of presets for image classification models that have been
trained on different datasets, such as ImageNet, COCO, and Pascal VOC. These
presets can be used to create models that can achieve state-of-the-art results
on a variety of image classification tasks. To use a pretrained task model
with a preset, one simply needs to import the `keras_cv.models` or
`keras_nlp.models` module and then call the `from_preset()` method on the
desired model class.

### Specifying a Backbone in a Task Model

It is possible to specify a backbone for task models. This is done by passing
the backbone argument to the `from_preset()` method. By specifying a different
backbone, users can change the features that are extracted. This can be useful
if one wants to improve the performance of the model on a specific task. User
can also specify their own custom backbones. To do this, one simply need to
create a subclass of the `models.Backbone` class.

### Fine-Tuning a Task Model

Fine-tuning a task model is the process of adapting a pretrained model to a
specific task. This is done by training the model on a dataset of labeled data
for the specific task.

⬇

semantic_segmentation_model = keras_cv.models.DeepLabV3Plus.from_preset(

"resnet50_v2_imagenet", num_classes=NUM_CLASSES

)

## References

  * Abadi et al. (2015) Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng.  TensorFlow: Large-scale machine learning on heterogeneous systems, 2015.  URL https://www.tensorflow.org/.  Software available from tensorflow.org.
  * Bradbury et al. (2018) James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang.  JAX: composable transformations of Python+NumPy programs, 2018.  URL http://github.com/google/jax.
  * Chollet et al. (2015) François Chollet et al.  Keras.  https://keras.io, 2015.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In Jill Burstein, Christy Doran, and Thamar Solorio, editors, _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.  doi: 10.18653/v1/N19-1423.  URL https://aclanthology.org/N19-1423.
  * Jiang et al. (2023) Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed.  Mistral 7b, 2023.
  * Kirillov et al. (2023) Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollar, and Ross Girshick.  Segment anything.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)_ , pages 4015–4026, October 2023.
  * Murray et al. (2021) Derek Gordon Murray, Jiri Simsa, Ana Klimovic, and Ihor Indyk.  tf.data: A machine learning data processing framework.  _CoRR_ , abs/2101.12127, 2021.  URL https://arxiv.org/abs/2101.12127.
  * 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.  Pytorch: An imperative style, high-performance deep learning library.  In _Advances in Neural Information Processing Systems 32_ , pages 8024–8035. Curran Associates, Inc., 2019.  URL http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf.
  * Sabne (2020) Amit Sabne.  XLA : Compiling machine learning for peak performance, 2020.
  * Team et al. (2024) Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, Pier Giuseppe Sessa, Aakanksha Chowdhery, Adam Roberts, Aditya Barua, Alex Botev, Alex Castro-Ros, Ambrose Slone, Amélie Héliou, Andrea Tacchetti, Anna Bulanova, Antonia Paterson, Beth Tsai, Bobak Shahriari, Charline Le Lan, Christopher A. Choquette-Choo, Clément Crepy, Daniel Cer, Daphne Ippolito, David Reid, Elena Buchatskaya, Eric Ni, Eric Noland, Geng Yan, George Tucker, George-Christian Muraru, Grigory Rozhdestvenskiy, Henryk Michalewski, Ian Tenney, Ivan Grishchenko, Jacob Austin, James Keeling, Jane Labanowski, Jean-Baptiste Lespiau, Jeff Stanway, Jenny Brennan, Jeremy Chen, Johan Ferret, Justin Chiu, Justin Mao-Jones, Katherine Lee, Kathy Yu, Katie Millican, Lars Lowe Sjoesund, Lisa Lee, Lucas Dixon, Machel Reid, Maciej Mikuła, Mateo Wirth, Michael Sharman, Nikolai Chinaev, Nithum Thain, Olivier Bachem, Oscar Chang, Oscar Wahltinez, Paige Bailey, Paul Michel, Petko Yotov, Rahma Chaabouni, Ramona Comanescu, Reena Jana, Rohan Anil, Ross McIlroy, Ruibo Liu, Ryan Mullins, Samuel L Smith, Sebastian Borgeaud, Sertan Girgin, Sholto Douglas, Shree Pandya, Siamak Shakeri, Soham De, Ted Klimenko, Tom Hennigan, Vlad Feinberg, Wojciech Stokowiec, Yu hui Chen, Zafarali Ahmed, Zhitao Gong, Tris Warkentin, Ludovic Peran, Minh Giang, Clément Farabet, Oriol Vinyals, Jeff Dean, Koray Kavukcuoglu, Demis Hassabis, Zoubin Ghahramani, Douglas Eck, Joelle Barral, Fernando Pereira, Eli Collins, Armand Joulin, Noah Fiedel, Evan Senter, Alek Andreev, and Kathleen Kenealy.  Gemma: Open models based on gemini research and technology, 2024.
  * Watson et al. (2024) Matthew Watson, Chen Qian, Bischof Jonathan, François Chollet, et al.  Kerasnlp.  https://github.com/keras-team/keras-nlp, 2024.
  * Wolf et al. (2020) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations_ , pages 38–45, 2020.
  * Wood et al. (2024) Luke Wood, Zhenyu Tan, Ian Stenbit, Jonathan Bischof, Scott Zhu, François Chollet, Divyashree Sreepathihalli, Ramesh Sampath, et al.  Kerascv.  https://github.com/keras-team/keras-cv, 2024.