# Learning Better Contrastive View from Radiologist’s Gaze

Sheng Wang, Zixu Zhuang, Xi Ouyang, Lichi Zhang, Zheren Li, Chong Ma,
Tianming Liu, Dinggang Shen, , Qian Wang Sheng Wang, Zixu Zhuang, Lichi Zhang
and Zheren Li are with the School of Biomedical Engineering, Shanghai Jiao
Tong University, Shanghai, China. E-mail: {wsheng, zixu.zhuang, lichizhang,
<EMAIL_ADDRESS>Liu is with the Department of Computer Science,
University of Georgia, GA, USA. E-mail<EMAIL_ADDRESS>Ma is with the
School of Automation, Northwestern Polytechnical University, Xi’an, China.
E-mail<EMAIL_ADDRESS>Shen and Qian Wang are with the
School of Biomedical Engineering, ShanghaiTech University, Shanghai, China.
E-mail:{dgshen<EMAIL_ADDRESS>Xi Ouyang, Zheren Li and
Dinggang Shen are with the Department of Research and Development, Shanghai
United Imaging Intelligence Co., Ltd., Shanghai, China. E-mail:xi.ouyang@uii-
ai.com

###### Abstract

Recent self-supervised contrastive learning methods greatly benefit from the
Siamese structure that aims to minimizing distances between positive pairs.
These methods usually apply random data augmentation to input images,
expecting the augmented views of the same images to be similar and positively
paired. However, random augmentation may overlook image semantic information
and degrade the quality of augmented views in contrastive learning. This issue
becomes more challenging in medical images since the abnormalities related to
diseases can be tiny, and are easy to be corrupted (e.g., being cropped out)
in the current scheme of random augmentation. In this work, we first
demonstrate that, for widely-used X-ray images, the conventional augmentation
prevalent in contrastive pre-training can affect the performance of the
downstream diagnosis or classification tasks. Then, we propose a novel
augmentation method, i.e., FocusContrast, to learn from radiologists’ gaze in
diagnosis and generate contrastive views for medical images with guidance from
radiologists’ visual attention. Specifically, we track the gaze movement of
radiologists and model their visual attention when reading to diagnose X-ray
images. The learned model can predict visual attention of the radiologists
given a new input image, and further guide the attention-aware augmentation
that hardly neglects the disease-related abnormalities. As a plug-and-play and
framework-agnostic module, FocusContrast consistently improves state-of-the-
art contrastive learning methods of SimCLR, MoCo, and BYOL by 4.0$\sim$7.0% in
classification accuracy on a knee X-ray dataset.

###### Index Terms:

Contrastive Learning, Eye-tracking, Medical Image, Human Visual Attention.

##  I Introduction

Self-supervised contrastive learning has seen great progress recently. It has
begun to match supervised pre-training in performance on several downstream
tasks [1, 2, 3].

The basic idea of contrastive learning is to encourage the representations
from the positive pairs of the samples to be close in the embedding space,
while the representations from the negative pairs are pushed apart. Here, two
views that are transformed or augmented from the same training sample are
usually defined as a positive pair, while two views from different training
samples are defined as a negative pair.

The generation of positive pairs, which is essential for the quality of the
learned representations, is a major focus in contrastive learning. [3, 4, 5].
A common approach in computer vision generates the positive pair by randomly
augmenting a certain image twice, e.g., random cropping, color distorting,
rotating, cutout, adding noise, and so on [6, 7].

However, because this process does not take image semantics into account, it
sometimes misses the regions of interest that are essential to the subsequent
tasks. For example, the abnormality related to knee osteoarthritis (OA)
diagnosis in Fig. 1 (highlighted by the yellow box) may be easily lost during
the conventional scheme of random cropping or cutout.

Figure 1: We propose to learn and generate positive contrastive views from the
radiologist’s gaze. Our proposed framework can preserve highly informative
area (e.g., the abnormality in knee OA diagnosis, as highlighted by the yellow
boxes) throughout image augmentation, which may be lost by conventional random
augmentation. Figure 2: Improper contrastive views in medical image caused by
strong random augmentation. Here we take knee X-ray as an example. On the
left, one can observe that the decision-related abnormality can be very subtle
with a small number of pixels. In comparison with natural images, the object
of interest usually occupies a large number of pixels. We also demonstrate
increasing aggressiveness of the augmentation applied to the knee X-ray image
and the dog image. When the X-ray image is strongly augmented, the semantic
cues (i.e., abnormalities) are hard to preserve and recognize, implying that
the augmented images may deviate significantly from the subsequent task of
diagnosis. In comparison, when the dog image is heavily augmented, the
semantic cues are much easier to distinguish and the classification task can
be accomplished accordingly.

There are recent debates about how to design the augmentation strategies and
how to control their strength [8, 9, 10]. For example, 2$\times$ random
resized cropping is regarded as a stronger augmentation than 1.2$\times$
random resized cropping, while an excessively strong augmentation may not
always yield a good outcome. And the augmentation strength should be
emphasized especially for some medical image scenarios. As shown in Fig. 2,
the visual cues related to the diseased abnormality in the X-ray image may
only occupy a tiny number of pixels (e.g., 4.12% of the image area for the
abnormality of femur-tibia space narrowing), which are however expected to
provide major information to support clinical decision or diagnosis. These few
pixels are likely to be removed especially when the augmentation strength
becomes excessive. And the resulting positive pairs may corrupt the
contrastive representations learned from the images, since the visual cues for
subsequent diagnosis are no longer available. On the contrary, a natural image
depicting a dog (with 35.50% foreground pixels in Fig. 2) is less susceptible
to the distortion induced by image augmentation. That is, given similar strong
image augmentation (at the right-most of the figure), one may still
distinguish that the image is for a dog rather than other objects.

Furthermore, this problem may become more challenging due to the limited size
of medical image data. The positive pairs generated by improper augmentation
contribute to incorrect labels (or noisy labels in machine learning
literature) during the contrastive pre-training process. These errors or noisy
labels can be addressed with bigger datasets [11] and more network parameters
[12, 13, 14]. However, it is not always the case for medical images where the
size of the dataset is relatively small and the application is complicated.

Contrastive learning does not show “close-to-supervised” performance on
medical images as it does on natural images, which is partially due to the
above-mentioned reasons. To evaluate the representation power and the
contrastive learning performance, the linear probing protocol is widely used
[1, 2]. That is, a simple linear classifier is trained by freezing the
backbone pre-trained through contrastive learning, and the accuracy attained
by the linear classifier is used as a measure for the quality of the self-
supervised contrastive representations. With the linear probing protocol,
there typically exists a performance gap between the above self-supervised
learning paradigm and the conventional fully-supervised learning (i.e., by the
end-to-end training of all the parameters of the entire network, instead of
only the final linear classifier). The gap is noticeably bigger on medical
images, very different from natural images. For example, the performance
difference is about 5% when applying Momentum Contrast (MoCo) to ImageNet [2];
however, it expands to nearly 10% when MoCo is applied to CheXpert (i.e., the
chest X-ray dataset) [15, 16].

We thus propose a hypothesis that semantic unawareness in image augmentation
degrades the quality of contrastive learning and accounts for the above
performance gap. We design controlled experiments on the publicly available
Knee Osteoarthritis Severity Grading Dataset (OAI) of X-ray images in Section
III-B. We show that, for the exemplar augmentation operations of cutout and
random resized crop, the sweet spot (i.e., the optimal augmentation setting)
is weaker in strength for medical images than the commonly used settings in
natural images. And the performance starts to drop earlier if the augmentation
strength exceeds the sweet spot and becomes too much aggressive for the pre-
training. To this end, we conclude that the augmentation recipe currently
popularly used in natural images may not be suitable for medical images.

Further, this paper aims to propose attention-based augmentation to preserve
image semantics that is critical to downstream diagnosis tasks. We argue that
optimal contrastive learning should follow the human visual system, i.e., not
to lose the salient parts in the images that provide critical visual cues to
clinical diagnoses. We thus design FocusContrast, an augmentation strategy
guided by radiologists’ vision, to replace the current handcrafted strategies
for random image augmentation. The expert visual attention is collected by the
eye-tracker attached to the monitor. The gaze collection is seamless and
effortless in the daily clinical workflow [17]. Supervised by radiologists’
gaze, we first learn to predict visual attention (or the gaze map) when a
human expert reads a medical image. Based on the predicted gaze map,
FocusContrast can filter out improper data augmentation, i.e., the
transformation that possibly corrupts the salient part of the image.

The implementation of FocusContrast is straightforward. Specifically, the
learned augmentations can be summarized into two categories.

  * •

We propose the focus cutout and focus crop to only remove the non-salient area
of the image selectively. In this way, we can preserve the diseased
abnormalities in positive pairs, even though they are small and easy to lose
in previous random augmentation.

  * •

By learning from the radiologists who have the expertise in reading medical
images efficiently, the focus mask suppresses these non-salient areas. It
makes the network easier to distinguish the difference between the negative
pairs.

FocusContrast is plug-and-play and a framework-agnostic module for current
contrastive learning frameworks. Experimental results show that the proposed
semantic-aware augmentation consistently outperforms the conventional way of
random augmentation or using hand-tuned hyper-parameters in MoCo [2], SimCLR
[1], and BYOL [3].

This paper is organized in the following manner. In Section II, we briefly
cover the literature on contrastive learning and focus on the issue of
positive pairs particularly. We also introduce the relevant attempts to use
eye-tracking for radiology. In Section III, we design and analyze controlled
experiments to illustrate why the current augmentation strategy may not be
optimal for medical images. In Section IV, we present FocusContrast, which to
our knowledge is among the earliest works to utilize human visual attention to
guide self-supervised learning. We will conclude this paper with more
discussion in Section V. Our code is publicly
available111https://github.com/JamesQFreeman/MICEYE..

##  II Background

In this section, we first review self-supervised learning and contrastive
learning for medical images. Then, we introduce the related works about
radiologists’ gaze and visual attention.

###  II-A Contrastive Learning and its Usage to Medical Images

Contrastive learning has emerged as the front-runner for self-supervised
learning, which has attained superior performance on downstream tasks. These
pre-trained networks appended with linear probing can surpass the same fully-
supervised networks, while only the appended linear layer is trainable [18].
Researchers also use linear probing to inspect the quality of the learned
representations, i.e., by fixing the extracted features of the pre-trained
models and evaluating their performance in the subsequent tasks. In general,
large-scale contrastive pre-training with linear probing has become popular
due to generalizability to many scenarios and robustness against overfitting
[19].

There are several attempts to utilize contrastive pre-training in medical
image analyses. Sowrirajan et al. [15] proposed MoCo-CXR, which was an
adaptation of the contrastive learning method of MoCo, to produce the models
with better representations and initializations for abnormality detection in
chest X-rays. Azizi et al. [16] utilized multiple images per patient to
construct more informative positive pairs for self-supervised learning. With
their proposed self-supervision paradigm, the network was able to outperform
the supervised baseline. Zhou et al. [20] presented comparing-to-learn (C2L),
providing pre-trained 2D deep models for radiograph-related tasks from
unannotated data. All these works have adopted a conventional augmentation
strategy similar to SimCLR [1].

In the early days of contrastive learning, high-quality representation
requires a large number of negative pairs in a batch (e.g. 4096). SimCLR [1]
has demonstrated that contrastive learning benefits from larger batch size,
which allows for more negative pairs. In order to increase the batch size,
Momentum Contrast (MoCo) [2] updates a momentum memory bank of negative
representations to get rid of GPU memory restriction. However, more recently,
negative pairs are shown to be less necessary for representation learning.
Several works show that the number of negative pairs may have a limited
influence on the representation quality when the learning framework is
designed in a sophisticated way. For example, BYOL [3] uses two Siamese
networks and an extra non-linear transform, which enables contrastive learning
with fewer negative samples in a single batch. DINO [5] has further extended
this idea and utilized the transformer architecture, with a small batch size
of 8. As a summary, recent researches tend to pay more attention to positive
pairs in contrastive learning. And in this paper, we will also focus on their
roles in terms of the quality of the learned representations.

###  II-B Positive Pair Augmentation in Contrastive Learning

One of the major design choices in contrastive learning is how to generate
positive pairs. An intuitive approach is to create different views from the
same sample using augmentation. And most contrastive learning frameworks apply
augmentation by adapting from the conventional way in supervised learning [1,
2, 3, 21, 22]. Chen et al. [1] comprehensively studied the augmentation
operations. They found that more aggressive augmentation was usually needed in
contrastive pre-training than supervised learning. Tian et al. [8] proposed an
InfoMin principle to catch a sweet spot of mutual information between
augmented views, which indicated the maximal learning performance.

There are also recent works for semantic-aware augmentation. Selvaraju et al.
[23] proposed CAST to use unsupervised saliency maps to sample the cropping.
Peng et al. [24] proposed ContrastiveCrop for the augmentation of semantic-
aware cropping. However, as we will demonstrate in the experiment section,
medical images have very different semantic patterns compared to natural
images, implying that those methods may not be suitable.

###  II-C Eye-tracking in Radiology

Visual attention is a useful tool to understand and interpret radiologists’
reasoning and clinical decision. In 1981, Carmody et al. [25] published one of
the first eye-tracking studies in the field of radiology, where they studied
the detection of lung nodules in chest X-ray films. Four radiologists
participated and their eye movements were recorded using special glasses based
on the corneal reflection technique. The participating radiologists were
instructed to press a key when they found a nodule in the X-ray film. The
study found that the false-negative errors in reading the X-ray data could be
impacted by the eye-scanning strategies used by individual radiologists.

Eye-tracking studies are also conducted on other specialties such as
mammography. Kundel et al. [26] gathered eye-tracking data and found that 57%
of cancer lesions were located within the first second of viewing. Voisin et
al. [27] investigated the association between gaze patterns and diagnostic
performance for lesion detection in mammograms. It is found that gaze
fixations are highly correlated with radiologists’ diagnostic errors.

There are also studies that focus on volumetric CT and MRI images. Bertram et
al. [28] investigated the image markers of visual expertise using abdominal
CT. The eye-tracking data shows that specialists react with longer fixations
and shorter saccades when encountering the presence of lesions. Mallett et al.
[29] focused on CT colonography videos, which were interpreted by 27
experienced radiologists and 38 inexperienced radiologists. The eye-tracking
data indicates that experienced readers have higher rates of polyp
identification than inexperienced readers as evidenced by multiple pursuits
when examining polyps. Stember et al. [30] used eye-tracking data to label
brain tumors in MRI scans. More related studies are surveyed in [31, 32].

Figure 3: Relationship between the linear-probing classification performance
and the augmentation strength in knee X-ray OA diagnosis. Random resized crop
and random cutout are investigated. In the left panel, we report the linear-
probing accuracy curves for 10% label fraction (578 image-label pairs) and
100% label fraction (5778 image-label pairs), with varying augmentation
strength. The yellow areas in the curves denote the common augmentation
settings in the natural image. In the right panel, we offer visual examples
corresponding to different augmentation strength.

##  III What Are the Optimal Contrastive Views for Medical Image?

In this section, we first recall the InfoMin principle [8], which offers an
explanation of why different tasks need different augmentation. Then, we
illustrate this issue by using knee X-rays as an example, which is a common
diagnostic examination in clinical practice.

###  III-A The Sweet Spot for Contrastive Augmentation

When adopting contrastive learning for pre-training, two views $v_{1}$ and
$v_{2}$ augmented from an image $x$ are often used to get $z_{1}$ and $z_{2}$.
The objective of contrastive self-supervised learning maximizes the lower
bound on the mutual information (MI) estimate between the two views, or
$I(v_{1};v_{2})$ [33, 34]. However, in a downstream task such as
classification for disease diagnosis, the objective is very different from the
above in contrastive learning. For the classification task to predict the
label $y$, the optimal representation $z^{*}$ can be encoded from the input
image $x$. Ideally, a model built with $z^{*}$ has all the information
necessary to predict $y$ as accurately as accessing $x$ directly, or
$I(z^{*};y)=I(x;y)$. That is, no information loss occurs in mapping $x$ to $z$
concerning the task of predicting the label $y$ under the minimal sufficient
statistic assumption.

The augmented views are critical to encoding the representations, which focus
on the shared information between the views in contrastive learning [35]. If
the shared information between a pair of weakly augmented images is too much,
or $I(v_{1};v_{2})>I(z^{*};y)$, then the captured representations cannot
discard enough nuisance information from the inputs [8]. That is, the code of
$z$ carries many task-independent noises, which affect performance in the
subsequent task due to reduced signal-to-noise ratio of $z$. On the other
hand, if the key information to solve the classification task is missing,
e.g., $I(v_{1};v_{2})<I(z^{*};y)$, due to excessively strong augmentation, the
code of $z$ also degrades the classification performance. As a result, Tian et
al. [8] defined the sweet spot in terms of the downstream performance where
the useful information between the views is properly preserved following

$I(v_{1};v_{2})=I(z^{*};y).$ (1)

We further argue that, for many classification tasks in medical images, the
measure of $I(z^{*};y)$ can be high whereas the diseased abnormalities are
often tiny. Note that $z$ is usually computed as a spatial integration of
localized features, e.g., by pooling from the convolutional network outputs.
Thus, the entropy of $z$ conditioned on the label $y$ is small, since $z$ is
dependent on restricted spatial locations that are abnormal and thus
associated with the disease diagnosis. Meanwhile, as the conditional entropy
of $z$ upon $y$ is an inverse to $I(z^{*};y)$, we conclude that $I(z^{*};y)$
is high particularly for the medical images such as in Fig. 2.

Therefore, it leads to our assumption that augmentation strength should be
relatively weaker corresponding to the sweet spot for medical images than the
common settings in natural images. Referring to equation (1), a too strong
augmentation may reduce mutual information between $v_{1}$ and $v_{2}$
significantly, which induces information loss to the learned representations
and affects the downstream classification.

To make sure that $I(v_{1};v_{2})$ is kept high throughout image augmentation,
besides controlling the augmentation strength, it is necessary to preserve the
region of the diseased abnormality in the image. In this way, $I(z_{1};z_{2})$
should be close to $I(v_{1};v_{2})$, as the mapping from the input image to
its latent representation may suffer little from the information loss. In this
way, one can guarantee the key information to pass to the downstream
classification task. Meanwhile, the redundant features that are influence
subsequent task are removed from learned representations.

###  III-B Augmentation Recipes for Knee X-rays

We investigate two widely-adopted augmentations to illustrate the shifting
sweet spots by using knee X-rays as the example. We pay special attention to
two operations of random resized crop (to crop a random portion of the image
and then resize the cropped image to the original size) and random cutout (to
randomly remove a portion of the image) [36]. The two operations could affect
the learned representation, since they could easily remove the semantic cues
in the images when the diseased abnormalities are small. We first briefly
introduce our settings and then present the experimental results.

Dataset and follow-up task. The knee X-ray images are acquired from the OAI
repository, which is made publicly available by [37]. The dataset used in this
paper includes 5778 training images and 1656 testing images. In the
contrastive pre-training stage, all the training images are used. Then, to
evaluate the learned representations of the contrastive pre-training, the
follow-up linear probing task is trained with “10% label fraction” and “ 100%
label fraction” setting. That is, we use 578 and 5778 training images and
their class labels to train the classifier. Then, we apply the trained
classifier to the test images, and yield their labels prediction in five
classes. Specifically, the knee OA severity is measured by the Kellgren-
Lawrence (KL) grading system [38], which is a 5-point semi-quantitative
progressive ordinal scale ranging from Grade 0 (normal) to 4 (high severity).
We use the KL scores as target of the classification task.

Contrastive learning framework. We choose BYOL as the contrastive learning
framework here. Specifically, We train at batch size 128 for 100 epochs on
Nvidia RTX Titan. Furthermore, we use linear warm-up for first 10 epochs, and
decay the learning rate with cosine decay schedule. Regarding the loss, we use
InfoNCE optimized by LARS optimizer with the learning rate of 0.2 and weight
decay of $1.5e^{-6}$. We use ResNet-50 [39] as our backbone. Our experimental
settings may not ensure the best possible performance (e.g., bigger backbone,
longer training and larger batch size that can further boost the performance),
but they do allow us to validate our hypothesis.

Although we are interested in the augmentations of random resized crop and
random cutout, we also deploy other popular image augmentations for
contrastive learning. The reason is that, if only random resized crop or
random cutout is used, the strictly controlled augmentation setting would be
too far away from the current real case. To this end, in our experiments, we
resize all images to 224$\times$224, followed by random flip, random color
distortion, random rotation (30 degrees), random cutout (32px$\times$32px) and
random resized crop (1.2$\times$ zoom-in) by default. When we investigate the
impact of random resized crop for example, we will change it from the default
to our specific setting, while keeping other augmentation settings unchanged.

Results for Random Resized Crop. In Fig. 3 Row 1, we report the results when
random resized crop varies its strength. Specifically, we test 0.6$\times$,
0.8$\times$, 1$\times$, 1.2$\times$, 1.4$\times$, 1.6$\times$, and 2$\times$
for random resized crop. Given the 1.4$\times$ setting as an example, the
224$\times$224 image is first resized to 313$\times$313, and then cropped back
to a 224$\times$224 output with a randomly chosen center location. The
accuracy curve for the “578 image-label pairs” shows a peak around the
1.0$\times$-1.2$\times$ settings. And accuracy starts to drop for stronger
augmentation such as 1.4$\times$. By comparison, the sweet spot is usually set
to 2$\times$ for natural image in contrastive learning, which is significantly
stronger than the optimal augmentation strength acquired for knee X-rays in
our experiment. As a reference, we also provide “5778 image-label pairs” curve
in the figure. And the same trend can be observed, implying that adding more
supervision to the classification task cannot easily compensate for the bad
representations caused by improper augmentation to contrastive pre-training.

In the top-right of the figure, we show examples of the augmented images with
varying strength of random resized crop. One may easily notice that the images
change their appearances significantly when only random resized crop is
enabled and its strength is increasing. While other augmentations are also
turned on as in our real experiments, the effect of random resized crop is
still clear visually. For example, at 2$\times$, the OA-related region of the
femur-tibia interface is cropped out, implying that the augmentation is
excessively strong. Meanwhile, one may notice that 0.6$\times$-0.8$\times$
settings can shrink the original images by padding zero-valued background. In
this way, the very weak augmentation induces much distortion in the augmented
images, which also echos the relatively low classification performance
achieved by linear probing.

Results for Random Cutout. In Fig. 3 Row 2, we report the linear-probing
accuracy when random cutout varies its strength. Specifically, we test 8px,
16px, 32px, 48px, 64px, 96px, 128px, as the value indicates the size of the
cutout region. One can notice that random cutout larger than 64px (i.e.,
larger than 8.2% of the whole image) will degrade the representations learned
from the Siamese contrastive learning. In comparison, random cutout in the
latest works like SwinTransformer [40] and ConvMixer [41] can be tuned up to
1/3 of the field-of-views for natural images, which is much larger than our
finding in the knee X-ray images.

In the right half of the figure, we show examples of augmented images with
varying cutout sizes. One may easily observe that the knee joint, which is
critical for diagnosis, can be masked by a large cutout area when the strength
of random cutout only is creasing. Despite turning on other augmentations like
random resized crop (1.2$\times$ in our example) will increase the area of the
diseased abnormality, random cutout still masks out too much. For example, at
96px and 128px, the knee joint and cartilage are masked out, resulting in
potentially incorrect positive pairs in contrastive training.

Conclusion. We conduct controlled experiments on two widely-used data
augmentations of random resized crop and random cutout by testing different
augmentation strength. Particularly for knee X-ray images, we empirically find
the optimal augmentation hyper-parameters, which are (1$\times$,1.2$\times$)
for random resized crop and (16px, 48px) for random cutout. Both settings are
significantly weaker than those in natural images. The results demonstrate
that aggressive augmentations noticeably degrade the representations from
contrastive learning. This is in line with our expectation since (1) subtle
lesions can be removed by aggressive random augmentation, and (2) the
transformed image is out of distribution and thus cannot benefit the follow-up
task.

Figure 4: The overview of FocusContrast augmentation framework. In (a), a
network called gaze map predictor is trained to predict the radiologist’s gaze
map. In (b), when an image is augmented into different views, the gaze map
predictor helps distinguish whether the salient parts in the images are
preserved or not. If the salient part is preserved, the augmented view will be
fed to the encoder for contrastive learning; otherwise, it will be treated as
a false-positive view and discarded.

##  IV Learning Views from Radiologists for Contrastive Learning

In Section III, we demonstrate that it is not appropriate to simply borrow the
contrastive augmentation transforms in natural images and then apply them to
medical images. And we conduct extensive experiments on two popular
augmentations by searching for their optimal hyper-parameters in brutal force.
However, an exhaustive search over all possible hyper-parameters is not always
a good solution in that it is time-consuming and expensive.

To solve this problem, we propose a straightforward yet efficient method to
learn the semantic-aware augmentations from the visual attention of the
radiologists for contrastive views. As shown in Fig. 4, we first collect
hundreds of gazes from radiologists, and train a gaze map predictor. Then, we
use the trained gaze map predictor to judge each augmented view and decide
whether the view is qualified to feed into contrastive learning. In
particular, the gaze map predictor tells which part in the (augmented) image
contains salient semantics according to the learned radiologists’ gaze. If the
salient part is preserved in the augmented view, it will be used in
contrastive learning; otherwise, it will be discarded.

In the next, we use the same images and task as the last section to compare
our learnable FocusContrast with the brutal-force ones. We will first
introduce how to collect the gaze data, followed by the training of the gaze
map predictor. Then we will introduce the way to filter out the inappropriate
views to eliminate false positives for contrastive learning. Finally, we will
demonstrate that our method can be highly effective in contrastive
representation learning.

###  IV-A Gaze Collection and Post-processing

Figure 5: (a) Radiologists look at a knee X-ray image until they make a
diagnosis decision. From the moment they are shown this image to the moment
they type down their decision, their eye movement is recorded by the eye-
tracker. (b) The eye movement is recorded discretely in coordinate-timestamp
form. The gaze points are first processed to remove the distortion. Then a
Gaussian kernel is applied to the gaze points, which generates the gaze map
(or the saliency heatmap). The image and the gaze map are used as the ground-
truth pair to train the gaze map predictor. Figure 6: On the left, we
demonstrate an image and its corresponding gaze map from a radiologist. On the
right, the image is augmented to generate the positive pairs. We present the
exemplar images after augmenting by FocusContrast (top row), and compare them
with the conventional counterparts (bottom row). We use the orange boxes to
indicate that the semantic cues (i.e., abnormal areas related to disease
diagnosis) are preserved in the augmented images. Contrarily, the red boxes
indicate the key areas that are missing or partially corrupted after
conventional augmentation.

We collect the eye-tracking data with the Tobii 4C remote eye-tracker that
records binocular gaze data at 90Hz. We implement customized data collection
software in Python using the manufacturer-provided SDK. The software is made
publicly available with this paper. Readers are seated in front of a 27-inch
LCD screen, to simulate the clinical working condition. They can adjust their
seat distances to the screen for comfort. Our software logs the reader’s on-
screen gaze locations with the corresponding timestamps.

Our data collection paradigm is illustrated in Fig. 5(a). Note that we make
the demonstration by reading knee X-ray images for OA assessment. First,
radiologists log in their profiles of identity. And then we calibrate the eye-
tracker using the standard 5-point calibration routine [42] for the
participating radiologists. Next, the radiologist follows a cycle of two
steps, namely reading and diagnosis, to complete the diagnosis task on an
X-ray image. During the reading step, an image is randomly drawn from our
training dataset and shown to the radiologist, as the radiologist reads the
image until they feel confident to reach a diagnosis decision. In the
diagnosis step, the radiologist types in the decision by pressing the number
keys on the keyboard, e.g., “1-4” used for representing KL-Grades 1-4, and
“Enter” for normal or Grade 0. We record eye movement during the entire cycle
of reading and diagnosis steps. We have a rest of 2 minutes for every 20
images to reduce radiologists’ fatigue. In this paper, we use 354 gaze tracks,
containing 154 images of Grade 0, 55 images of Grade 1, 81 images of Grade 2,
40 images of Grade 3, and 24 images of Grade 4.

###  IV-B Training of the Gaze Map Predictor

The training procedure is demonstrated briefly in Fig. 5(b). As the figure
shows, the collected gaze point sequence is filtered by a post-processing
algorithm [17] to remove the gaze points when the radiologist is looking at
the GUI elements or other distractions. We further acquire the gaze map from
the filtered gaze points per training image, by modeling them in the Gaussian
mixtures. The Gaussian kernel is sized to 99, determined by the human para-
central vision area [43], image size ($25\times 25cm^{2}$) and distance from
eyes to screen ($\approx 50cm$) [17].

The X-ray image and the gaze map are used as the input and output to supervise
the training of a U-Net [44]. The trained U-Net of gaze map predictor can
predict where radiologists will pay attention when reading a knee X-ray image.
The U-Net consists of 4 downsample encoders and 4 upsample decoders, and has
the same size for the input and output. The network is trained using Adam [45]
for 30 epochs and the weight of the last epoch is used in our experiment.
Empirically, a minor mistake of the gaze map predictor will not influence the
follow-up view-filtering process as we only need it to give us a rough
location of where the radiologist would pay attention to.

###  IV-C Generating Image Views using Gaze Map Predictor

The intuition of our method is simple – when radiologists think the specific
area of the image is informative, these areas should not be lost during the
augmentation to prevent potential information loss. Based on this idea, we
propose three novel augmentations to generate the positive views using our
trained gaze map predictor: 1) focus cutout, 2) focus crop, and 3) focus mask.
These three augmentations are designed based on the classical cutout, random
resized crop and random mask, respectively. See Fig. 6 for visual examples of
how these augmentations are done.

