# ATLANTIS: A Benchmark for Semantic Segmentation of Waterbody Images

Mohammad H. Erfani
Department of Civil and Environmental Engineering
University of South Carolina
Columbia, SC
<EMAIL_ADDRESS>
Zhenyao Wu, Xinyi Wu
Department of Computer Science and Engineering
University of South Carolina
Columbia, SC
{zhenyao<EMAIL_ADDRESS>
Song Wang
Department of Computer Science and Engineering
University of South Carolina
Columbia, SC
<EMAIL_ADDRESS>
Erfan Goharian
Department of Civil and Environmental Engineering
University of South Carolina
Columbia, SC
<EMAIL_ADDRESS>

###### Abstract

Vision-based semantic segmentation of waterbodies and nearby related objects
provides important information for managing water resources and handling
flooding emergency. However, the lack of large-scale labeled training and
testing datasets for water-related categories prevents researchers from
studying water-related issues in the computer vision field. To tackle this
problem, we present ATLANTIS, a new benchmark for semantic segmentation of
waterbodies and related objects. ATLANTIS consists of 5,195 images of
waterbodies, as well as high quality pixel-level manual annotations of 56
classes of objects, including 17 classes of man-made objects, 18 classes of
natural objects and 21 general classes. We analyze ATLANTIS in detail and
evaluate several state-of-the-art semantic segmentation networks on our
benchmark. In addition, a novel deep neural network, AQUANet, is developed for
waterbody semantic segmentation by processing the aquatic and non-aquatic
regions in two different paths. AQUANet also incorporates low-level feature
modulation and cross-path modulation for enhancing feature representation.
Experimental results show that the proposed AQUANet outperforms other state-
of-the-art semantic segmentation networks on ATLANTIS. We claim that ATLANTIS
is the largest waterbody image dataset for semantic segmentation providing a
wide range of water and water-related classes and it will benefit researchers
of both computer vision and water resources engineering.

Figure 1: ATLANTIS: ArTificiaL And Natural waTer-bodIes dataSet.

##  1 Introduction

Every year, floods claim tens of billions of US dollars losses and thousands
of lives globally [17]. Accurate detection, measurement, and track of the
waterbodies can help both the public and decision-makers to take appropriate
actions to minimize the risk and losses [19, 15]. With the popularity of smart
phones and drones, various data at flooding sites can be collected rapidly and
continuously to provide more useful and heterogeneous information source [12,
11, 5, 38], compared to the conventional gauge sensing and remote sensing [12,
27]. As a fundamental step to leverage such collected images for modeling and
decision-making, we need first to conduct a refined semantic segmentation of
included waterbodies and related objects in such scenes, which we focus on in
this paper.

With the advancement of deep neural networks, semantic segmentation has
achieved a great success in recent years on various kinds of images, such as
natural images [26], street images [8, 46], and medical images [2, 22].
However, waterbody images pose many new unique challenges for semantic
segmentation. In some forms, water preserves the intrinsic properties such as
reflection, transparency, shapeless and colorless visual features; which in
turn, brings difficulties to semantic segmentation of water and related
objects. Moreover, in some other forms, these properties can be affected by
illumination sources from surroundings, turbidity, and turbulence. Different-
labeled waterbodies, such as river and canal, or lake and reservoir, often
have similar visual characteristics that make task of semantic segmentation
even harder. As shown in our later experiments, these unique challenges may
significantly affect the performance of the existing semantic segmentation
networks.

Meanwhile, deep-learning based approaches always require large-scale training
data with necessary ground-truth annotations. Lack of such a public dataset
for waterbody segmentation significantly impedes the research on this problem.
The collection and annotation of such a dataset can be very laborious and
time-consuming to cover a wide range of waterbodies and related objects. There
is no specific repository providing relevant images. In addition, team members
and annotators are required to have prior knowledge on water resources
engineering to be capable of selecting and precisely annotating the images.

In this paper, we present a new benchmark, ATLANTIS (ArTificiaL And Natural
waTer-bodIes dataSet). For the first time, this dataset has covered a wide
range of natural and man-made (artificial) waterbodies such as sea, lake,
river, canal, reservoir, and dam. ATLANTIS includes 5,195 pixel-wise annotated
images split to 3,364 training, 535 validation and 1,296 testing images. As
shown in the Table 1, in addition to 35 waterbody and water-related objects,
ATLANTIS also covers 21 general labels. Moreover, we construct ATLANTIS
Texture (ATeX) dataset, which consists of 12,503 patches for the water-bodies
texture classification, sampled from 15 kinds of waterbodies in ATLANTIS.

Table 1: List of the ATLANTIS labels. Artificial |  | breakwater; bridge; canal; culvert; dam; ditch; levee; lighthouse; pipeline;
---
pier; offshore platform; reservoir; ship; spillway; swimming pool;
water tower; water well.
Natural |  | cliff; cypress tree; fjord; flood; glaciers; hot spring; lake; mangrove; marsh;
---
puddle; rapids; river; river delta; sea; shoreline; snow; waterfall; wetland.
General |  | road; sidewalk; building; wall; fence; pole; traffic sign; vegetation; terrain;
---
sky; train; person; car; bus; truck; bicycle; parking meter; motorcycle;
fire hydrant; boat; umbrella.

In order to tackle the inherent challenges in the segmentation of waterbodies,
AQUANet is developed which takes an advantage of two different paths to
process the aquatic and non-aquatic regions, separately. Each path includes
low-level feature and cross-path modulation, to adjust features for better
representation. The results show that the proposed AQUANet outperforms other
ten state-of-the-art semantic-segmentation networks on ATLANTIS, and the
ablation studies justify the effectiveness of the components of the proposed
AQUANet.

##  2 Related Work

###  2.1 Semantic segmentation dataset

Large-scale annotated datasets, such as COCO [26], PASCAL Context [30], ADE20K
[51], and more recently Mapillary Vistas Dataset [32] and BDD100K [46], make
it possible for researchers to develop deep-learning based models for real-
world applications. Considering the most related dataset to ATLANTIS,
Gebrehiwot et al. [15] collected a small number of top-view waterbody dataset
(100 images) using Unmanned Aerial Vehicles (UAVs) which contains only four
categories (i.e., water, building, vegetation and road). In addition, Sazara
et al. [37] introduced a larger dataset (253 images) which just focuses on
flood region segmentation. More recently, Sarp et al. [36] provided a dataset
that consists of 441 annotated roadway flood images. However, these datasets
have limitations in either the number of annotated images, or the categories
they covered, and none of those considers more complex classes of waterbodies
such as sea, lake and waterfall. Therefore, we develop a new dataset,
ATLANTIS, as the first large-scale annotated dataset to provide a wide-range
of waterbodies and water-related objects.

###  2.2 Semantic segmentation network

All of existing semantic segmentation approaches share the same goal to
classify each pixel of a given image but differ in the network design,
including low-resolution representations learning [28, 6], high-resolution
representations recovering [1, 33, 25], contextual aggregation schemes [47,
50, 48], feature fusion and refinement strategy [25, 20, 23, 52, 14].
Typically, method designs are dependent on their respective datasets and all
the mentioned networks are developed by training on benchmark datasets such as
Cityscapes [8], COCO [26] and VOC [13] where the inter-class boundary is clear
even for the within-group categories (e.g., car and truck). As mentioned
above, waterbody images pose new challenges to semantic segmentation. Previous
works on waterbody segmentation mainly use satellite imagery [31, 24, 10]. In
this work, we focus on natural waterbody images terrestrially captured by
various cameras, and design AQUANet, a new two-path semantic segmentation
network, by including an aquatic branch explicitly for waterbody classes.

(a)

(b)

Figure 2: (a) The frequency distribution of images for different waterbodies
in ATLANTIS (b) The Percentage of pixels for waterbody labels.

##  3 ATLANTIS Dataset

The ATLANTIS dataset is designed and developed with the goal of capturing a
wide-range of water-related objects, either those exist in natural environment
or the infrastructure and man-made (artificial) water systems. In this
dataset, labels were first selected based on the most frequent objects, used
in water-related studies or can be found in real-world scenes. Aside from the
background objects, total of 56 labels, including 17 artificial, 18 natural
waterbodies, and 21 general labels, are selected (Table 1). These general
labels are considered for providing contextual information that most likely
can be found in water-related scenes. After finalizing the selection of
waterbody labels, a comprehensive investigation on each individual label was
performed by annotators to make sure all the labels are vivid examples of
those objects in real-world. Moreover, sometimes some of the water-related
labels, e.g., levee, embankment, and floodbank, have been used interchangeably
in water resources field; thus, those labels are either merged into a unique
group or are removed from the dataset to prevent an individual object receives
different labels.

In order to gather a corpus of images, we have used Flickr API to query and to
collect “medium-sized” unique images for each label based on eight commonly
used “Creative Commons”, “No Known Copyright Restrictions” and “United States
Government Work" licenses. Downloaded images were then filtered by a two-stage
hierarchical procedure. In the first stage, each annotator was assigned to
review a specific list of labels and remove irrelevant images based on that
specific list of labels. In the second stage, several meetings were held
between the entire annotation team and the project coordinator to finalize the
images which appropriately represent each of 56 labels.

This sieving procedure has been applied four times in order to meet the limit
and to reach the current number of images. The percentage of image acceptance
rate for the third and fourth phases are 14.41$\%$ and 5.06$\%$, respectively.
It means if we want to add 1000 more images to the dataset, we should process
at least 20,000 images. Finally, images were annotated by annotators who have
solid water resources engineering background as well as experience working
with the CVAT [39], which is a free, open source, and web-based image/video
annotation tool.

###  3.1 Dataset statistics

Figure 2 shows the frequency distribution of the number of images and the
percentage of pixels for waterbody labels. Such a long-tailed distribution is
common for semantic segmentation datasets [26, 51] even if the number of
images that contain specific label are pre-controlled. Such frequency
distribution for pixels would be inevitable for objects existing in real-
world. Taking “water tower” as an example, despite having 219 images, the
percentage of pixels are less than many other labels in the dataset. Figure 3
shows the positive but weak correlation between number of images for each
label and the corresponding pixels. In total, only $4.89\%$ of pixels are
unlabeled, and $34.17\%$ and $60.94\%$ of pixels belong to waterbodies
(natural and man-made) and general labels, respectively. As it is evident, the
main proportion of pixels belongs to general labels. This clearly shows the
importance of general labels for better scene understanding [3] and accurate
object classification in semantic segmentation network.

Figure 3: Correlation between number of images for each label and the
corresponding pixels

####  3.1.1 Spatial analysis

Following ADE20K dataset [51], a spatial analysis, known as “mode of the
object segmentation” has been done on the ground truth segmentation map for
each label. Specifically, considering a waterbody label $L$ with $n$ images in
ATLANTIS, we resize their corresponding $n$ ground-truth segmentation maps to
$512\times 512$ pixels. We then count the most frequent label at each pixel of
the map and construct a “mode of segmentation” map for label $L$, as shown in
Figure 4. This map demonstrates the spatial distributions of the most frequent
co-occurred labels with respect to a given waterbody label. Based on this, we
equipped our proposed network with cross-path feature modulation to cope with
the difficulties associated with the recognition of waterbody labels having
visual similarities.

Figure 4: Spatial analysis of four different waterbody labels.

###  3.2 Image annotation

####  3.2.1 Annotation pipeline

ATLANTIS was annotated by six annotators having prior knowledge in the area of
water resources. The goal in the annotation task is to balance speed and
quality. Generally, time spent on a single image can range from 4 minutes to
25 minutes depending on the image complexity. In this project, each kind of
waterbody was assigned to a specific annotator. Before annotation of a label,
all the images of that label are scrutinized and discussed by a group of
experts in water resources engineering. We can see that the annotation of
complex flood scenes takes more time since such images are usually captured in
urban areas and have more elaborated components as shown in Figure 5.

Figure 5: Two samples of complex flood scenes

####  3.2.2 Consistency analysis

While one image is annotated by one annotator for ATLANTIS, we perform
additional consistency analysis across annotators and over time for an
annotator. We choose 52 images from ATLANTIS, by including both images that
are highly susceptible to wrong labelling and those contain objects prone to
be either left unannotated or wrongly annotated. We ask three annotators to
annotate them again and compare the results against the already approved
ground truth in ATLANTIS. The accuracy and mIoU in terms of all 52 images
(total) and the subsets of images that had been annotated by themselves before
(individual) are shown in Table 2. We can see that an annotator can process
the images that he/she annotated before with much better consistency.

Table 2: Consistency analysis for three annotators. |  | Annotator 1 | Annotator 2 | Annotator 3
---|---|---|---|---
Total | acc | 84.00 | 79.93 | 79.00
mIoU | 62.29 | 59.33 | 57.34
Individual | acc | 91.11 | 90.44 | 94.69
mIoU | 72.83 | 76.14 | 75.69

###  3.3 ATLANTIS Texture (ATex)

Waterbodies usually bear texture appearance and it is an interesting problem
to study whether different kinds of waterbodies may show subtle differences
associated with the texture features. For this purpose, we construct a new
waterbody texture dataset, ATeX, by cropping patches from ATLANTIS and take
the corresponding annotated waterbody label as the label of the patch. We set
patch size to $32\times 32$ pixels and all pixels in a cropped patch must have
the same waterbody label in the original image. We also ensure there is no
partial overlap between any two patches. In total we collected 12,503 patches
with 15 waterbody labels: Two waterbody labels “estuary” and “swamp” are added
based on the nearby tree species – mangrove “estuary” and cypress for “swamp”,
while four waterbody labels “canal”, “ditch”, ”reservoir” and “fjord” are
omitted because of high visual similarities with other labels. Sample images
of ATeX dataset are shown in Figure 6. We split ATeX into 8,753 for training,
1,252 for validation and 2,498 for testing. In the later experiment, we train
different models to evaluate their classification performance on ATeX images.

Figure 6: Samples of ATeX texture images.

##  4 AQUANet

Typically, existing semantic segmentation networks are designed based on a
strong backbone (e.g. ResNet [16]) to extract features from images with
additional feature aggregation schemes such as ASP-OC [47] and PPM [50].
Because of difficulties associated with semantic segmentation of waterbodies,
we design AQUANet to segment aquatic and non-aquatic categories, separately,
as shown in Figure 7.

Figure 7: The network architecture of proposed AQUANet.

###  4.1 Network architecture

According to Figure 7, the input image is first fed into a ResNet-101
(pretrained on ImageNet [9]) to extract the feature $F$ with a size of
$C\times H\times W$. Then, the feature is sent into two separate paths for
further processing. The aquatic path is to segment different types of
waterbodies including sea, river, waterfall, wetland and etc., while the non-
aquatic path is to segment other categories such as ship and bridge. In each
path, the feature $F$ is first modulated by the low-level feature $F_{l}$
extracted from the third convolutional layer of ResNet-101, and then passed
into ASP-OC [47] to produce the probability map. In the last step, two cross-
path modulation blocks are applied to adjust the probability maps $P_{1}$ and
$P_{2}$ in parallel. Finally, the resulted probability maps are concatenated
and upsampled to the size of the original image.

###  4.2 Feature modulation

The goal of the feature modulation $\mathcal{M}$ is to adjust a feature map
$F_{1}$ given feature map $F_{2}$ to represent the adjusted feature
$F_{1}^{\prime}$. It can be formulated as:

$F_{1}^{\prime}=\mathcal{M}(F_{1}|F_{2}).$ (1)

To generate the modulated feature $F_{1}^{\prime}$, the parameters $\alpha$
and $\beta$ are learned from $F_{2}$ via the feature modulation that consists
of three downsampling layers, six $1\times 1$ convolutional layers and two
leakyReLU layers as shown in Figure 8. The learned parameters $\alpha$ and
$\beta$ have the shape as the $F_{1}$. Then, the resulting feature
$F_{1}^{\prime}$ is constructed as follows according to [43, 34]:

$F_{1}^{\prime}=\alpha\cdot F_{1}+\beta+F_{1}.$ (2)

####  4.2.1 Low-level feature modulation

Figure 8: An illustration of the feature modulation.

To enhance the low-level texture representation of the water-bodies, we
propose to use the low-level feature $F_{l}$ to modulate the feature $F$ and
the resulted feature $F^{\prime}$ is defined as:

$F^{\prime}=\mathcal{M}(F|F_{l}).$ (3)

Note that different channels of the receptive field in the third convolutional
layer of ResNet-101 are used to construct the low-level feature $F_{l}$ for
the two paths.

####  4.2.2 Cross-path modulation

We also propose a cross-path modulation to aggregate the outputs of the
probability maps. The adjusted probability maps resulted from this module are
defined as:

$P^{\prime}_{1}=\mathcal{M}(P_{1}|P_{2}),$ (4)
$P^{\prime}_{2}=\mathcal{M}(P_{2}|P_{1}),$ (5)

where the $P^{\prime}_{1}$ and $P^{\prime}_{2}$ represent the final
probability maps for the aquatic and non-aquatic objects, respectively. Note
that the modulation layer is not sharing weights across the two paths.

###  4.3 Loss function

The proposed network is trained in a fully supervised fashion constrained by
the widely-used cross-entropy loss function on both final prediction $P$ (main
loss) and the intermediate feature produced by the fourth block of the
ResNet-101 (auxiliary loss). Following Zhao et al. [50], the weights of the
main and the auxiliary loss are set to 1 and 0.4, respectively.

##  5 Experiments

To demonstrate the effectiveness of the proposed method for water-bodies
semantic segmentation, we train AQUANet on the proposed ATLANTIS dataset. For
performance evaluation, we take the mean of class-wise intersection over union
(mIoU) and the per-pixel accuracy (acc) as the main evaluation metrics. To
further evaluate the performance on the waterbodies, we calculate the mean IoU
for aquatic categories (A-mIoU) and the accuracy in the aquatic region
(A-acc). Aquatic categories include 17 labels showing just water content in
different forms and bodies, e.g., sea, river, lake, etc.

Table 3: The per-category results on the ATLANTIS test set by current state-
of-the-art methods and our AQUANet. The best and the second best results are
highlighted with bold font and underline, respectively.

Method canal ditch fjord flood glaciers hot spring lake puddle rapids
reservoir river river delta sea snow swimming pool waterfall wetland A-acc (%)
A-mIoU acc (%) mIoU PSPNet 53.8 29.0 42.9 46.5 57.2 53.9 29.7 54.7 38.2 29.8
28.8 65.5 63.5 49.9 47.7 48.4 47.5 66.19 46.29 72.72 40.85 DeepLabv3 52.5 27.2
52.3 43.8 58.7 42.5 31.1 54.2 46.0 32.4 27.1 51.1 61.5 46.3 53.6 52.8 52.9
65.83 46.25 69.21 36.23 DANet 50.5 34.1 37.1 37.0 51.0 61.6 23.8 51.5 42.8
30.2 31.5 63.5 60.4 50.8 43.1 55.2 54.6 62.00 45.80 74.12 39.60 CCNet 41.1
17.4 35.2 26.9 43.7 47.9 18.6 43.8 29.9 16.6 23.7 48.3 53.3 47.6 38.4 51.1
34.1 51.98 36.33 70.84 36.11 EMANet 46.1 16.6 27.1 23.0 53.8 63.7 17.2 43.6
42.2 17.2 21.0 68.6 53.5 47.3 36.1 52.1 36.2 55.88 39.13 71.93 36.43 ANNet
50.9 22.8 31.6 32.0 53.1 58.1 25.6 52.9 48.4 20.8 28.6 56.8 60.4 51.1 43.9
57.9 51.4 61.51 43.90 74.06 39.79 GCNet 56.6 19.0 44.7 34.8 46.9 36.1 35.8
39.4 39.9 41.6 32.4 67.0 62.2 46.4 42.9 50.7 59.7 69.89 44.48 68.64 37.73
DNLNet 54.4 26.3 48.8 36.3 63.2 55.3 35.5 52.3 40.4 32.1 31.3 37.1 61.7 48.3
52.4 48.7 54.6 67.72 45.80 71.95 39.97 OCNet 56.4 33.6 48.0 37.3 57.7 55.2
29.2 50.6 43.8 35.1 35.6 65.9 62.7 47.2 47.9 53.1 54.9 67.97 47.89 73.54 41.19
OCRNet 52.4 19.4 46.9 34.9 48.3 58.8 30.4 39.7 42.5 29.8 31.9 55.5 55.4 47.3
43.6 56.8 51.5 65.90 43.83 71.66 36.17 Ours 55.0 27.7 53.4 47.0 63.1 60.5 33.2
54.4 46.3 39.0 34.7 63.2 64.2 50.3 44.9 53.0 66.1 68.63 50.34 75.18 42.22

Figure 9: Visualisation comparison of AQUANet and four well-known methods on
the ATLANTIS validation set.

Image

OCNet

GCNet

PSPNet

DANet

AQUANet

GT

###  5.1 Experimental settings

The AQUANet is implemented using PyTorch. During training, the base learning
rate is set to $2.5\times 10^{-4}$ and it is decayed following the poly policy
[50]. The network is optimized using SGD with a momentum of 0.9 and weight
decay of 0.0001. In total, we train the network for 30 epochs, around 80K
iterations with a batch size of 2. The training data are augmented with random
horizontal flipping, random scaling ranging from 0.5 to 2.0 and random
cropping with the size of $640\times 640$.

###  5.2 Comparisons

We use several state-of-the-art networks to perform training and testing on
ATLANTIS, including PSPNet [50], DeepLabv3 [7], CCNet [20], EMANet [23], ANNet
[52], DANet [14], DNLNet [45], GCNet [4], OCNet [47], OCRNet [48]. For a fair
comparison, we train all the networks with the same backbone (ResNet-101) for
30 epochs. As shown in Table 3, the proposed AQUANet outperforms all these
networks on waterbody image semantic segmentation. Figure 9 shows the
visualization results of some samples from ATLANTIS validation set.
Considering the ground truth and comparing with other networks’ outputs, the
boundaries between different classes are better preserved in AQUANet output.
Compared with Chen et al. [7], Zhao et al. [50], Fu et al. [14], and Yuan and
Wang [47], our method achieves better results on both the aquatic and non-
aquatic regions.

###  5.3 Failure cases

Due to the challenges associated with segmentation of waterbody images, there
are still many failure cases we found in the testing stage. Three failure
examples are shown Figure 10, from which we observe that many aquatic classes
are vulnerable to mis-classification – here sea is mis-classified to lake and
river (row 1-2) and river is mis-classified to canal (row 3).

Image

Ours

GroundTruth

Figure 10: Failure cases from the ATLANTIS val set.

###  5.4 Ablation studies

We also conduct ablation studies to compare a number of different model
variants of the proposed network, including the design of aquatic and non-
aquatic paths, and the two feature modulations. The results are shown in Table
4. We can see that the design of two paths can improve the performance of
waterbody image semantic segmentation. Moreover, both the proposed low-level
feature modulation (LM) and the cross-path modulation (CM) can achieve certain
performance gains in terms of acc and mIoU.

Table 4: Ablation study of each proposed component of AQUANet on the ATLANTIS dataset. Two Paths | LM | CM | A-acc | A-mIoU | acc | mIoU
---|---|---|---|---|---|---
|  |  | 67.27 | 44.53 | 73.28 | 38.81
$\surd$ |  |  | 67.73 | 47.89 | 75.29 | 40.28
$\surd$ | $\surd$ |  | 68.11 | 47.90 | 75.29 | 40.28
$\surd$ |  | $\surd$ | 67.21 | 46.63 | 75.81 | 40.57
$\surd$ | $\surd$ | $\surd$ | 68.85 | 48.42 | 76.18 | 40.83

Comparing to other state-of-the-art semantic segmentation models, AQUANet
provides highest mIoU and accuracy. In addition, by considering just labels
which includes water content, AQUANet still works better than others
(mIoU=50.34$\%$). In this regard, OCNet provides the second highest
performance (mIoU=47.89$\%$) and GCNet also provides the best results for
canal, lake and reservoir. These results approved AQUANet is well customized
for analysis of waterbodis and water-related scenes. However, considering mIoU
as a more informative metric for semantic segmentation task, all approaches
could only achieve 36.11%$\sim$42.22% mIoU on the proposed dataset. It shows
that the semantic segmentation of waterbodies and related objects is a
challenging task and needs more researches.

###  5.5 ATeX Experimental results

We further train ten well-known classification models including VGG [40],
ResNet [16], SqueezeNet [21], DenseNet [18], GoogLeNet [41], ShuffleNet v2
[29], MobileNetV2 [35], ResNeXt [44], Wide ResNet [49] and EfficientNet [42]
on the proposed ATeX dataset. All models are implemented using PyTorch. Cross-
entropy loss function is applied for training networks. We train all the
networks with the same 30 training epochs, SGD optimizer with a momentum of
0.9 and weight decay of 0.0001, and batch size is set to 64. For all networks,
the learning rate is first set to $2.5\times 10^{-4}$, then it is adjusted
based on decaying rate of the resulted loss function during training. Table 5
shows the training time and learning rate for each models over certain 30
epochs.

Table 5: The perfomance result on ATeX test set by well-known classification models. Networks | Time | LR | Val | Test
---|---|---|---|---
[mm:ss] | Acc. | Prec. | Recall | F1
Wide-ResNet-50-2 | 06:56 | 2.5E-4 | 91 | 77 | 75 | 75
VGG-16 | 04:38 | 2.5E-4 | 90 | 75 | 72 | 72
SqueezeNet | 00:47 | 7.5E-4 | 82 | 81 | 81 | 81
ShuffleNet V2 $\times 1.0$ | 01:46 | 1.0E-2 | 90 | 90 | 90 | 90
ResNeXt-50-$32\times 4d$ | 03:15 | 2.5E-4 | 90 | 77 | 75 | 75
ResNet-18 | 01:28 | 2.5E-4 | 87 | 74 | 72 | 72
MobileNet V2 | 01:35 | 2.5E-4 | 88 | 74 | 72 | 72
GoogLeNet | 01:51 | 5.0E-3 | 89 | 88 | 88 | 88
EffNet-B7 | 12:42 | 1.0E-2 | 90 | 91 | 91 | 91
EffNet-B0 | 02:38 | 7.5E-3 | 91 | 90 | 90 | 90
DenseNet-161 | 06:15 | 2.5E-4 | 91 | 81 | 79 | 79

Three common performance metrics including Precision, Recall and F1-score are
reported to evaluate the performance of the models on ATeX. Table 5 shows
weighted average (averaging the support-weighted mean per label) of these
three metrics on the test set. Accordingly, EffNet-B7, EffNet-B0 and
ShuffleNet V2 $\times 1.0$ provide the best results. Considering training
time, ShuffleNet V2 $\times 1.0$ can be presented as the most efficient
network.

##  6 Conclusion

In this paper, we introduced ATLANTIS, a large-scale dataset for semantic
segmentation of waterbodies and water-related scenes, by carefully collecting
images of diverse area from the internet and providing high-quality
annotations with the help of annotators major in water resources engineering.
We further provided comprehensive analysis of the characteristic of ATLANTIS
and reported the performance of current state-of-the-arts by training and
testing the networks on our dataset. A novel baseline network AQUANet is also
proposed for waterbody image semantic segmentation and achieves the best
performance on ATLANTIS. Additionally, we constructed ATLANTIS Texture (ATeX)
dataset that are sampled from ATLANTIS for texture-based classification and
evaluated several baseline methods on it.

ATLANTIS posed significant challenges for semantic segmentation which we
believe will boost new insights in both water resources engineering and
computer vision communities.

## References

  * Badrinarayanan et al. [2015] Vijay Badrinarayanan, Ankur Handa, and Roberto Cipolla.  Segnet: A deep convolutional encoder-decoder architecture for robust semantic pixel-wise labelling.  _arXiv preprint arXiv:1505.07293_ , 2015.
  * Bernal et al. [2017] Jorge Bernal, Nima Tajkbaksh, Francisco Javier Sánchez, Bogdan J Matuszewski, Hao Chen, Lequan Yu, Quentin Angermann, Olivier Romain, Bjørn Rustad, Ilangko Balasingham, et al.  Comparative validation of polyp detection methods in video colonoscopy: results from the miccai 2015 endoscopic vision challenge.  _IEEE Transactions on Medical Imaging_ , 36(6):1231–1249, 2017.
  * Caesar et al. [2018] Holger Caesar, Jasper Uijlings, and Vittorio Ferrari.  Coco-stuff: Thing and stuff classes in context.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 1209–1218, 2018\.
  * Cao et al. [2019] Yue Cao, Jiarui Xu, Stephen Lin, Fangyun Wei, and Han Hu.  Gcnet: Non-local networks meet squeeze-excitation networks and beyond.  In _Int. Conf. Comput. Vis. Worksh._ , pages 0–0, 2019.
  * Cervone et al. [2016] Guido Cervone, Elena Sava, Qunying Huang, Emily Schnebele, Jeff Harrison, and Nigel Waters.  Using twitter for tasking remote-sensing data collection and damage assessment: 2013 boulder flood case study.  _International Journal of Remote Sensing_ , 37(1):100–124, 2016.
  * Chen et al. [2017a] Liang-Chieh Chen, George Papandreou, Iasonas Kokkinos, Kevin Murphy, and Alan L Yuille.  Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs.  _IEEE Trans. Pattern Anal. Mach. Intell._ , 40(4):834–848, 2017a.
  * Chen et al. [2017b] Liang-Chieh Chen, George Papandreou, Florian Schroff, and Hartwig Adam.  Rethinking atrous convolution for semantic image segmentation.  _arXiv preprint arXiv:1706.05587_ , 2017b.
  * Cordts et al. [2016] Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele.  The cityscapes dataset for semantic urban scene understanding.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 3213–3223, 2016\.
  * Deng et al. [2009] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 248–255. IEEE, 2009.
  * Duan and Hu [2019] Lunhao Duan and Xiangyun Hu.  Multiscale refinement network for water-body segmentation in high-resolution satellite imagery.  _IEEE Geoscience and Remote Sensing Letters_ , 17(4):686–690, 2019.
  * Eltner et al. [2018] Anette Eltner, Melanie Elias, Hannes Sardemann, and Diana Spieler.  Automatic image-based water stage measurement for long-term observations in ungauged catchments.  _Water Resources Research_ , 54(12):10–362, 2018\.
  * Eltner et al. [2021] Anette Eltner, Patrik Olã Bressan, Thales Akiyama, Wesley Nunes Gonçalves, and José Marcato Junior.  Using deep learning for automatic water stage measurements.  _Water Resources Research_ , 57(3):e2020WR027608, 2021.
  * Everingham et al. [2010] Mark Everingham, Luc Van Gool, Christopher KI Williams, John Winn, and Andrew Zisserman.  The pascal visual object classes (voc) challenge.  _Int. J. Comput. Vis._ , 88(2):303–338, 2010\.
  * Fu et al. [2019] Jun Fu, Jing Liu, Haijie Tian, Yong Li, Yongjun Bao, Zhiwei Fang, and Hanqing Lu.  Dual attention network for scene segmentation.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 3146–3154, 2019\.
  * Gebrehiwot et al. [2019] Asmamaw Gebrehiwot, Leila Hashemi-Beni, Gary Thompson, Parisa Kordjamshidi, and Thomas E Langan.  Deep convolutional neural network for flood extent mapping using unmanned aerial vehicles data.  _Sensors_ , 19(7):1486, 2019.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 770–778, 2016\.
  * Hirabayashi et al. [2013] Yukiko Hirabayashi, Roobavannan Mahendran, Sujan Koirala, Lisako Konoshima, Dai Yamazaki, Satoshi Watanabe, Hyungjun Kim, and Shinjiro Kanae.  Global flood risk under climate change.  _Nature Climate Change_ , 3(9):816–821, 2013\.
  * Huang et al. [2017] Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger.  Densely connected convolutional networks.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 4700–4708, 2017\.
  * Huang et al. [2018] Xiao Huang, Cuizhen Wang, and Zhenlong Li.  A near real-time flood-mapping approach by integrating social media and post-event satellite imagery.  _Annals of GIS_ , 24(2):113–123, 2018.
  * Huang et al. [2019] Zilong Huang, Xinggang Wang, Lichao Huang, Chang Huang, Yunchao Wei, and Wenyu Liu.  Ccnet: Criss-cross attention for semantic segmentation.  In _Int. Conf. Comput. Vis._ , pages 603–612, 2019.
  * Iandola et al. [2016] Forrest N Iandola, Song Han, Matthew W Moskewicz, Khalid Ashraf, William J Dally, and Kurt Keutzer.  Squeezenet: Alexnet-level accuracy with 50x fewer parameters and $<$0.5 mb model size.  _arXiv preprint arXiv:1602.07360_ , 2016.
  * Jha et al. [2020] Debesh Jha, Pia H Smedsrud, Michael A Riegler, Pål Halvorsen, Thomas de Lange, Dag Johansen, and Håvard D Johansen.  Kvasir-seg: A segmented polyp dataset.  In _International Conference on Multimedia Modeling_ , pages 451–462. Springer, 2020.
  * Li et al. [2019a] Xia Li, Zhisheng Zhong, Jianlong Wu, Yibo Yang, Zhouchen Lin, and Hong Liu.  Expectation-maximization attention networks for semantic segmentation.  In _Int. Conf. Comput. Vis._ , pages 9167–9176, 2019a.
  * Li et al. [2019b] Ziyao Li, Rui Wang, Wen Zhang, Fengmin Hu, and Lingkui Meng.  Multiscale features supported deeplabv3+ optimization scheme for accurate water semantic segmentation.  _IEEE Access_ , 7:155787–155804, 2019b.
  * Lin et al. [2017] Guosheng Lin, Anton Milan, Chunhua Shen, and Ian Reid.  Refinenet: Multi-path refinement networks for high-resolution semantic segmentation.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 1925–1934, 2017\.
  * Lin et al. [2014] Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick.  Microsoft coco: Common objects in context.  In _Eur. Conf. Comput. Vis._ , pages 740–755. Springer, 2014.
  * Lo et al. [2015] Shi-Wei Lo, Jyh-Horng Wu, Fang-Pang Lin, and Ching-Han Hsu.  Visual sensing for urban flood monitoring.  _Sensors_ , 15(8):20006–20029, 2015.
  * Long et al. [2015] Jonathan Long, Evan Shelhamer, and Trevor Darrell.  Fully convolutional networks for semantic segmentation.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 3431–3440, 2015\.
  * Ma et al. [2018] Ningning Ma, Xiangyu Zhang, Hai-Tao Zheng, and Jian Sun.  Shufflenet v2: Practical guidelines for efficient cnn architecture design.  In _Eur. Conf. Comput. Vis._ , pages 116–131, 2018.
  * Mottaghi et al. [2014] Roozbeh Mottaghi, Xianjie Chen, Xiaobai Liu, Nam-Gyu Cho, Seong-Whan Lee, Sanja Fidler, Raquel Urtasun, and Alan Yuille.  The role of context for object detection and semantic segmentation in the wild.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 891–898, 2014\.
  * Muñoz et al. [2021] David F Muñoz, Paul Muñoz, Hamed Moftakhari, and Hamid Moradkhani.  From local to regional compound flood mapping with deep learning and data fusion techniques.  _Science of The Total Environment_ , 782:146927, 2021.
  * Neuhold et al. [2017] Gerhard Neuhold, Tobias Ollmann, Samuel Rota Bulo, and Peter Kontschieder.  The mapillary vistas dataset for semantic understanding of street scenes.  In _Int. Conf. Comput. Vis._ , pages 4990–4999, 2017.
  * Noh et al. [2015] Hyeonwoo Noh, Seunghoon Hong, and Bohyung Han.  Learning deconvolution network for semantic segmentation.  In _Int. Conf. Comput. Vis._ , pages 1520–1528, 2015.
  * Park et al. [2019] Taesung Park, Ming-Yu Liu, Ting-Chun Wang, and Jun-Yan Zhu.  Semantic image synthesis with spatially-adaptive normalization.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , 2019.
  * Sandler et al. [2018] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen.  Mobilenetv2: Inverted residuals and linear bottlenecks.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 4510–4520, 2018\.
  * Sarp et al. [2020] Salih Sarp, Murat Kuzlu, Mecit Cetin, Cem Sazara, and Ozgur Guler.  Detecting floodwater on roadways from image data using mask-r-cnn.  In _2020 International Conference on INnovations in Intelligent SysTems and Applications (INISTA)_ , pages 1–6. IEEE, 2020.
  * Sazara et al. [2019] Cem Sazara, Mecit Cetin, and Khan M Iftekharuddin.  Detecting floodwater on roadways from image data with handcrafted features and deep transfer learning.  In _2019 IEEE Intelligent Transportation Systems Conference (ITSC)_ , pages 804–809. IEEE, 2019.
  * Schnebele and Cervone [2013] Emily Schnebele and Guido Cervone.  Improving remote sensing flood assessment using volunteered geographical data.  _Natural Hazards and Earth System Sciences_ , 13(3):669–677, 2013.
  * Sekachev et al. [2020] Boris Sekachev, Nikita Manovich, Maxim Zhiltsov, Andrey Zhavoronkov, Dmitry Kalinin, Ben Hoff, TOsmanov, Dmitry Kruchinin, Artyom Zankevich, DmitriySidnev, Maksim Markelov, Johannes222, Mathis Chenuet, a andre, telenachos, Aleksandr Melnikov, Jijoong Kim, Liron Ilouz, Nikita Glazov, Priya4607, Rush Tehrani, Seungwon Jeong, Vladimir Skubriev, Sebastian Yonekura, vugia truong, zliang7, lizhming, and Tritin Truong.  opencv/cvat: v1.1.0, August 2020.  URL https://doi.org/10.5281/zenodo.4009388.
  * Simonyan and Zisserman [2014] Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_ , 2014.
  * Szegedy et al. [2015] Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich.  Going deeper with convolutions.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 1–9, 2015.
  * Tan and Le [2019] Mingxing Tan and Quoc Le.  Efficientnet: Rethinking model scaling for convolutional neural networks.  In _Int. Conf. on Mach. Learn._ , pages 6105–6114. PMLR, 2019.
  * Wang et al. [2018] Xintao Wang, Ke Yu, Chao Dong, and Chen Change Loy.  Recovering realistic texture in image super-resolution by deep spatial feature transform.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 606–615, 2018\.
  * Xie et al. [2017] Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He.  Aggregated residual transformations for deep neural networks.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 1492–1500, 2017\.
  * Yin et al. [2020] Minghao Yin, Zhuliang Yao, Yue Cao, Xiu Li, Zheng Zhang, Stephen Lin, and Han Hu.  Disentangled non-local neural networks.  In _Eur. Conf. Comput. Vis._ , pages 191–207. Springer, 2020.
  * Yu et al. [2020] Fisher Yu, Haofeng Chen, Xin Wang, Wenqi Xian, Yingying Chen, Fangchen Liu, Vashisht Madhavan, and Trevor Darrell.  Bdd100k: A diverse driving dataset for heterogeneous multitask learning.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 2636–2645, 2020\.
  * Yuan and Wang [2018] Yuhui Yuan and Jingdong Wang.  Ocnet: Object context network for scene parsing.  _arXiv preprint arXiv:1809.00916_ , 2018.
  * Yuan et al. [2020] Yuhui Yuan, Xilin Chen, and Jingdong Wang.  Object-contextual representations for semantic segmentation.  In _Eur. Conf. Comput. Vis._ , 2020.
  * Zagoruyko and Komodakis [2016] Sergey Zagoruyko and Nikos Komodakis.  Wide residual networks.  _arXiv preprint arXiv:1605.07146_ , 2016.
  * Zhao et al. [2017] Hengshuang Zhao, Jianping Shi, Xiaojuan Qi, Xiaogang Wang, and Jiaya Jia.  Pyramid scene parsing network.  In _IEEE Conf. Comput. Vis. Pattern Recog._ , pages 2881–2890, 2017\.
  * Zhou et al. [2019] Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba.  Semantic understanding of scenes through the ade20k dataset.  _Int. J. Comput. Vis._ , 127(3):302–321, 2019\.
  * Zhu et al. [2019] Zhen Zhu, Mengde Xu, Song Bai, Tengteng Huang, and Xiang Bai.  Asymmetric non-local neural networks for semantic segmentation.  In _Int. Conf. Comput. Vis._ , pages 593–602, 2019.


# Learning Slice-Aware Representations with Mixture of Attentions

Cheng Wang Sungjin Lee Sunghyun Park Han Li Young-Bum Kim Ruhi Sarikaya
Amazon Alexa AI
<EMAIL_ADDRESS>

###### Abstract

Real-world machine learning systems are achieving remarkable performance in
terms of coarse-grained metrics like overall accuracy and F-1 score. However,
model improvement and development often require fine-grained modeling on
individual data subsets or slices, for instance, the data slices where the
models have unsatisfactory results. In practice, it gives tangible values for
developing such models that can pay extra attention to critical or interested
slices while retaining the original overall performance. This work extends the
recent slice-based learning (SBL) Chen et al. (2019) with a mixture of
attentions (MoA) to learn slice-aware dual attentive representations. We
empirically show that the MoA approach outperforms the baseline method as well
as the original SBL approach on monitored slices with two natural language
understanding (NLU) tasks.

⬇ def SF_Long(sentence,k=10): n = len(sentence .split(’ ’)) return n > k def
SF_Question(sentence): return sentence[-1] == ’?’
---
Table 1: The designed slice functions (SFs)333The SFs are task-dependent and
not assumed to be perfectly accurate. They can be noisy or from weak
supervision sources Ratner et al. (2016). Here, for the task in sec.4.2, SFs
are defined to improve the slices where the model has unsatisfactory results
as compared to the overall performance. For the task in sec.4.3, we define the
SFs for the slices of interest.. Left: We monitor three data slices - short
utterances, those involving “time”, and those involving “email”. Right: We
monitor long sentences and questions444Alternatively, 5W1H rule for questions
Kim et al. (2019)..

##  1 Introduction

Though machine learning systems have been achieving excellent performance in
terms of coarse-grained metrics like accuracy, they perform poorly or even
fail on some individual data subsets (i.e., slices). For instance, many models
have difficulties when learning for classes with only a few samples or samples
with challenging structures.

Figure 1:  The slice-aware architecture with MoA. It consists of six
components: (1) slice functions define the special data slices that we want to
monitor; (2) backbone model for feature extraction (e.g., BERT); (3) slice
indicators are membership functions to predict if a sample belongs to the
slice; (4) slice experts aim to learn slice-specific representations; (5)
shared head is the base task predictive layer across experts and (6) the
proposed mixture of attentions (MoA) learns to attend to the slices of
interest. It contains two different attention mechanism (red boxes): a
membership attention and a dot-product attention. The MoA learns to re-weight
the expert representation $\mathbf{r}$ to a slice-aware representation
$\mathbf{s}$ and the original representation $\mathbf{x}$ to $\mathbf{s}$
(yellow lines). The slice distributions are computed in deterministic
(weighted sum of slices) or stochastic (sampling) way in re-weighting
$\mathbf{r}$ and $\mathbf{x}$.

Inspecting particular data slices can serve as an important component in model
development cycles. A recently proposed slice-based learning (SBL) exhibited
compelling results with more than 3% improvements on pre-defined slices Chen
et al. (2019) in the task of binary classification. However, one potential
limitation of the existing attention mechanism in SBL is that in multi-class
cases, the attention suffers from the difficulty in using the experts’
confidences appropriately for computing slice distributions (refer to Sec. 3).

In this paper, we extend SBL with a mixture of attentions (MoA) mechanism. Two
different attention mechanisms are learned to jointly attend to the defined
slices from different representations in different latent subspaces. The first
attention is based on slice membership likelihood and/or experts confidence as
in SBL Chen et al. (2019), which we call membership attention. The second one
is dot-product attention that is based on the backbone model (e.g., BERT
Devlin et al. (2019)) extracted representations. The MoA approach is akin to
multi-head attention Vaswani et al. (2017) but with different attention types
that receive different inputs.

As presented in Figure 4, the two attentions in MoA can work jointly to attend
to (1) the expert representation $\mathbf{r}$ and (2) the backbone model
extracted representation $\mathbf{x}$, and finally form an attentive
representation $\mathbf{s}$. The $\mathbf{s}$ is a slice-aware featurization
of the samples in the particular data slices and will be used for making a
final model prediction.

We argue that learning joint attention with MoA from different resources for
computing slice distributions is beneficial Vaswani et al. (2017); Li et al.
(2018). We evaluate the effectiveness of our proposed approach on intent
detection Liu et al. (2019) and linguistic acceptability Warstadt et al.
(2018) tasks.

Our main contributions are twofold:

  * •

We extend SBL with MoA. The MoA approach has the ability to attend to slices
in deterministic (weighted summation) and stochastic (sampling) ways.

  * •

We conduct extensive experiments on two NLU tasks. The results show that MoA
outperforms the baseline and vanilla SBL by average up to 9% and 6%
respectively on defined slices.

##  2 Architecture

Figure 1 presents the slice-aware architecture based on SBL Chen et al.
(2019). Let $\\{x^{n},y^{n}\\}_{n}^{N}$ be a dataset with N samples. We aim to
learn slice-aware representation $\mathbf{s}$ from slice-experts-learned
representation $\mathbf{r}$ and backbone-model-extracted representation
$\mathbf{x}$.

We first define slice functions (SFs) as in Table 4 to split the dataset into
$k$ slices of interests. Each sample is assigned with a slice label
$\gamma\in[0,1]$ in $\\{\gamma_{1},\gamma_{2},...,\gamma_{k}\\}$ as
supervision data555$s_{1}$ is the base slice, and $s_{2}$ to $s_{k}$ are the
slices of interest..

Second, we use a backbone model like BERT to extract representation
$\mathbf{x}\in\mathbf{R}^{d}$ for a given sample. Then, slice indicators
$f_{i}(\mathbf{x};\mathbf{w}^{f}_{i})$,
$\mathbf{w}^{f}_{i}\in\mathcal{R}^{d\times 1}$, $i\in\\{1,..,k\\}$ map
$\mathbf{x}$ to a prediction $h_{i}$. $f_{i}$ are trained with
$\\{\mathbf{x}^{n},\gamma^{n}\\}_{n}^{N}$ to predict whether a sample belongs
to a particular slice. They are learned with cross entropy loss

$\zeta_{1}=\sum_{i}^{k}\mathcal{L}_{CE}(h_{i},\gamma_{i})$ (1)

Then, slice experts $g_{i}(\mathbf{x};\mathbf{w}^{g}_{i})$,
$\mathbf{w}^{g}_{i}\in\mathcal{R}^{d\times d}$ learn a mapping from
$\mathbf{x}$ to a slice vector $r_{i}\in\mathcal{R}^{d}$ with the samples that
only belong to the slice, followed by a shared head, which is shared across
all experts and maps $r_{i}$ to a prediction
$\hat{y}=\varphi(r_{i};\mathbf{w}_{s})$. $g_{i}$ and $\varphi$ are learned on
the base (original) task with ground-truth label $y$ by

$\zeta_{2}=\sum_{i}^{k}\gamma_{i}\mathcal{L}_{CE}(\hat{y},y)$ (2)

Finally, a mixture of attentions(MoA) (as in Sec. 3) re-weights $\mathbf{r}$
and $\mathbf{x}$ to form $\mathbf{s}$. The $\mathbf{s}$ goes through a final
prediction function $\eta$ on the base task. The loss function is

$\zeta_{3}=\mathcal{L}_{CE}(\eta(\mathbf{s};\mathbf{w}_{p}),y)$ (3)

The total loss is a combination of the loss for slice indicators, slice
experts and base task prediction function:

$\zeta=\zeta_{1}+\zeta_{2}+\zeta_{3}$ (4)

The whole model is optimised with back-propagation Rumelhart et al. (1986) in
an end-to-end way.

##  3 Methodology

The SBL approach Chen et al. (2019) proposed a slice-residual attention
modules (SRAMs) that are directly based on stacked membership likelihood
$H\in\mathbf{R}^{k}$ and experts’ prediction confidence
$\left|Y\right|\in\mathbf{R}^{c\times k}$, $c=1$ (i.e., binary
classification). Then, slice distribution (attention weights) is computed with
$a=\textsc{softmax}(H+\left|Y\right|)$. One potential limitation of this
mechanism is that the above formulation can lead to mismatch shape in element-
wise addition when $c>2$ (i.e., multi-class classification). To circumvent
this, we propose a mixture of attentions (MoA) to augment membership attention
with dot-product attention from different information resources.

###  3.1 Mixture of Attentions

Let $\mathbf{x}\in\mathbf{R}^{d}$ be the original representation from the
backbone model (e.g., BERT), $h_{i}\in\mathbf{R}^{c}$ ($c=1$) as i-th
indicator function’s prediction, and $r_{i}\in\mathbf{R}^{d}$ as i-th expert
learned representation. When stacking on $k$ slices, we have
$\mathbf{h}\in\mathbf{R}^{c\times k}$ and $\mathbf{r}\in\mathbf{R}^{d\times
k}$. MoA’s goal is to (1) attend to $\mathbf{r}$ based on indicator functions’
membership likelihood and/or experts confidence666In multi-class case, only
membership likelihood is used.; (2) attend to $\mathbf{x}$ with a dot-product
attention; (3) to form a new slice-aware attentive representation
$\mathbf{s}\in\mathbf{R}^{d}$ with weighted (sampled) $\mathbf{r}$ and
$\mathbf{x}$.

The slice distributions are computed differently. For membership attention,
the probability $\mathbf{p}_{1}=\textsc{softmax}(\mathbf{h})$ or
$\mathbf{p}_{1}=\textsc{softmax}(\mathbf{h}+\left|\mathbf{r}\right|)\in\mathbf{R}^{k}$
($d$=1 in binary classification). Then membership weighted slice
representation is computed:
$\mathbf{s}_{1}=\mathbf{r}\mathchoice{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\displaystyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\textstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptscriptstyle\bullet$}}}}}\mathbf{p}_{1},\mathbf{s}_{1}\in\mathbf{R}^{d}$.
For dot-product attention, we aim to learn an attention matrix
$\mathbf{A}=\\{\mathbf{a}_{1},...,\mathbf{a}_{k}\\},\mathbf{a}\in\mathbf{R}^{d},\mathbf{A}\in\mathbf{R}^{d\times
k}$ is randomly initialized and learned by the standard back-propagation.
Intuitively, each $\mathbf{a}$ is learned to be a slice prototype Wang and
Niepert (2019); Roy et al. (2020). The probability over slices is computed as:

$\mathbf{p}_{2}=\textsc{softmax}(\mathbf{A}^{\top}\mathchoice{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\displaystyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\textstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptscriptstyle\bullet$}}}}}\mathbf{x})\in\mathbf{R}^{k}$
(5)

A new attentive representation $\mathbf{s}_{2}$ is formed by weighting
$\mathbf{A}$ with $\mathbf{p}_{2}$:

$\mathbf{s}_{2}=\mathbf{A}\mathchoice{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\displaystyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\textstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptscriptstyle\bullet$}}}}}\mathbf{p}_{2},~{}~{}~{}\mathbf{s}_{2}\in\mathbf{R}^{d}$
(6)

or sampling from $\mathbf{A}$:

$\mbox{sample }\mathbf{s}_{2}\sim\\{\mathbf{a}_{1},...,\mathbf{a}_{k}\\}$ (7)

Then slice-aware vector $\mathbf{s}$ is computed by

$\mathbf{s}=\mathbf{s}_{1}\circledcirc\mathbf{s}_{2}$ (8)

where $\circledcirc$ is an operator (either $\oplus$: element-wise addition or
$\otimes$: element-wise multiplication). The eq.(8) can be extended into a
more general form – mixture of attentions (MoA):

$\mathbf{s}=\underbrace{\mathbf{r}\mathchoice{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\displaystyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\textstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptscriptstyle\bullet$}}}}}\phi(\mathbf{h})}_{\textbf{{membership}}}\circledcirc~{}\underbrace{\mathbf{A}\mathchoice{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\displaystyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\textstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptscriptstyle\bullet$}}}}}\phi(\mathbf{A}^{\top}\mathchoice{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\displaystyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\textstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptstyle\bullet$}}}}}{\mathbin{\vbox{\hbox{\scalebox{0.6}{$\scriptscriptstyle\bullet$}}}}}\mathbf{x})}_{\textbf{{dot-
product}}}$ (9)