Focus cutout. Cutout is a commonly used augmentation technique that improves
the robustness and overall performance of networks [36, 6]. It randomly masks
out square regions of the input, which could accidentally remove the important
semantic information, e.g. the diseased abnormality. To prevent such loss, we
propose focus cutout, which only randomly masks out regions that radiologists
may not pay attention to.

The implementation of focus cutout is detailed in Algorithm 1. Specifically,
we predict the gaze map for an input image $x$. Follwing a random cutout, we
get an augmented view $v_{1}$ as well as its corresponding gaze map. Then, we
calculate the intersection (i.e., IOU) between the non-zero gaze-map regions
before and after random cutout. If the IOU scores for two augmented views
($v_{1}$ and $v_{2}$) are both above a threshold (e.g., 0.9), the two views
are preserved as a valid positive pair.

Algorithm 1 Example of focus cutout

$x$ is an input image

network for gaze map predictor

$gaze$ = Predictor($x$)

do$\triangleright$ Same cutout for both views

$v_{1}$, $gaze_{v1}$ = random_cutout([$x$, $gaze$])

$v_{2}$, $gaze_{v2}$ = random_cutout([$x$, $gaze$])

while IOU($gaze_{v1}$,$gaze$)$>0.9$ & IOU($gaze_{v2}$,$gaze$)$>0.9$

return $v_{1}$, $v_{2}$$\triangleright$ A positive pair after augmentation

Focus crop. In SimCLR, random resized crop is performed for “context
aggregation” and “neighboring view prediction”. However, an aggressive resized
crop ($\geq$1.4$\times$ in Fig. 3) may remove the critical abnormality. In
Fig. 6, random resized crop also has lost the abnormality in the knee joint
(indicated by red boxes). To address this problem, we propose focus crop. Our
implementation is similar to focus cutout – zoom-in resizing is only performed
when it does not remove too much area of radiologists’ interest. In
particular, the intersection threshold is set to 0.8 for focus crop.

Focus mask. The idea of using dropout or random mask [46] for augmentation is
natural and applicable in self-supervised learning. However, while this simple
idea has achieved success in natural language processing [47], it is not
common in computer vision until very recently [13]. In this paper, we propose
focus mask, which masks the area that the reader is not much paying attention
to. A threshold is applied to the gaze map to mask out 80% less-informative
areas.

###  IV-D Experimental settings

We inherit the knee X-ray dataset and the KL-grading task from Section III-B.
In our experiments, following [1, 2, 3], we use ResNet50 as our backbone. If
not otherwise specified, our models here are trained for 200 epochs with batch
size of 512. We use MMSelfSup [48] framework to pre-train our backbone and
adopt the recommended training recipe. During pre-training, all 5778 images in
the training set are used, which is also the same with Section III.