Note eq.(9) entails the following transformations ($\to$) and captures the
representational differences from $\mathbf{r}$ to $\mathbf{s}$ and from
$\mathbf{x}$ to $\mathbf{s}$:

$\displaystyle\mathbf{x}\leavevmode\hbox to6.67pt{\vbox
to6.67pt{\pgfpicture\makeatletter\hbox{\hskip 3.33301pt\lower-3.33301pt\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}}
{{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{0.0pt}{0.0pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox
to0.0pt{}{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}\to\mathbf{r}\leavevmode\hbox
to6.67pt{\vbox to6.67pt{\pgfpicture\makeatletter\hbox{\hskip
3.33301pt\lower-3.33301pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}}
{{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{0.0pt}{0.0pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox
to0.0pt{}{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}\to\mathbf{p}_{1}\to\mathbf{s}_{1}\leavevmode\hbox
to6.67pt{\vbox to6.67pt{\pgfpicture\makeatletter\hbox{\hskip
3.33301pt\lower-3.33301pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}}
{{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{0.0pt}{0.0pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox
to0.0pt{}{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}\to\leavevmode\hbox
to6.67pt{\vbox to6.67pt{\pgfpicture\makeatletter\hbox{\hskip
3.33301pt\lower-3.33301pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}}
{{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{0.0pt}{0.0pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox
to0.0pt{}{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}\mathbf{s}$ (10)
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}\leavevmode\hbox
to6.67pt{\vbox to6.67pt{\pgfpicture\makeatletter\hbox{\hskip
3.33301pt\lower-3.33301pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}}
{{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{0.0pt}{0.0pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox
to0.0pt{}{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}\mathbf{p}_{2}\to\mathbf{s}_{2}\leavevmode\hbox
to6.67pt{\vbox to6.67pt{\pgfpicture\makeatletter\hbox{\hskip
3.33301pt\lower-3.33301pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}}
{{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{0.0pt}{0.0pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox
to0.0pt{}{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}\leavevmode\hbox
to1.5pt{\vbox to3.2pt{\pgfpicture\makeatletter\hbox{\hskip
1.19998pt\lower-1.59998pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ } { {}{}{}{}{}}{}{{}}{}{ {}{}{}{}{}}
{}{}\pgfsys@beginscope\pgfsys@invoke{
}\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@invoke{
}\pgfsys@color@gray@fill{0}\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0,0,0}{}{{
{\pgfsys@beginscope\pgfsys@setlinewidth{0.32pt}\pgfsys@setdash{}{0.0pt}\pgfsys@roundcap\pgfsys@roundjoin{}
{}{}{} {}{}{}
\pgfsys@moveto{-1.19998pt}{1.59998pt}\pgfsys@curveto{-1.09998pt}{0.99998pt}{0.0pt}{0.09999pt}{0.29999pt}{0.0pt}\pgfsys@curveto{0.0pt}{-0.09999pt}{-1.09998pt}{-0.99998pt}{-1.19998pt}{-1.59998pt}\pgfsys@stroke\pgfsys@endscope}}
}{}{}{{}}{}{}{{}}\pgfsys@moveto{3.0pt}{0.0pt}\pgfsys@lineto{3.45999pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
}{{}{{}}{}{}{{}}{{{}}{{{}}{\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-3.45999pt}{0.0pt}\pgfsys@invoke{
}\pgfsys@invoke{ \lxSVG@closescope }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}}{{}}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox
to0.0pt{}{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}\leavevmode\hbox
to0.4pt{\vbox to0.4pt{\pgfpicture\makeatletter\hbox{\hskip
0.2pt\lower-0.2pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ } { {}{}{}{}{}}{}{ {}{}{}{}{}}
{{{{{}}{ {}{}}{{}{}} {}{}}{}{}}}{}{{{{{}}{ {}{}}{{}{}}
{}{}}{}{}}}{{}}{}{}{}\pgfsys@beginscope\pgfsys@invoke{
}\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@invoke{
}\pgfsys@color@gray@fill{0}\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0,0,0}{}{}{}{}{{}}\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@lineto{0.45999pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
}{{}{{}}{}{}{{}}{{{}}{{{}}{\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-0.45999pt}{0.0pt}\pgfsys@invoke{
}\pgfsys@invoke{ \lxSVG@closescope }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}}{{}}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox
to0.0pt{}{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}\leavevmode\hbox
to16.5pt{\vbox to8.45pt{\pgfpicture\makeatletter\hbox{\hskip
8.24765pt\lower-8.24765pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ } { {}{}{}{}{}}{}{{}}{} {
{}{}{}{}{}}{{}{}}{{}} {{{}}{{}}}{{}}{ {}{}{}{}{}}{{{}}{{}}}{
{}{}{}{}{}}{}{{}}{}{}{}\pgfsys@beginscope\pgfsys@invoke{
}\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@invoke{
}\pgfsys@color@gray@fill{0}\pgfsys@invoke{
}\definecolor[named]{pgffillcolor}{rgb}{0,0,0}{}{}{}{}{{}}{}{}{{}}\pgfsys@moveto{2.12132pt}{-2.12132pt}\pgfsys@curveto{8.04765pt}{-8.04765pt}{-8.04765pt}{-8.04765pt}{-2.44658pt}{-2.44658pt}\pgfsys@stroke\pgfsys@invoke{
}{{}{{}}{}{}{{}}{{{}}{{{}}{\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{0.7071}{0.7071}{-0.7071}{0.7071}{-2.44658pt}{-2.44658pt}\pgfsys@invoke{
}\pgfsys@invoke{ \lxSVG@closescope }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}}{{}}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hbox
to0.0pt{}{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}\vspace{-3mm}$ (11)

The $\phi(\cdot)$ is either softmax:
$p_{i}=\frac{\exp(z_{i})}{\sum_{j}^{k}\exp(z_{j})}$ that deterministically
computes slice distributions or a Monte-Carlo gradient estimator: gumbel-
softmax Gumbel (1954); Jang et al. (2017); Maddison et al. (2017):

$p_{i}=\frac{\exp[(\log(z_{i})+\pi_{i})/\tau]}{\sum_{j}^{k}\exp[\log(z_{j})+\pi_{j})/\tau]}$
(12)

The $\pi_{i}$ are i.i.d. samples from the Gumbel$(0,1)$, that is,
$\pi=-\log(-\log(u)),u\sim\textsc{uniform}(0,1)$. $\tau$ is temperature which
controls the concentration of slice distribution, and small $\tau$ leads to
more confident prediction over slices. It aims to stochastically compute slice
distribution. With Gumbel-softmax, the slice distribution is a soft sampling
from:

$\displaystyle\mathbf{p}_{1}\sim\textsc{Gumbel-Softmax}(\mathbf{h})$ (13)
$\displaystyle\mathbf{p}_{2}\sim\textsc{Gumbel-
Softmax}(\mathbf{A}^{\top}\cdot\mathbf{x})$ (14)

or a hard sampling (but differentiable) from:

$\displaystyle\mathbf{p}_{1}\sim\textsc{One-hot}(\arg\max(\mathbf{p}_{1}))$
(15) $\displaystyle\mathbf{p}_{2}\sim\textsc{One-
hot}(\arg\max(\mathbf{p}_{2}))$ (16)

for membership and dot-product attention respectively.

(a)  Figure 2:  Distributions over slices (base, $s1$= Length, $s2$=Time and
$s3$=Email) of random test samples that are from membership and dot-product
attention mechanisms. In (a)(c) membership attention shows higher confidence
while dot-product attention gives higher confidence in (b)(d). Top rows are
the utterances from the test set.

##  4 Experiments

We performed our experiments on a binary classification task with linguistic
acceptability and on a multi-class classification task with intent detection.

| F1 | F1 Lift(%) | MCC | MCC Lift(%)
---|---|---|---|---
Methods | Overall | S1 | S2 | Avg. | Max. | Overall | S1 | S2 | Avg. | Max.
Baseline | 0.70 | 0.60 | 0.65 | — | — | 0.24 | 0.18 | 0.22 | — | —
SBL | 0.69 | 0.71 | 0.72 | 9.0 % | 11.0% | 0.23 | 0.20 | 0.24 | 2.0% | 2.0%
SBL-MoA $\oplus$ | 0.70 | 0.71 | 0.68 | 7.0% | 11.0% | 0.25 | 0.24 | 0.20 | 2.0% | 6.0%
SBL-MoA $\otimes$ | 0.69 | 0.72 | 0.71 | 9.0% | 12.0% | 0.24 | 0.24 | 0.25 | 4.5% | 6.0%
SBL-MoA-S $\oplus$ | 0.69 | 0.67 | 0.68 | 5.0% | 7.0% | 0.24 | 0.25 | 0.18 | 1.5% | 7.0%
SBL-MoA-S $\otimes$ | 0.69 | 0.69 | 0.71 | 7.5% | 9.0% | 0.26 | 0.28 | 0.22 | 5.0% | 10.0%
SBL-MoA-H $\oplus$ | 0.70 | 0.69 | 0.70 | 7.0% | 9.0% | 0.25 | 0.28 | 0.26 | 8.0% | 10.0%
SBL-MoA-H $\otimes$ | 0.69 | 0.69 | 0.65 | 4.5% | 9.0% | 0.25 | 0.22 | 0.32 | 7.0% | 10.0%
Table 2: The results on CoLA test datasets. F1-score and MCC are reported
(averaged on 5 random runs for each model). $s1$=Long, and $s2$=Question.The
lift is the averaged relative improvement across slices over baseline. The
largest improvement is in bold and second largest lift number is in underline
(same to Table 3).

###  4.1 Experimental Setup

Datasets and Metrics. The CoLA Warstadt et al. (2018) dataset has 8551 train
and 527 development in domain samples777https://nyu-mll.github.io/CoLA/. We
randomly split it into train/val/test with 7200/878/1000 samples. As in Chen
et al. (2019), we ensure the sample proportion in ground-truth are consistent
across splits. We use F1-score and Matthews correlation coefficient (MCC)
Matthews (1975) as our metrics. The NLU dataset Liu et al. (2019) for intent
detection contains 25k user utterances across 64 intents. We randomly split it
into train/val/test with ratio 0.7:0.1:0.2. We use the accuracy and F1-score
as our metrics.

|  | Acc | Acc Lift(%) | F1 | F1 Lift(%)
---|---|---|---|---|---
Methods | Overall | S1 | S2 | S3 | Avg. | Max. | Overall | S1 | S2 | S3 | Avg. | Max.
Baseline | 0.7413 | 0.73 | 0.74 | 0.73 | — | — | 0.7404 | 0.74 | 0.72 | 0.74 | — | —
SBL | 0.7422 | 0.75 | 0.76 | 0.75 | 2.0% | 2.0% | 0.7418 | 0.74 | 0.72 | 0.75 | 0.3% | 1.0%
SBL-MoA $\oplus$ | 0.7414 | 0.74 | 0.77 | 0.74 | 1.7% | 3.0% | 0.7390 | 0.74 | 0.73 | 0.74 | 0.3% | 1.0%
SBL-MoA $\otimes$ | 0.7440 | 0.80 | 0.73 | 0.76 | 3.0% | 7.0% | 0.7411 | 0.77 | 0.74 | 0.74 | 1.7% | 3.0%
SBL-MoA-S $\oplus$ | 0.7403 | 0.74 | 0.75 | 0.73 | 0.7% | 1.0% | 0.7403 | 0.73 | 0.73 | 0.76 | 0.7% | 2.0%
SBL-MoA-S $\otimes$ | 0.7424 | 0.75 | 0.72 | 0.75 | 0.7% | 2.0% | 0.7421 | 0.75 | 0.74 | 0.74 | 1.0% | 2.0%
SBL-MoA-H $\oplus$ | 0.7405 | 0.73 | 0.74 | 0.73 | 0.0% | 0.0% | 0.7397 | 0.76 | 0.74 | 0.76 | 2.0% | 2.0%
SBL-MoA-H $\otimes$ | 0.7418 | 0.74 | 0.73 | 0.75 | 0.7% | 2.0% | 0.7401 | 0.75 | 0.74 | 0.74 | 1.0% | 2.0%
Table 3: The results on intent detection. Accuracy and F1 scores are reported.
$s1$=Length, $s2$=Time, $s3$=Email are the slices that we monitor and aim to
improve. The experts’ confidence scores are not used as discussed in Sec.3.

Compared Methods. We implemented and compared the following methods:

  * •

Baseline: A three-layer feed-forward network.

  * •

SBL: Slice-based learning Chen et al. (2019).

  * •

SBL-MoA: Our approach that extends SBL with a mixture of attentions (MoA).

For SBL-MoA, we developed multiple variants with Gumbel-Softmax. SBL-MoA-S
(SBL-MoA-H) are the variant models with soft (hard) sampling from a Gumbel-
Softmax distributions. We also tested the way that membership attention and
dot-product interact with each other with $\oplus$ (element-wise addition) and
$\otimes$ (element-wise multiplication).

Implementation Details. BERT-base Devlin et al. (2019) in sentence-transformer
Thakur et al. (2020) is used as the backbone model. We use 128 hidden units
for all models, which are implemented with Pytorch Paszke et al. (2019). A
dropout (p=0.5)888As the data size is relatively small, we use strong dropout
regularization to prevent overfitting. is applied after input layer. The
models are trained with Adam (0.001) Kingma and Ba (2014), with weight decay
of 0.01 and 0.001 for the two tasks, respectively. All models are trained with
a maximum of 500 epochs with early stopping (patience=50). The best models are
selected based on model performance on the validation sets. The temperature
$\tau=1.0$ is fixed in all the experiments.

###  4.2 Results on Linguistic Acceptability

Table 2 presents the results on CoLA. First, slice-based models (i.e., SBL,
SBL-MoA, and its variants) show that they can maintain (or improve) the
original overall performance. Second, we observe that they achieve obvious
performance lift on the monitored slices. For instance, SBL achieves an
average 9% F1 score over the baseline. The proposed method (SBL-MoA,
$\otimes$) achieves an average of 9% and maximum 12% lift. For MCC, the best
performer is SBL-MoA-H, which achieves an average $\geqslant$7% and maximum
10% as compared to the baseline. It outperforms SBL by $\geqslant$ 5%. Also,
we notice that using operator $\otimes$ (element-wise multiplication) between
the attention mechanisms lead to better performance as compared to $\oplus$.

###  4.3 Results on Intent Detection

Table 3 demonstrates that both SBL and SBL-MoA improve model performance on
the monitored slices, with a similar (slightly better) overall performance on
the base task999Note the lift on slice can be negligible to overall due to
small size of slice data, e.g., For SBL-MoA $\otimes$, $s1$ with 122 samples,
7.0% lift only contributes to 122$\times$0.07/5124 $\approx$ 0.0017.. SBL-MoA
variants achieve the best scores and outperform SBL by average 1% accuracy and
1.7% F1.

Figure 2 illustrates the slice distributions given some random samples. We
denote $\mathbf{p}_{1}$ and $\mathbf{p}_{2}$ for membership and dot-product
attention respectively. The experiments show that $\mathbf{p}_{1}$ and
$\mathbf{p}_{2}$ reach an agreement on predicting the correct slices.
Interestingly, the sample in (d) — “write sms to our friends”, in principle,
should be sliced as “base”, but both attentions exhibit high confidence to
$s3$=“Email”. We conjecture the reason is that all utterances are encoded with
BERT which captures the similarity between the sample and the utterances in
the “Email” slice.

##  5 Related Work

SBL Chen et al. (2019) is a novel programming model for critical data slices.
It is an instance of weakly supervised learning Zhou (2018); Medlock and
Briscoe (2007). The weak supervision data are generated from pre-defined
labeling functions Ratner et al. (2016). SBL has shown better predictive
performance compared to the mixture of experts Jacobs et al. (1991) and multi-
task learning Caruana (1997), with reduced run-time cost and parameters Chen
et al. (2019). The concept of SBL has been recently used in many applications.
Penha et al. Penha and Hauff (2020) proposed to adapt SBL to improve ranking
performance and capture the failures of the ranker model. Wang et al. Wang et
al. (2021) recently implemented SBL in a commercial conversational AI system
in order to handle the long-tail problem of imbalanced distribution in
customer queries and further improved the performance of the conversational
skill routing components Li et al. (2021); Kim et al. (2018b, a).

Our proposed mixture of attention (MoA) is an instance of multi-head attention
Vaswani et al. (2017) but with different attention types. MoA can also be
extended to include other attention types. We have shown the effectiveness of
this mechanism in determining the slice distributions.

##  6 Conclusion

This paper extends SBL with MoA (SBL-MoA) to improve model performance on
particular data slices. We empirically show that SBL-MoA yields better slice
level performance lift to baseline and vanilla SBL with two NLU tasks:
linguistic acceptability and intent detection.

## References

  * Caruana (1997) Rich Caruana. 1997.  Multitask learning.  _Machine learning_ , 28(1):41–75.
  * Chen et al. (2019) Vincent Chen, Sen Wu, Alexander J Ratner, Jen Weng, and Christopher Ré. 2019\.  Slice-based learning: A programming model for residual learning in critical data slices.  In _Advances in neural information processing systems_ , pages 9397–9407.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  Bert: Pre-training of deep bidirectional transformers for language understanding.  In _NAACL-HLT_.
  * Gumbel (1954) Emil Julius Gumbel. 1954.  Statistical Theory of Extreme Values and Some Practical Applications. A Series of Lectures.  _Number 33. US Govt. Print. Office_.
  * Jacobs et al. (1991) Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. 1991\.  Adaptive mixtures of local experts.  _Neural computation_ , 3(1):79–87.
  * Jang et al. (2017) Eric Jang, Shixiang Gu, and Ben Poole. 2017.  Categorical reparameterization with gumbel-softmax.  _ICLR_.
  * Kim et al. (2019) Najoung Kim, Roma Patel$\varphi$, Adam Poliak, Alex Wang, Patrick Xia, R Thomas McCoy, Ian Tenney, Alexis Ross, Tal Linzen, Benjamin Van Durme, et al. 2019.  Probing what different nlp tasks teach machines about function word comprehension.  _NAACL HLT 2019_ , page 235.
  * Kim et al. (2018a) Young-Bum Kim, Dongchan Kim, Joo-Kyung Kim, and Ruhi Sarikaya. 2018a.  A scalable neural shortlisting-reranking approach for large-scale domain classification in natural language understanding.  _arXiv preprint arXiv:1804.08064_.
  * Kim et al. (2018b) Young-Bum Kim, Dongchan Kim, Anjishnu Kumar, and Ruhi Sarikaya. 2018b.  Efficient large-scale neural domain classification with personalized attention.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 2214–2224.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba. 2014.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_.
  * Li et al. (2021) Han Li, Sunghyun Park, Aswarth Dara, Jinseok Nam, Sungjin Lee, Young-Bum Kim, Spyros Matsoukas, and Ruhi Sarikaya. 2021.  Neural model robustness for skill routing in large-scale conversational ai systems: A design choice exploration.  _arXiv preprint arXiv:2103.03373_.
  * Li et al. (2018) Jian Li, Zhaopeng Tu, Baosong Yang, Michael R Lyu, and Tong Zhang. 2018.  Multi-head attention with disagreement regularization.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pages 2897–2903.
  * Liu et al. (2019) Xingkun Liu, Arash Eshghi, Pawel Swietojanski, and Verena Rieser. 2019.  Benchmarking natural language understanding services for building conversational agents.  In _10th International Workshop on Spoken Dialogue Systems Technology 2019_.
  * Maddison et al. (2017) Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. 2017.  The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables.  In _5th International Conference on Learning Representations (ICLR)_.
  * Matthews (1975) Brian W Matthews. 1975.  Comparison of the predicted and observed secondary structure of t4 phage lysozyme.  _Biochimica et Biophysica Acta (BBA)-Protein Structure_ , 405(2):442–451.
  * Medlock and Briscoe (2007) Ben Medlock and Ted Briscoe. 2007.  Weakly supervised learning for hedge classification in scientific literature.  In _Proceedings of the 45th annual meeting of the association of computational linguistics_ , pages 992–999.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019.  Pytorch: An imperative style, high-performance deep learning library.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, _Advances in Neural Information Processing Systems 32_ , pages 8024–8035.
  * Penha and Hauff (2020) Gustavo Penha and Claudia Hauff. 2020.  Slice-aware neural ranking.  In _Proceedings of the 5th International Workshop on Search-Oriented Conversational AI (SCAI)_ , pages 1–6.
  * Ratner et al. (2016) Alexander J Ratner, Christopher M De Sa, Sen Wu, Daniel Selsam, and Christopher Ré. 2016.  Data programming: Creating large training sets, quickly.  _Advances in neural information processing systems_ , 29:3567–3575.
  * Roy et al. (2020) Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. 2020.  Efficient content-based sparse attention with routing transformers.  _arXiv preprint arXiv:2003.05997_.
  * Rumelhart et al. (1986) David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. 1986.  Learning representations by back-propagating errors.  _nature_ , 323(6088):533–536.
  * Thakur et al. (2020) Nandan Thakur, Nils Reimers, Johannes Daxenberger, and Iryna Gurevych. 2020.  Augmented sbert: Data augmentation method for improving bi-encoders for pairwise sentence scoring tasks.  _arXiv preprint arXiv:2010.08240_.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In _Advances in neural information processing systems_ , pages 5998–6008.
  * Wang et al. (2021) Cheng Wang, Sun Kim, Taiwoo Park, Sajal Choudhary, Sunghyun Park, Young-Bum Kim, Ruhi Sarikaya, and Sungjin Lee. 2021.  Handling long-tail queries with slice-aware conversational systems.  _arXiv preprint arXiv:2104.13216_.
  * Wang and Niepert (2019) Cheng Wang and Mathias Niepert. 2019.  State-regularized recurrent neural networks.  In _International Conference on Machine Learning_ , pages 6596–6606.
  * Warstadt et al. (2018) Alex Warstadt, Amanpreet Singh, and Samuel R Bowman. 2018.  Neural network acceptability judgments.  _arXiv preprint arXiv:1805.12471_.
  * Zhou (2018) Zhi-Hua Zhou. 2018.  A brief introduction to weakly supervised learning.  _National science review_ , 5(1):44–53.