To evaluate the effectiveness of the pre-trained representations, we freeze
the backbone model and train a linear classifier on top using the labeled data
(denoted as linear-probe in our experimental results). In addition, we can
unfreeze all layers and fine-tune the entire model end-to-end, to compare
transferability on the overall performance (denoted as fine-tune). Both the
linear-probing and the fine-tuning schemes under comparison are trained for
100 epochs with batch size of 128.

We have validated five different sizes of the classification training sets: 1%
label fraction (55 images evenly across 5 classes contributing their labels to
train the classification task), 5% label fraction (288 images), 10% label
fraction (577 images), 20% label fraction (1154 images) and 100% label
fraction (5778 images). All evaluation is conducted on 1656 testing images,
and the best classification accuracy (ACC) is reported with the mean absolute
error (MAE) of the predicted grades at the same epoch.

###  IV-E Experimental results

Comparisons to other attention prediction models. In this research, we use
gaze to supervise the training for visual attention prediction. There are many
different ways to predict human visual attention. However, professional
radiologists have very different visual attention compared to non-
professionals. Therefore, the attention prediction models designed for the
general population and common visual tasks cannot function well in medical
images. To verify the above, in this experiment, we compare our gaze map
predictor with the commonly adopted methods of SpectralResidual (SR) [49],
Montabone et al. [50], and Pyramid Feature Attention Network (PFAN) [51]. Note
that PFAN is regarded as a state-of-the-art saliency prediction algorithm
trained on natural images, ranking top on PASCAL-S and DUT-OMRON [51].

The results are presented in Fig. 7. Our attention prediction model trained on
the radiologists’ gaze easily outperforms other models, which are designed for
the general-purpose human visual system. This is in line with our early
finding, as radiologists are trained to seek the areas with high disease-
related risks in an image and focus on them more [17].

In contrast, the general-purpose human visual system tends to look at highly
informative areas (e.g., with large intensity changes). And the compared
methods largely follow or mimic the human visual system as reflected by their
predicted attention maps. For example, SR [49] summarizes from natural image
statistics that intensity-frequency has a log-log relationship on average.
Then, the area out of this log-log distribution leads to a high response in
attention. On the other hand, PFAN uses a learning-based approach, training on
paired natural images and eye-tracking data [51]. The prediction, however,
cannot localize the critical areas effectively.

Comparisons to handcrafted augmentations. In this section, we compare the
learning-based FocusContrast with the handcrafted augmentations, which are
widely used for contrastive pre-training. We also compare with the optimal
augmentation hyper-parameters searched exhaustively (i.e., 1.2$\times$ for
random resized crop and 48px for cutout, see Section III-B). The three
settings are denoted by the superscripts Learned, Default, and Searched in
Table I, respectively. Note that we use BYOL as the common contrastive
learning framework for the comparisons here.

In addition, there are three baselines to compare in the table, including: a
randomly initialized ResNet-50 that is trained in the end-to-end manner
(denoted as from-scratch); a backbone fixed after ImageNet-based pre-training
and a linear-probing classifier (denoted as IN linear); a network pre-trained
by ImageNet and then fine-tuned end-to-end (denoted as IN fine-tune). To
investigate the training efficiency when utilizing classification labels, we
also test our models on different label fractions of the training data.

Figure 7: Comparison of different attention prediction models on knee X-ray
images. Red denotes high attention while blue denotes low attention. Our model
focuses more compactly on the diagnosis-related areas in the images.

There are several observations to derive from Table I and Fig. 8.

  * •

The pre-training procedure, especially contrastive learning, improves the
performance of the final classification task effectively. In the first row of
Table I, the accuracy of the network trained from-scratch is significantly
lower than other settings. Specifically, the ACC gap is +5.67% between from-
scratch and IN fine-tune when 100% label fraction is available; the gap
becomes even larger (+10.85% ACC in average) when the label fraction is
reduced to below 20%. Compared to pre-training upon ImageNet, contrastive
learning constantly improves the learned representations and the model’s
transferability over different label fractions of the training data. These
observations are in line with other recent literatures [16].

  * •

FocusContrast outperforms default contrastive learning in linear-probing
performance. FocusContrast (denoted as Learned, yellow line in Fig. 8 bottom
panel) constantly outperforms the widely adopted default random augmentations
(denoted as Default, green line in Fig. 8 bottom panel) by a large margin,
averaging in a +4.90% ACC gap. These results prove that FocusContrast helps to
learn a higher quality representation than the default random augmentation
does.

  * •

Compared to default augmentation, FocusContrast is more label-efficient to the
subsequent task. For example, linear-probeLearned at 1% label fraction (ACC:
39.07%) performs even better than linear-probeDefault at the 5% label fraction
(ACC: 36.05%). And linear-probeLearned at 20% label fraction (ACC: 56.39%)
performs better than linear-probeDefault at 100% label fraction (ACC: 55.81%).

  * •

As for transferability, FocusContrast (denoted as fine-tuneLearned, Row 6)
still introduces considerable improvement compared to the widely adopted
network parameters pre-trained from ImageNet (denoted as IN fine-tune, Row 2)
and contrastive pre-training by default augmentations (denoted as fine-
tuneDefault, Row 4).

  * •

FocusContrast has a similar performance as the exhaustively searched optimal
augmentation setting. On the lower panel of Fig. 8, FocusContrast (yellow
line) has a larger linear-probing performance lead (compared to searched
augmentation, orange line) when fewer labeled data is available. On the upper
panel of the fine-tune performance, the performance gap is smaller.

In summary, FocusContrast, being learning-based, outperforms other widely-
adopted approaches (ImageNet-based pre-training and contrastive learning with
conventional augmentations). It can even have a slightly better overall
performance than the exhaustively searched (and thus optimal) augmentation
setting. In that collecting gaze is seamless and effortless in our
configuration while searching hyper-parameters is time-consuming and
computationally intense, FocusContrast can be practical and useful to develop
many future computer-aided diagnosis systems.

TABLE I: OAI KL-Grade prediction performance of models trained with different number of labels. ResNet-50 is used as the backbone. | Label fraction
---|---
Method | 1% | 5% | 10% | 20% | 100%
| ACC | MAE | ACC | MAE | ACC | MAE | ACC | MAE | ACC | MAE
from-scratch | 26.03 | 1.370 | 37.62 | 1.216 | 39.79 | 1.163 | 39.25 | 1.189 | 58.10 | 0.614
IN fine-tune | 33.94 | 1.028 | 38.95 | 1.007 | 55.19 | 0.624 | 58.03 | 0.619 | 63.77 | 0.502
IN linear | 30.43 | 1.272 | 33.88 | 1.002 | 43.12 | 0.986 | 44.38 | 0.938 | 49.40 | 0.790
fine-tuneDefault | 34.84 | 0.980 | 42.61 | 0.958 | 56.52 | 0.606 | 57.55 | 0.624 | 65.70 | 0.445
fine-tuneSearched | 36.11 | 0.974 | 44.38 | 0.874 | 58.51 | 0.579 | 60.45 | 0.565 | 67.33 | 0.429
fine-tuneLearned | 37.01 | 0.930 | 45.17 | 0.908 | 60.45 | 0.559 | 61.54 | 0.557 | 67.09 | 0.434
linear-probeDefault | 32.85 | 1.053 | 36.05 | 1.097 | 50.57 | 0.755 | 51.85 | 0.763 | 55.81 | 0.685
linear-probeSearched | 35.75 | 1.018 | 38.40 | 1.026 | 53.44 | 0.697 | 55.62 | 0.651 | 58.93 | 0.586
linear-probeLearned | 39.07 | 0.961 | 40.76 | 0.872 | 54.71 | 0.658 | 56.39 | 0.649 | 59.66 | 0.583
Figure 8: Classification accuracy using different augmentation strategies.

Comparisons of different contrastive learning frameworks. As a plug-and-play
and framework-agnostic augmentation algorithm, FocusContrast can work on
different contrastive learning frameworks naturally. In this section, we
further test over three widely adopted contrastive learning frameworks,
namely, MoCo v2 [7] (denoted as MoCo), SimCLR [1], and BYOL [3] in Table II.
The superscript Learned indicates that we integrate the proposed FocusContrast
with the specific contrastive learning framework. We also compare to the
randomly initialized (random-init) and ImageNet pre-trained (ImageNet)
backbones in the table.

We observe that our FocusContrast improves over conventional random
augmentation by +2.72% (ACC) for MoCo, +1.76% for SimCLR, and +3.14% for BYOL,
respectively, when 10% label fraction is available to linear probing. A
similar and even larger improvement is also achieved when 100% label fraction
is used (i.e., +3.26%, +3.86%, +3.50%, respectively). The consistent gains
show that our FocusContrast is effective and generalizes well with respect to
existing contrastive learning frameworks. And the performance gains become
more significant when compared to non-contrastive pre-training methods (such
as random-init and ImageNet), which are in line with our early experimental
results.

TABLE II: Linear Evaluation of our method and other contrastive learning frameworks | Label fraction | Label fraction
---|---|---
Pretrain method | 10% | 100%
| ACC | MAE | ACC | MAE
random-init | 38.65 | 1.244 | 39.13 | 1.216
ImageNet | 43.12 | 0.986 | 44.44 | 0.930
MoCo | 50.84 | 0.746 | 53.68 | 0.674
MoCoLearned | 53.56 | 0.701 | 56.94 | 0.627
SimCLR | 44.25 | 0.962 | 49.58 | 0.781
SimCLRLearned | 46.01 | 0.878 | 53.44 | 0.697
BYOL | 51.57 | 0.755 | 55.31 | 0.685
BYOLLearned | 54.71 | 0.658 | 58.81 | 0.607

Ablation study on the proposed augmentation operators. To further verify the
contributions of the three proposed image augmentations, we investigate the
performance of linear probing when applying the augmentations individually or
in pairs. Three baselines are here for comparison in Table III, including
ImageNet pre-trained (IN pre-trained), BYOL, and BYOL with the optimal
augmentation setting exhaustively searched (BYOLSearched). Our proposed
augmentations are added to the baseline of BYOL.

Table III shows the linear probing results using 10% and 100% label fractions,
respectively. For 10% label fraction, focus mask improves the linear probing
accuracy from 51.57% to 53.62%. Also, when only focus crop is applied, the
linear probing accuracy is boosted to 53.44%. In comparison, focus cutout
brings in a smaller improvement compared to the other two operators. Similar
conclusions are also reported recently in masked image model self-supervised
learning [13]. We believe that it is because the cutout only removes a small
part of the image, which makes the contrastive learning task too easy compared
to the other two augmentations. As a summary, by adding one or more
augmentations to our FocusContrast, the capability of contrastive learning
continuously improves as reflected by our experimental results.

TABLE III: Ablation study of three proposed augmentation. This table reports the linear evaluation performance. Focus Crop is abbreviated as F Crop and the same applies for Focus Cutout and Focus Mask. | Label fraction | Label fraction
---|---|---
Baseline Methods | 10% | 100%
| ACC | MAE | ACC | MAE
IN pre-trained | 43.12 | 0.986 | 44.44 | 0.930
BYOL | 51.57 | 0.755 | 55.31 | 0.685
BYOLSearched | 53.44 | 0.697 | 58.93 | 0.586
FocusContrast Opertators |  |  |  |
F Crop | F Cutout | F Mask |  |  |  |
|  | $\checkmark$ | 53.62 | 0.699 | 58.81 | 0.607
$\checkmark$ |  |  | 53.44 | 0.696 | 56.94 | 0.620
| $\checkmark$ |  | 52.16 | 0.723 | 55.89 | 0.652
$\checkmark$ |  | $\checkmark$ | 53.56 | 0.701 | 56.64 | 0.640
$\checkmark$ | $\checkmark$ |  | 53.08 | 0.695 | 56.64 | 0.636
| $\checkmark$ | $\checkmark$ | 54.12 | 0.674 | 59.43 | 0.603
$\checkmark$ | $\checkmark$ | $\checkmark$ | 54.71 | 0.658 | 59.66 | 0.583

##  V Conclusion and Discussion

In this paper, we start by questioning whether the current augmentation
strategies popular in general-purpose contrastive learning may hurt the
representations learned for medical images. Then we demonstrate that the
strong augmentation that works well in natural images cannot fit the pre-
training of the medical images, because the small semantic-related areas can
be easily corrupted to degrade the learned representations.

Moreover, we propose three learnable augmentation operations guided by
radiologists’ gaze. Particularly, we use eye-tracker to record the experts’
gaze and train a visual attention prediction network to help filter out
inappropriate views. The proposed FocusContrast boosts the network’s
performance significantly in the follow-up tasks compared to default
augmentation settings. Our solution requires no extra code or package and
works with different contrastive learning frameworks naturally

Our paper shows the feasibility to learn augmentation paradigms from human
experts instead of handcrafting them when applying contrastive learning to
unseen domains. In the future, we will further extend the gaze-guided approach
to more self-supervised frameworks such as masked image models, to guide the
network to learn from unlabeled images more efficiently.

There are limitations of this work which we plan to further address in future
works. First, using a gaze map predicting network to predict the visual
attention may potentially slow down the augmentation process. The current
workaround is calculating and storing a salience map for each image before
contrastive pre-training. Then, the do-while statement (Algorithm 1, Line 4)
introduces branches thus making the transform hard to be computed in parallel.
Finally, the knee X-ray image is relatively a simple medical image. We are
planning to extend this idea to more 2D/3D imaging modalities and clinical
applications.

## References

  * [1] T. Chen, S. Kornblith, M. Norouzi, and G. Hinton, “A simple framework for contrastive learning of visual representations,” in _International conference on machine learning_. PMLR, 2020, pp. 1597–1607.
  * [2] K. He, H. Fan, Y. Wu, S. Xie, and R. Girshick, “Momentum contrast for unsupervised visual representation learning,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2020, pp. 9729–9738.
  * [3] J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. H. Richemond, E. Buchatskaya, C. Doersch, B. A. Pires, Z. D. Guo, M. G. Azar _et al._ , “Bootstrap your own latent: A new approach to self-supervised learning,” _arXiv preprint arXiv:2006.07733_ , 2020.
  * [4] X. Chen and K. He, “Exploring simple siamese representation learning,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2021, pp. 15 750–15 758.
  * [5] M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin, “Emerging properties in self-supervised vision transformers,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2021, pp. 9650–9660.
  * [6] Z. Zhong, L. Zheng, G. Kang, S. Li, and Y. Yang, “Random erasing data augmentation,” in _Proceedings of the AAAI Conference on Artificial Intelligence_ , vol. 34, no. 07, 2020, pp. 13 001–13 008.
  * [7] X. Chen, H. Fan, R. Girshick, and K. He, “Improved baselines with momentum contrastive learning,” _arXiv preprint arXiv:2003.04297_ , 2020.
  * [8] Y. Tian, C. Sun, B. Poole, D. Krishnan, C. Schmid, and P. Isola, “What makes for good views for contrastive learning,” _CoRR_ , vol. abs/2005.10243, 2020\. [Online]. Available: https://arxiv.org/abs/2005.10243
  * [9] T. Xiao, X. Wang, A. A. Efros, and T. Darrell, “What should not be contrastive in contrastive learning,” _arXiv preprint arXiv:2008.05659_ , 2020.
  * [10] R. Dangovski, L. Jing, C. Loh, S. Han, A. Srivastava, B. Cheung, P. Agrawal, and M. Soljačić, “Equivariant contrastive learning,” _arXiv preprint arXiv:2111.00899_ , 2021.
  * [11] D. Rolnick, A. Veit, S. Belongie, and N. Shavit, “Deep learning is robust to massive label noise,” _arXiv preprint arXiv:1705.10694_ , 2017.
  * [12] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell _et al._ , “Language models are few-shot learners,” _Advances in neural information processing systems_ , vol. 33, pp. 1877–1901, 2020.
  * [13] K. He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick, “Masked autoencoders are scalable vision learners,” _arXiv preprint arXiv:2111.06377_ , 2021.
  * [14] Z. Li, Z. Cui, S. Wang, Y. Qi, X. Ouyang, Q. Chen, Y. Yang, Z. Xue, D. Shen, and J.-Z. Cheng, “Domain generalization for mammography detection via multi-style and multi-view contrastive learning,” in _International Conference on Medical Image Computing and Computer-Assisted Intervention_. Springer, 2021, pp. 98–108.
  * [15] H. Sowrirajan, J. Yang, A. Y. Ng, and P. Rajpurkar, “Moco-cxr: Moco pretraining improves representation and transferability of chest x-ray models,” in _Proc. International Conference on Medical Imaging with Deep Learning (MIDL)_ , 2021.
  * [16] S. Azizi, B. Mustafa, F. Ryan, Z. Beaver, J. Freyberg, J. Deaton, A. Loh, A. Karthikesalingam, S. Kornblith, T. Chen _et al._ , “Big self-supervised models advance medical image classification,” _arXiv preprint arXiv:2101.05224_ , 2021.
  * [17] S. Wang, X. Ouyang, T. Liu, Q. Wang, and D. Shen, “Follow my eye: Using gaze to supervise computer-aided diagnosis,” _IEEE Transactions on Medical Imaging_ , 2022.
  * [18] N. Tomasev, I. Bica, B. McWilliams, L. Buesing, R. Pascanu, C. Blundell, and J. Mitrovic, “Pushing the limits of self-supervised resnets: Can we outperform supervised learning without labels on imagenet?” _arXiv preprint arXiv:2201.05119_ , 2022.
  * [19] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark _et al._ , “Learning transferable visual models from natural language supervision,” in _International Conference on Machine Learning_. PMLR, 2021, pp. 8748–8763.
  * [20] H.-Y. Zhou, S. Yu, C. Bian, Y. Hu, K. Ma, and Y. Zheng, “Comparing to learn: Surpassing imagenet pretraining on radiographs by comparing image representations,” in _International Conference on Medical Image Computing and Computer-Assisted Intervention_. Springer, 2020, pp. 398–407.
  * [21] E. D. Cubuk, B. Zoph, D. Mane, V. Vasudevan, and Q. V. Le, “Autoaugment: Learning augmentation policies from data,” _arXiv preprint arXiv:1805.09501_ , 2018.
  * [22] E. D. Cubuk, B. Zoph, J. Shlens, and Q. V. Le, “Randaugment: Practical automated data augmentation with a reduced search space,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops_ , 2020, pp. 702–703.
  * [23] R. R. Selvaraju, K. Desai, J. Johnson, and N. Naik, “Casting your model: Learning to localize improves self-supervised representations,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2021, pp. 11 058–11 067.
  * [24] X. Peng, K. Wang, Z. Zhu, and Y. You, “Crafting better contrastive views for siamese representation learning,” _arXiv preprint arXiv:2202.03278_ , 2022\.
  * [25] D. P. Carmody, C. F. Nodine, and H. L. Kundel, “Finding lung nodules with and without comparative visual scanning,” _Perception & psychophysics_, vol. 29, no. 6, pp. 594–598, 1981.
  * [26] H. L. Kundel, C. F. Nodine, E. A. Krupinski, and C. Mello-Thoms, “Using gaze-tracking data and mixture distribution analysis to support a holistic model for the detection of cancers on mammograms,” _Academic radiology_ , vol. 15, no. 7, pp. 881–886, 2008.
  * [27] S. Voisin, F. Pinto, S. Xu, G. Morin-Ducote, K. Hudson, and G. D. Tourassi, “Investigating the association of eye gaze pattern and diagnostic error in mammography,” in _Medical Imaging 2013: Image Perception, Observer Performance, and Technology Assessment_ , vol. 8673. International Society for Optics and Photonics, 2013, p. 867302\.
  * [28] R. Bertram, J. Kaakinen, F. Bensch, L. Helle, E. Lantto, P. Niemi, and N. Lundbom, “Eye movements of radiologists reflect expertise in ct study interpretation: A potential tool to measure resident development,” _Radiology_ , vol. 281, no. 3, pp. 805–815, 2016.
  * [29] S. Mallett, P. Phillips, T. R. Fanshawe, E. Helbren, D. Boone, A. Gale, S. A. Taylor, D. Manning, D. G. Altman, and S. Halligan, “Tracking eye gaze during interpretation of endoluminal three-dimensional ct colonography: visual perception of experienced and inexperienced readers,” _Radiology_ , vol. 273, no. 3, pp. 783–792, 2014.
  * [30] J. N. Stember, H. Celik, D. Gutman, N. Swinburne, R. Young, S. Eskreis-Winkler, A. Holodny, S. Jambawalikar, B. J. Wood, P. D. Chang _et al._ , “Integrating eye tracking and speech recognition accurately annotates mr brain images for deep learning: Proof of principle,” _Radiology: Artificial Intelligence_ , vol. 3, no. 1, p. e200047, 2020.
  * [31] T. T. Brunyé, T. Drew, D. L. Weaver, and J. G. Elmore, “A review of eye tracking for understanding and improving diagnostic interpretation,” _Cognitive research: principles and implications_ , vol. 4, no. 1, pp. 1–16, 2019.
  * [32] C.-C. Wu and J. M. Wolfe, “Eye movements in medical image perception: a selective review of past, present and future,” _Vision_ , vol. 3, no. 2, p. 32, 2019.
  * [33] B. Poole, S. Ozair, A. Van Den Oord, A. Alemi, and G. Tucker, “On variational bounds of mutual information,” in _International Conference on Machine Learning_. PMLR, 2019, pp. 5171–5180.
  * [34] S. Soatto and A. Chiuso, “Visual representations: Defining properties and deep approximations,” _arXiv preprint arXiv:1411.7676_ , 2014.
  * [35] A. v. d. Oord, Y. Li, and O. Vinyals, “Representation learning with contrastive predictive coding,” _arXiv preprint arXiv:1807.03748_ , 2018\.
  * [36] T. DeVries and G. W. Taylor, “Improved regularization of convolutional neural networks with cutout,” _arXiv preprint arXiv:1708.04552_ , 2017.
  * [37] P. Chen, L. Gao, X. Shi, K. Allen, and L. Yang, “Fully automatic knee osteoarthritis severity grading using deep neural networks with a novel ordinal loss,” _Computerized Medical Imaging and Graphics_ , vol. 75, pp. 84–92, 2019.
  * [38] J. Kellgren and J. Lawrence, “Radiological assessment of osteo-arthrosis,” _Annals of the rheumatic diseases_ , vol. 16, no. 4, p. 494, 1957.
  * [39] K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 770–778.
  * [40] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo, “Swin transformer: Hierarchical vision transformer using shifted windows,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2021, pp. 10 012–10 022.
  * [41] A. Trockman and J. Z. Kolter, “Patches are all you need?” _arXiv preprint arXiv:2201.09792_ , 2022.
  * [42] N. Karessli, Z. Akata, B. Schiele, and A. Bulling, “Gaze embeddings for zero-shot image classification,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2017, pp. 4525–4534.
  * [43] M. Carrasco, “Visual attention: The past 25 years,” _Vision research_ , vol. 51, no. 13, pp. 1484–1525, 2011.
  * [44] O. Ronneberger, P. Fischer, and T. Brox, “U-net: Convolutional networks for biomedical image segmentation,” in _International Conference on Medical image computing and computer-assisted intervention_. Springer, 2015, pp. 234–241.
  * [45] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” _arXiv preprint arXiv:1412.6980_ , 2014.
  * [46] N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov, “Dropout: a simple way to prevent neural networks from overfitting,” _The journal of machine learning research_ , vol. 15, no. 1, pp. 1929–1958, 2014.
  * [47] T. Gao, X. Yao, and D. Chen, “Simcse: Simple contrastive learning of sentence embeddings,” _arXiv preprint arXiv:2104.08821_ , 2021.
  * [48] M. Contributors, “MMSelfSup: Openmmlab self-supervised learning toolbox and benchmark,” https://github.com/open-mmlab/mmselfsup, 2021.
  * [49] X. Hou and L. Zhang, “Saliency detection: A spectral residual approach,” in _2007 IEEE Conference on computer vision and pattern recognition_. Ieee, 2007, pp. 1–8.
  * [50] S. Montabone and A. Soto, “Human detection using a mobile platform and novel features derived from a visual saliency mechanism,” _Image and Vision Computing_ , vol. 28, no. 3, pp. 391–402, 2010.
  * [51] T. Zhao and X. Wu, “Pyramid feature attention network for saliency detection,” in _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , 2019, pp. 3085–3094.


# Neural Language Modeling With Implicit Cache Pointers

###### Abstract

A cache-inspired approach is proposed for neural language models (LMs) to
improve long-range dependency and better predict rare words from long
contexts. This approach is a simpler alternative to attention-based pointer
mechanism that enables neural LMs to reproduce words from recent history.
Without using attention and mixture structure, the method only involves
appending extra tokens that represent words in history to the output layer of
a neural LM and modifying training supervisions accordingly. A memory-
augmentation unit is introduced to learn words that are particularly likely to
repeat. We experiment with both recurrent neural network- and Transformer-
based LMs. Perplexity evaluation on Penn Treebank and WikiText-2 shows the
proposed model outperforms both LSTM and LSTM with attention-based pointer
mechanism and is more effective on rare words. $N$-best rescoring experiments
on Switchboard indicate that it benefits both very rare and frequent words.
However, it is challenging for the proposed model as well as two other models
with attention-based pointer mechanism to obtain good overall WER reductions.

Index Terms: RNNLM, Transformer, cache model, pointer component, automatic
speech recognition

##  1 Introduction

Neural language models (LMs) are an important module in automatic speech
recognition (ASR) [1, 2, 3]. Standard recurrent neural network language models
(RNNLMs) make predictions based on a fix-sized hidden vector, making modeling
long-range dependency challenging. Although LSTMs outperform vanilla RNNs, it
has been observed that they usually retain only a relatively short span of
context [4, 5]. Memory augmented models and attention mechanism have been
proposed to increase the hidden state’s capacity to retrieve information from
hidden states in the more distant past. Though improved performance has been
reported, RNNLMs with the standard softmax output still struggle with rare or
unknown words, even with attention.

Since the self-attention architecture was proposed [6], deep Transformers have
demonstrated state-of-the-art performance on natural language processing tasks
[7, 8, 9]. Transformer-based LMs have outperformed RNNLMs on large corpora and
been used in rescoring stage in ASR systems [10, 11]. However, their ability
to capture long-term dependency, e.g. self-trigger effects (word repetitions),
remains unclear.

In real scenarios, especially in conversations, after a word or phrase is
spoken, it is highly likely to be spoken again [12, 13]. These self-triggers
or topic-word effects can be captured by cache models, which stores the
unigram distribution of recently seen words. Cache models adapt pre-trained
LMs to local contexts (decoded hypotheses) in ASR systems and hence can
improve ASR performance [14, 15]. Usually, cache models are integrated in pre-
trained models at test time. It is a lightweight approach as no model
retraining is required, while it may not be optimal. Effectively incorporating
them in training stages and enabling neural LMs to learn to adapt to recent
history remain to be explored.

In this work, we propose a cache-inspired approach for neural LMs to improve
the capability of modeling long-term dependency, especially for rare words.
The output is extended by a predefined size $L$ to represent $L$ preceding
words in history. The pre-softmax activation of the $L$ units, like the other
pre-softmax units, is computed by a linear transformation of the hidden state
or context vector, and then appended to the output before the softmax layer,
as shown in Figure 1. The training loss is still cross entropy. However,
unlike standard training, wherein supervision comes from a vocabulary-sized
one-hot vector encoding the predicted word, the supervision vector is now $L$
bits longer and contains additional ones in each history position where the
word is the same as the predicted one.

The extended output and modified supervision implicitly enable learning _from
where_ in history _to copy_. While it may still be difficult for the model to
learn which words are particularly likely to be self-triggers, i.e. _when to
copy_. To provide a mechanism for this, one additional unit is introduced in
the pre-softmax layer (but not included in the softmax computation) to capture
the probability that the current word may be a self-trigger. At each word
position, activations from these additional units in the $L$ previous
positions are added to the $L$ extended output units.

Though cache-inspired neural LMs for improving long-range dependency have been
proposed and demonstrated superior performance than LSTMs in terms of
perplexity, to our best knowledge, their effect on ASR accuracy remains to be
explored. In this study, we evaluate neural LMs on ASR tasks. We also apply
the proposed approach to Transformer architecture to verify if cache-based
information is still beneficial.

##  2 Related Work

In this section, we briefly introduce related work about approaches to improve
performance on rare words and long-term dependency for sequence modeling
problems including neural language modeling and machine translation [16, 17,
18, 19, 20, 21]. Vinyals et al. [16] introduces an attention-based pointer
network to select items from the input as output. It has been shown to help on
geometric problems [16]. The pointer network can also improve performance of
text summarization [17, 18] and alleviate issues of rare or unknown words in
neural machine translation [18].

For neural LMs, similar ideas have been proposed to better model long-range
dependency [19, 20]. The most relevant work is the pointer sentinel mixture
model (PSMM), a mixture model of a standard LSTM and an auxiliary pointer
network which captures the unigram distribution of history words via attention
[19]. The mixture weight is jointly optimized. A similar mixture model, neural
cache model [20], differs from PSMM in aspects such as the query vector for
computing attention scores is hidden state itself instead of a projected
version and it does not require model retraining. The motivation of dynamic
evaluation [21] is similar to the neural cache model, but the implementation
is different: it adjusts model parameters via gradient updates based on
partial predicted sequences during test time. It may be viewed as a modified
version of the dynamic updating method proposed by Mikolov et al. [1].

##  3 Proposed Model

Language modeling can be framed as predicting the next word (target) given
preceding words (history). It usually can be observed that some words tend to
be much more likely targets once they have occurred in the history. PSMM
learns to “reproduce” a word from recent history by an attention-based pointer
network. And its mixture weight is computed by a specially designed gating
mechanism. Though the PSMM achieves lower perplexity than standard LSTM, the
attention and gating mechanisms are relatively complex and not the only
approach to do so. We aim to achieve a similar effect with simpler models.

###  3.1 Pointer Component

Let us denote the hidden state of an RNNLM at time step $t$ as
$\mathbf{h}_{t}$. Conventionally, the RNNLM output $\mathbf{y}_{t}$ is
determined as

$\mathbf{y}_{t}=\text{softmax}(\mathbf{W}\mathbf{h}_{t}+\mathbf{b}),$ (1)

where $\mathbf{W}\in\mathbb{R}^{V\times H}$, $\mathbf{b}\in\mathbb{R}^{V}$,
and $\mathbf{y}_{t}\in\mathbb{R}^{V}$, with $V$ and $H$ being the vocabulary
size and hidden state dimension, respectively.

We extend the output dimension by a predefined size $L$. The extended part
represents the $L$ immediately preceding words in history. Activation of these
$L$ extended units, denoted as $\mathbf{p}_{t}$ in Figure 1, is computed via a
linear projection of $\mathbf{h}_{t}$ from the last hidden layer of the RNNLM.
We thus have

$\displaystyle\mathbf{p}_{t}=\mathbf{W}_{p}\mathbf{h}_{t},$ (2)
$\displaystyle\mathbf{z}_{t}=\text{concat}(\mathbf{W}\mathbf{h}_{t}+\mathbf{b},\mathbf{p}_{t}),$
(3) $\displaystyle\mathbf{y}_{t}=\text{softmax}\left(\mathbf{z}_{t}\right),$
(4)

where $\mathbf{W}_{p}\in\mathbb{R}^{L\times H}$,
$\mathbf{z}_{t}\in\mathbb{R}^{V+L}$, and $\mathbf{p}_{t}\in\mathbb{R}^{L}$.
Applying softmax on $\mathbf{z}_{t}$ generates the extended output
$\mathbf{y}_{t}\in\mathbb{R}^{V+L}$.

Since the $L$ extended outputs indicate _where_ to copy _from_ the history, we
call $\mathbf{p}_{t}$ the _pointer_ component of our model. It only introduces
$L$ $\times$ $H$ additional parameters.

The objective for training a neural LM is to maximize the log likelihood of
training data. The loss function is written as

$L(\theta)=-\frac{1}{T}\sum_{t=1}^{T}\log(\mathbf{y}_{t}\bm{\cdot}\mathbf{s}_{t}),$
(5)

where $T$ is the total number of words in the training data, $\mathbf{s}_{t}$
is the supervision vector, and $\bm{\cdot}$ is vector dot product operation.
In conventional training, $\mathbf{s}_{t}$ is a one-hot vector with 1 in the
index of the target word. To train the proposed neural LM with the pointer
component, the supervision vector $\mathbf{s}_{t}$ is set to have _additional_
ones in history positions where the the target was previously seen. So
$\mathbf{s}_{t}$ is an at-least-one-hot vector.

Figure 1: Neural LMs with implicit cache pointers.

###  3.2 Memory Augmented Pointer Component

The pointer component and the modified training supervision makes an RNNLM be
aware of where to copy from the history. However, it may still be challenging
for the model to memorize which words are particularly likely to reoccur, i.e.
are ”bursty”. To learn the burstiness of words, one additional unit, denoted
by $m_{t}$ is introduced alongside the pointer $\mathbf{p}_{t}$, as shown in
Figure 1.

This additional unit is computed by a vector dot product of the hidden state
$\mathbf{h}_{t}$ and a parameter vector with the same dimension as
$\mathbf{h}_{t}$, but is not used in computing the softmax. It influences the
probability that a word may repeat through $\mathbf{p}_{t}$. Specifically,
these additional units from the $L$ immediately preceding word positions are
concatenated to form

$\mathbf{m}_{t}=\text{concat}(m_{t-(L-1)},...,m_{t}),$ (6)

where $\mathbf{m}_{t}\in\mathbb{R}^{L}$, and $\mathbf{m}_{t}$ is element-
wisely added to the pointer component $\mathbf{p}_{t}$, i.e.

$\mathbf{p}_{t}:=\mathbf{p}_{t}+\mathbf{m}_{t}.$ (7)

Thus $m_{t}$ influences the output $\mathbf{y}_{t}$ indirectly by modifying
$\mathbf{p}_{t}$ in (3), which in turn is a part of $\mathbf{z}_{t}$ in (4).

Compared with an RNNLM, the memory augmented pointer component only has
$(L+1)$ $\times$ $H$ total additional parameters while a PSMM has extra
$H^{2}+2H$ parameters. Without attention and gating mechanism, the proposed
model is simpler than PSMM and has fewer extra parameters when $L$ $\leq$ $H$.

This pointer mechanism described above is for RNNLMs, while it can also be
easily incorporated into Transformer-based LMs. In the latter, the context
vector from the last Transformer block is treated as the hidden state in
RNNLMs.

##  4 Experimental Setup

We conduct experiments on two text datasets, Penn Treebank (PTB) and
WikiText-2 [19], and two ASR corpora, Switchboard (SWBD) and Wall Street
Journal (WSJ). We use Kaldi RNNLM [3] for data preprocessing on SWBD, e.g.
including the English Fisher corpus, and WSJ. Sentences in SWBD+Fisher
interleave conversation turns, as derived from time-information in
transcriptions. Statistics of the datasets are shown in Table 1 (“sent len” is
average sentence length).

Table 1: Statistics of datasets used in experiments.

Dataset |  # words* |  $|$Vocab$|$ | sent len | OOV (train / dev / test) | Style
---|---|---|---|---|---
PTB | 929K | 10K | 21 | 4.8% / 4.7% / 5.8% | written
WikiText-2 | 2M | 33K | 22 | 2.6% / 5.4% / 6.2% | written
SWBD+Fisher | 34M | 30K | 10 | 8.9% / 0.0% / 5.8% | spoken
WSJ | 39M | 123K | 23 | 4.6% / 4.6% / 5.6% | written

  * *

The end of sentence token is included in the count of training words.

We develop baselines with both LSTM and Transformer-based LMs. Model details
are present in Table 2. Plain LSTMs are baselines from each dataset except for
WSJ. We has a stronger baseline for PTB and WikiText-2: AWD-LSTM [22] with
frequency-agnostic word embeddings [23], denoted by Frage-AWD-LSTM. For
SWBD+Fisher, the stronger baseline is a Transformer LM with self-attention
[6]. We only experiment with Transformer architecture on WSJ. Given our
academic computational resources, we were unable to make comparisons with even
stronger baselines, e.g. GPT and BERT, or with the optimized architectures of
[24], which requires industrial-strength resources.

All neural LMs are on word level, implemented with Pytorch, and optimized via
SGD111For the Transformer LMs, we tried Adam with the learning rate schedule
proposed in [6], but failed to get better performance than SGD.. We tie the
embedding and output matrices in all setups. The dropout rate for PTB and
WikiText-2 is 0.5, while for SWBD+Fisher it is 0.1 for both LSTM and
Transformer LMs. Parameters of Frage-AWD-LSTMs not listed in Table 2 follow
the settings in [23].

Table 2: Details of neural network dimensions for various LMs.

Model | Corpus | Layers | Units | Heads
---|---|---|---|---
Plain LSTM | All (except for WSJ) | 2 | 650 | -
Frage-AWD-LSTM | PTB/WikiText-2 | 3 | 1150 | -
Transformer | SWBD+Fisher/WSJ | 6 | 512/768 | 8

For ASR experiments on SWBD and WSJ, we use the Kaldi toolkit [25] to train
acoustic models and perform $N$-best rescoring. Acoustic models are factorized
TDNNs [26], trained using the LF-MMI objective [27]. We do not include Fisher
audio to train acoustic models for SWBD. To rescore each of the $N$ hypotheses
for an utterance, we find it useful to initialize the initial LM state with
the last LM state of the best hypothesis for the previous utterance.

##  5 Experiments

###  5.1 Perplexities on PTB and WikiText-2

We first compare the proposed model with PSMM and neural cache under the plain
LSTM setup. Perplexities on PTB and WikiText-2 are shown in Table 3. The
performance gap between the PSMM in the paper [19] and ours is mainly caused
by different implementations of truncated back-propagation through time
(BPTT). They use an explicit truncated BPTT while we follow the normal way
discussed in [19] considering efficiency and convenience of data
preprocessing. We first concatenate all text words and then chunk them with
fixed size $L$. So, if the truncated BPTT length is $L$, each training word on
average experiences $L$/2 instead of $L$ time-steps for back-propagation. This
means each training word sees $L$/2 history words on average.

Table 3: Perplexities on PTB and WikiText-2 (plain LSTMs).

Model | PTB | WikiText-2
---|---|---
#Params | Dev | Test | #Params | Dev | Test
5gram KN [28] | 2M | - | 141.2 | - | - | -
LSTM (medium) [29] | 20M | 86.2 | 82.7 | - | - | -
PSMM [19] | 21M | 72.4 | 70.9 |  47M222The stated number of parameters of the PSMM in [19] is 21M. Considering the hidden state size and vocabulary, we suspect that to be a typo. We obtain around 47M assuming no parameter tying. | 84.8 | 80.8
LSTM | 13.3M | 73.6 | 71.9 | 28.5M | 89.1 | 84.8
PSMM (Ours) | 13.7M | 73.5 | 71.6 | 28.9M | 86.8 | 82.8
LSTM + Neural Cache (L=50) | 13.7M | 69.3 | 68.5 | 28.9M | 81.3 | 77.0
Proposed w/o Memory Aug | 13.3M | 70.1 | 69.8 | 28.5M | 80.6 | 76.7
Proposed w/ Memory Aug | 13.3M | 68.1 | 67.8 | 28.5M | 78.2 | 74.3

In Table 3 “Memory Aug” refers to the memory augmented pointer. We set history
length as 100, equal to the truncated BPTT length. Setting $L$ = 50 for neural
cache is a fair comparison with others. Results on both datasets show that
memory augmentation provides further improvement on top of the pointer
component. And with memory augmentation, the proposed model outperforms the
rest on both datasets. In subsequent tables, “Proposed” refers to LMs with the
memory augmented pointer.

To verify whether the proposed approach is robust, we conduct experiments on a
stronger baseline Frage-AWD-LSTM setup [23]. We reproduced their results and
implemented the proposed approach on top of theirs, without tuning meta-
parameters. Perplexity results in Table 4 shows that the proposed model on
both datasets achieves better results than Frage-AWD-LSTM. Further
improvements are observed with increasing the history length from 50 to 100,
as expected. We also observe complementary effects of the proposed model and
neural cache model.

Table 4: Perplexities on PTB and WikiText-2 (Frage-AWD-LSTM setup).

Model | PTB | WikiText-2
---|---|---
#Params |  Dev |  Test | #Params |  Dev |  Test
Frage-AWD-LSTM [23] | 24M |  58.1 |  56.1 | 33M |  66.5 |  63.4
Frage-AWD-LSTM (Ours) | 24M |  57.7 |  55.3 | 33M |  64.6 |  62.1
Proposed (L = 50) | 24M |  55.2 |  53.9 | 33M |  60.7 |  58.5
Proposed (L = 100) | 24M |  54.2 |  53.5 | 33M |  60.2 |  57.5
Proposed + Neural Cache(L=100) | 24M |  53.0 |  52.5 | 33M |  58.0 |  55.6

###  5.2 Perplexities on SWBD and WSJ

We experiment with both LSTM- and Transformer-based LMs on SWBD. Perplexity on
dev set from Kaldi RNNLM is 50. The history length as well as BPTT length is
set to 100 for the proposed model and PSMM. Results of Pytorch trained models
are in Table 5.

Table 5: Perplexities on SWBD.

Model | #Params | Dev | Eval’00
---|---|---|---
LSTM | 26.5M | 47.0 | 41.7
PSMM | 26.9M | 45.6 | 39.5
LSTM + Neural Cache (L=50) | 26.5M | 45.1 | 39.6
LSTM + Neural Cache (L=100) | 26.5M | 44.6 | 39.1
LSTM + Proposed | 26.5M | 45.9 | 40.4
Transformer w/o positional embedding [6] | 25.0M | 51.6 | 44.4
Transformer with positional embedding | 25.1M | 46.8 | 41.5
Transformer + Proposed | 25.1M | 45.0 | 40.2

For LSTM-based models, the proposed approach outperforms the baseline LSTM,
but performs slightly worse than the PSMM and neural cache models. For
Transformer LMs, the proposed approach also achieves better perplexity than
the two Transformer baselines.

We notice the performance gains of the proposed model over both LSTM and
Transformer baselines on SWBD are smaller than on PTB and WikiText-2. To check
whether this may relate to style (only SWBD is spoken style) and average
sentence length (sentences on Switchboard are the shortest on average), we
experiment with Transformer-based LMs on WSJ. The proposed approach reduces
perplexity from 71.5 to 65.6 on test set (eval92), compared with a baseline
Transformer LM. The improvement is relative 8.2% and in similar range with
gains on WikiText-2. And this results in 0.1 absolute WER reduction from 1.5
to 1.4 on the same test set by $N$-best rescoring. While no conclusions can be
drawn yet, experiments show that the proposed model perform better on written-
style text which usually has longer average sentence length.

###  5.3 Analysis of Impact on Rare Words

One desired feature of the proposed model is that it may predict rare words
better than LSTMs. To verify this, we further examine LM performance on the
test set of each dataset. We split the vocabulary of each corpus into 10
buckets based on word frequencies in training data such that we get a roughly
equal number of test tokens in each bucket. We then compute the differences
between the test cross entropy of the proposed model and the LSTM baseline on
words in each bucket for each dataset. Figure 2 shows the results on
WikiText-2. As expected, larger reductions in cross-entropy are observed from
the proposed model on rare words. Similar trends are seen on PTB and SWBD,
though the overall perplexity improvement on SWBD is marginal. Figures for
them are omitted due to page limits.

Figure 2: Cross-entropy reduction from the proposed model w.r.t an LSTM on
WikiText-2.

###  5.4 Rescoring Evaluation and Analysis on SWBD

As Transformers show similar perplexity gains as LSTMs on SWBD, we only
experiment with LSTMs for $N$-best rescoring. WERs on the full HUB5’00
evaluation set (Eval’00), the SWBD subset (SWB), and Callhome subset (CH) are
in Table 6. “State-carry” means when scoring a hypothesis for current
utterance, the initial hidden state is copied from the last hidden state of
the best hypothesis for the previous utterance, instead of being zero
initialized. WER improvements by the LSTM with “state-carry” in Table 6
indicate that cross-sentence context is useful. Similar observations are
presented in [30]. If not specified with “w/o state-carry”, models are
evaluated in the state-carry way.

Table 6: WERs by N-best rescoring with baselines and the proposed model on
SWBD.

Model | Eval’00 | SWB | CH
---|---|---|---
Kaldi RNNLM (w/o state-carry) | 11.3 | 7.5 | 15.0
LSTM (w/o state-carry) | 11.2 | 7.3 | 15.1
LSTM | 10.9 | 7.1 | 14.5
PSMM | 10.9 | 7.1 | 14.6
LSTM + Neural Cache | 10.9 | 7.2 | 14.5
LSTM + Proposed | 10.8 | 7.1 | 14.4

To investigate the effect on WERs of rare words, we conduct a similar analysis
as Section 5.3 does. Words with errors on Eval’00 (<5000) are divided into 5
buckets based on frequency. Relative WER reductions on words in these buckets
in Figure 3 show that the proposed model improves performance on both
relatively rare words and very frequent ones. As expected, some rare words
occur within the context window, for example, masters and offered, are
recognized correctly. Decoded output also shows that frequent words such as
train, short, and were are correctly recognized. Though the overall WER
improvement by the proposed model is marginal, correctly recognizing
relatively rare words plays an important role in user experience of ASR-based
products or service.

Figure 3: Relative WER reduction on by the proposed model w.r.t an LSTM on
SWBD.

We also notice the proposed model sometimes introduces errors on words that
are wrongly recognized in first pass decoding. A possible reason is that the
supervision vectors for pointer components are from decoded hypotheses and
hence may contain errors. We verify this by using test transcription in
rescoring and observe a further 0.1 absolute WER reduction on Eval’00 of SWBD.
The mismatched condition between training and evaluation is a common issue for
the proposed approach, PSMM, and neural cache model. To alleviate the
mismatch, word level confidence scores and error adaptive training approaches
could be considered.

##  6 Conclusion and Future Work

In this work, we propose a cache-inspired pointer mechanism for neural LMs to
improve the capacity of modeling long-range dependency and better predict rare
words. It can be applied to both RNN- and Transformer-based models. Perplexity
evaluation show that the proposed approach generally outperforms LSTM and PSMM
and is more effective on rare words. Rescoring with the proposed model on SWBD
and WSJ gives marginal WER improvements. Analysis shows that the mismatch
between training and rescoring conditions (i.e. potentially incorrect
histories) may make it challenging for both the proposed model and models with
attention-based pointer network to achieve large overall WER reductions.
Future work is therefore focused on methods that can mitigate the mismatch
issue.

## References

  * [1] T. Mikolov, M. Karafiát, L. Burget, J. Černockỳ, and S. Khudanpur, “Recurrent neural network based language model,” in _Proc. of Interspeech_ , 2010.
  * [2] X. Chen, X. Liu, M. J. Gales, and P. C. Woodland, “Recurrent neural network language model training with noise contrastive estimation for speech recognition,” in _Proc. of ICASSP_ , 2015.
  * [3] H. Xu, K. Li, Y. Wang, J. Wang, S. Kang, X. Chen, D. Povey, and S. Khudanpur, “Neural network language modeling with letter-based features and importance sampling,” in _Proc. of ICASSP_ , 2018.
  * [4] U. Khandelwal, H. He, P. Qi, and D. Jurafsky, “Sharp nearby, fuzzy far away: How neural language models use context,” in _Proc. of ACL_ , 2018.
  * [5] C. Chelba, M. Norouzi, and S. Bengio, “N-gram language modeling using recurrent neural network estimation,” _arXiv preprint arXiv:1703.10724_ , 2017.
  * [6] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” in _Proc. of NeurIPS_ , 2017.
  * [7] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “Bert: Pre-training of deep bidirectional transformers for language understanding,” in _Proc. of NAACL_ , 2019.
  * [8] A. Radford, K. Narasimhan, T. Salimans, and I. Sutskever, “Improving language understanding by generative pre-training,” 2018.
  * [9] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. V. Le, and R. Salakhutdinov, “Transformer-xl: Attentive language models beyond a fixed-length context,” in _Proc. of ACL_ , 2019.
  * [10] K. Irie, A. Zeyer, R. Schlüter, and H. Ney, “Language modeling with deep transformers,” in _Proc. of Interspeech_ , 2019.
  * [11] K. Li, Z. Liu, T. He, H. Huang, F. Peng, D. Povey, and S. Khudanpur, “An empirical study of transformer-based neural language model adaptation,” in _Proc. of ICASSP_ , 2020.
  * [12] R. Lau, R. Rosenfeld, and S. Roukos, “Trigger-based language models: A maximum entropy approach,” in _Proc. of ICASSP_ , 1993.
  * [13] K. W. Church, “Empirical estimates of adaptation: the chance of two noriegas is closer to p/2 than $p^{2}$,” in _Proc. of COLING_ , 2000.
  * [14] R. Kuhn and R. De Mori, “A cache-based natural language model for speech recognition,” _IEEE transactions on pattern analysis and machine intelligence_ , vol. 12, no. 6, pp. 570–583, 1990.
  * [15] K. Li, H. Xu, Y. Wang, D. Povey, and S. Khudanpur, “Recurrent neural network language model adaptation for conversational speech recognition.” in _Proc. of Interspeech_ , 2018.
  * [16] O. Vinyals, M. Fortunato, and N. Jaitly, “Pointer networks,” in _Proc. of NeurIPS_ , 2015.
  * [17] J. Gu, Z. Lu, H. Li, and V. O. Li, “Incorporating copying mechanism in sequence-to-sequence learning,” in _Proc. of ACL_ , 2016.
  * [18] C. Gulcehre, S. Ahn, R. Nallapati, B. Zhou, and Y. Bengio, “Pointing the unknown words,” in _Proc. of ACL_ , 2016.
  * [19] S. Merity, C. Xiong, J. Bradbury, and R. Socher, “Pointer sentinel mixture models,” in _Proc. of ICLR_ , 2017.
  * [20] E. Grave, A. Joulin, and N. Usunier, “Improving neural language models with a continuous cache,” _Proc. of ICLR_ , 2017.
  * [21] B. Krause, E. Kahembwe, I. Murray, and S. Renals, “Dynamic evaluation of neural sequence models,” _Proc. of ICML_ , 2018.
  * [22] S. Merity, N. S. Keskar, and R. Socher, “Regularizing and optimizing lstm language models,” in _ICLR_ , 2018.
  * [23] C. Gong, D. He, X. Tan, T. Qin, L. Wang, and T.-Y. Liu, “Frage: Frequency-agnostic word representation,” in _Proc. of NeurIPS_ , 2018.
  * [24] C. Wang, M. Li, and A. J. Smola, “Language models with transformers,” _arXiv preprint arXiv:1904.09408_ , 2019.
  * [25] D. Povey, A. Ghoshal, G. Boulianne, L. Burget, O. Glembek, N. Goel, M. Hannemann, P. Motlicek, Y. Qian, P. Schwarz _et al._ , “The kaldi speech recognition toolkit,” in _Proc. of ASRU_ , 2011.
  * [26] D. Povey, G. Cheng, Y. Wang, K. Li, H. Xu, M. Yarmohammadi, and S. Khudanpur, “Semi-orthogonal low-rank matrix factorization for deep neural networks.” in _Proc. of Interspeech_ , 2018.
  * [27] D. Povey, V. Peddinti, D. Galvez, P. Ghahremani, V. Manohar, X. Na, Y. Wang, and S. Khudanpur, “Purely sequence-trained neural networks for asr based on lattice-free mmi.” in _Proc. of Interspeech_ , 2016.
  * [28] T. Mikolov and G. Zweig, “Context dependent recurrent neural network language model,” in _Proc. of SLT_ , 2012.
  * [29] W. Zaremba, I. Sutskever, and O. Vinyals, “Recurrent neural network regularization,” _arXiv preprint arXiv:1409.2329_ , 2014.
  * [30] K. Irie, A. Zeyer, R. Schlüter, and H. Ney, “Training language models for long-span cross-sentence evaluation,” in _Proc. of ASRU_ , 2019.