# A Graph Enhanced BERT Model for Event Prediction

Li Du, Xiao Ding, Yue Zhang, Kai Xiong, Ting Liu, and Bing Qin
Research Center for Social Computing and Information Retrieval
Harbin Institute of Technology, China
{ldu, xding, kxiong<EMAIL_ADDRESS>Corresponding author

###### Abstract

Predicting the subsequent event for an existing event context is an important
but challenging task, as it requires understanding the underlying relationship
between events. Previous methods propose to retrieve relational features from
event graph to enhance the modeling of event correlation. However, the
sparsity of event graph may restrict the acquisition of relevant graph
information, and hence influence the model performance. To address this issue,
we consider automatically building of event graph using a BERT model. To this
end, we incorporate an additional structured variable into BERT to learn to
predict the event connections in the training process. Hence, in the test
process, the connection relationship for unseen events can be predicted by the
structured variable. Results on two event prediction tasks: script event
prediction and story ending prediction, show that our approach can outperform
state-of-the-art baseline methods.

##  1 Introduction

Understanding the semantics of events and their underlying connections is a
long-standing task in natural language processing Minsky (1974); Schank
(1975). Much research has been done on extracting script knowledge from
narrative texts, and making use of such knowledge for predicting a likely
subsequent event given a set of context events.

A key issue to fulfilling such tasks is the modeling of event relation
information. To this end, early work exploited event pair relations Chambers
(2008); Jans et al. (2012); Granroth and Clark (2016) and temporal information
Pichotta (2016); Pichotta and Mooney (2016). The former has been used for
event prediction by using embedding methods, where the similarity between
subsequent events and context events are measured and used for candidate
ranking. The latter has been used for neural network methods, where models
such as LSTMs have been used to model a chain of context events. There has
also been work integrating the two methods Wang et al. (2017).

Figure 1: (a) An example for event prediction. (b) Given an event sequence,
retrieval-based methods lookup structural information of events from event
graph. However, in the test process, part of events may be not covered by the
event graph, hence their connection information is unavailable. Different from
retrieval-based methods, GraphBERT is able to predict the connection strength
between events.

Despite achieving certain effectiveness, the above methods do not fully model
the underlying connection between context events. As shown in Figure 1 (a),
given the facts that _Jason had been overstretched at work_ , _He decided to
change job_ and _Jason finds a new job_ , the subsequent event _Jason is
satisfied with his new job_ is more likely than _Jason feels much stressed at
his new job_ , which can be inferred by understanding the fact that the reason
for his new job search is stress in his job. Li et al. (2018b) and Koncel et
al. (2019) consider such context structure by building event evolutionary
graphs, and using network embedding models to extract relational features. For
these methods, event graphs serve as a source of external structured
knowledge, which are extracted from narrative texts and provide prior features
for event correlation.

One limitation of their methods is that the effectiveness of their methods
heavily relies on the coverage of the event graph. As shown in Figure 1 (b),
Li et al. (2018b) and Koncel et al. (2019)’s methods work by looking up the
event tuples in the event graph to _retrieve_ the connection information
between events for predicting the output. This is done by the standard
knowledge graph lookup operation. However, if the context events are not in
the event graph, the method cannot find relevant information. Figure 1 (b)
shows an extreme case. In event sequence $\beta$, although the context events
_be starving_ and _go for a meal_ are highly similar to the event graph
content _feel hungry_ and _go for lunch_ , the retrieval-based methods can
fail to match context events in the event graph and utilize the event graph
knowledge. However, in practice, it is infeasible to construct an event graph
that covers most of the possible events. As an event is the composition of
multiple arguments, so the same event can correspond to various semantically
equivalent expressions, such as “feel hungry” vs “be starving”, or “hunger”,
etc. This would limit the performance of the retrieval-based systems.

To address this issue, we consider automatically predicting the event links
using a graph-enhanced BERT model (GraphBERT). As shown in Figure 1 (b), we
collect event structure information into a BERT model with graph structure
extension. Given a set of event contexts, we use the GraphBERT model to
construct an event graph structure by predicting connection strengths between
context events, instead of retrieving them from a prebuilt event graph.
Specifically, we extend the BERT model by introducing a structured variable,
which captures the connection strengths between events. As shown in Figure 2,
during training, both context events and external event graph information are
used to train the structured variable. During testing, the structured variable
which describes connection strengths between events is obtained using the
context event only, which is used for finding the next event. Subsequently, we
encode the predicted link strength for making a prediction.

Experimental results on standard datasets show that our model outperforms
baseline methods. Further analysis demonstrates that GraphBERT can predict the
connection strengths for unseen events and improve the prediction accuracy.
The codes are publicly available at https://github.com/sjcfr.

##  2 Background

As shown in Figure 1 (a), the task of event prediction Mostafazadeh et al.
(2016); Li et al. (2018b) can be defined as choosing the most reasonable
subsequent event for an existing event context. Formally, given a candidate
event sequence $X=\\{X_{e_{1}},\dots,X_{e_{t}},X_{e_{c_{j}}}\\}$, where
$\\{X_{e_{1}},\dots,X_{e_{t}}\\}$are $t$ context events and $X_{e_{c_{j}}}$is
the $c_{j}$th candidate subsequent event, the prediction model is required to
predict a relatedness score $Y\in[0,1]$ for the candidate subsequent event
given the event context.

Event graphs Li et al. (2018b) have been used to represent relationships
between multiple events. Formally, an event graph could be denoted as
$G=\\{V,R\\}$, where $V$ is the node set, $R$ is the edge set. Each node
$V_{i}\in V$ corresponds to an event $X_{i}$, while each edge $R_{ij}\in R$
denotes a directed edge $V_{i}\rightarrow V_{j}$ along with a weight $W_{ij}$,
which is calculated by:

$W_{ij}=\frac{\mathrm{count}(V_{i},V_{j})}{\sum_{k}\mathrm{count}(V_{i},V_{k})}$
(1)

where $\mathrm{count}(V_{i},V_{j})$ denotes the frequency of a bigram
$(V_{i},V_{j})$. Hence, the weight $W_{ij}$ is the probability that $X_{j}$ is
the subsequent event of $X_{i}$.

Figure 2: Model Structure. (a) Architecture of the baseline system. Given an
event sequence, the baseline system retrieves event node features and
connection strength from a prebulit event graph. (b) In addition to the
baseline system, GraphBERT introduces an additional aggregator to obtain event
representation from the hidden states of BERT, and learns to predict the
connection strength between events in the training process using the inferer.
So that in the test process, the connection information can be predicted for
arbitrary event.

##  3 Baseline System

Before formally introducing the GraphBERT framework, we first introduce a
retrieval-based baseline system. As Figure 2 (a) shows, given an event
sequence $X=\\{X_{e_{1}},\dots,X_{e_{t}},X_{e_{c_{j}}}\\}$, the baseline
system retrieves the corresponding structural information for each event
within $X$ from a prebuilt event graph $G$, and then integrates the retrieved
structural information into the BERT frame for predicting the relatedness
score $Y$.

For an arbitrary event tuple $(X_{e_{i}},X_{e_{j}})$, if it is covered by the
event graph $G$ (i.e., both $X_{e_{i}}$ and $X_{e_{j}}$ are nodes of $G$),
then we can retrieve the corresponding node embeddings $e_{i}$ and $e_{j}$,
together with the edge weight $A_{ij}$ by matching the event tuple in the
event graph. The representation vector of the events within $X$ further form
into an embedding matrix $E$, and the edge weights form into an adjacency
matrix $A$. To make use of the retrieved structural information for enhancing
the prediction process, we first employ a graph neural network to combine the
event representation matrix and the adjacency matrix:

$\displaystyle E^{(U)}=\sigma(AEW_{U})$ (2)

where $W_{U}\in\mathbb{R}^{d\times d}$ is a weight matrix; $\sigma$ is a
sigmoid function; $E^{(U)}$ is the event representation matrix updated by $A$.

Then the combined event graph knowledge can be merged into the frame of BERT
for enhancing the prediction process. To this end, we employ an attention
operation to softly select relevant information from the updated event
representations $E^{(U)}$, and then update the hidden states of BERT.
Specifically, we take the hidden states of the $s_{1}$th Transformer layer of
BERT (denoted as $H^{s_{1}}$) as the query, and take the updated event
representation $E^{(U)}$ as the key:

${E^{(U)}}^{*}=\mathrm{MultiAttn}(H^{s_{1}},E^{(U)})$ (3)

where ${E^{(U)}}^{*}$ carries information selected from $E^{(U)}$ and relevant
to $H^{s_{1}}$.

Then we merge ${E^{(U)}}^{*}$ with $H^{s_{1}}$ through an addition operation,
and employ layer normalization to keep gradient stability:

${H^{s_{1}}}^{*}=\mathrm{LayerNorm}({E^{(U)}}^{*}+H^{s_{1}})$ (4)

${H^{s_{1}}}^{*}$ contains both the node feature information and the
connection information between events. By taking ${H^{s_{1}}}^{*}$ as the
input of the subsequent $(s_{1}+1)$th Transformer layers of BERT, the event
prediction process is enhanced with the predicted event graph knowledge.

This retrieval-based baseline system can be regarded as the adaption of Li et
al. (2018b) and Koncel et al. (2019)’s retrieval-based methods on a pretrained
model BERT.

##  4 GraphBERT

A critical weakness of the retrieval-based baseline system is that it heavily
relies on the coverage of the event graph. In other words, if an event is not
covered by the event graph, then the structural information (i.e., node
features and the adjacency matrix) would be absent from the constructed event
graph, which further limits the model performance.

In this paper, we propose a predictive-based framework GraphBERT. GraphBERT
uses the transformer layers of BERT as an encoder to obtain the representation
for arbitrary events, and then learns to predict the link strength between
events in the training process, so that the sparsity issues in the retrieval
process can be avoided.

To this end, as Figure 2 (b) shows, in contrast to the retrieval-based
baseline system, we introduce two more modules: (1) An aggregator to obtain
event representations from the BERT framework; (2) an inferer to predict the
link strength between events based on the event representations.

###  4.1 Event Encoding

Given an event sequence $X$, to calculate the event representations and
predict the link strength for events within $X$, GraphBERT first encodes $X$
into a set of _token_ -level distributed representations by taking the
$1$st-$s_{0}$th Transformer layers of BERT as an encoder. Then an aggregator
is employed to aggregate the token level representations into event
representations.

Token Level Representations For an event sequence
$X=\\{X_{1},\cdots,X_{t+1}\\}$, where $X_{i}=\\{x_{1},\dots,x_{l_{i}}\\}$is an
event within $X$ and with $l_{i}$ tokens, the $s_{0}$th Transformer layer of
BERT encodes these tokens into contextualized distributed representations
$H^{s_{0}}=\\{(h_{1}^{1},\dots,h^{1}_{l_{1}}),\cdots,(h_{1}^{t+1},\dots,h_{l_{t+1}}^{t+1})\\}$,
where $h_{j}^{i}\in\mathbb{R}^{1\times d}$ is the distributed representation
of the $j$th token of event $X_{i}$. Then we conduct the graph information
prediction as well as the prediction task based on the token representations.

Event Level Representations An aggregator module aggregates tokens
representation of events derived from the hidden states of BERT (i.e.,
$H^{s_{0}}$) to obtain the event level representations. For an arbitrary event
$X_{i}\in X$, we employ a multi-head attention operation Vaswani et al. (2017)
to aggregate information from the corresponding token representations
$H_{i}^{s_{0}}=(h_{1}^{i},\dots,h_{l_{i}}^{i})$ and obtain the vector
representation of $X_{i}$. Specifically, we define the query matrix of
attention operation as $q_{i}=\frac{1}{l_{i}}\sum h_{l}^{i}$, and take
$H_{i}^{s_{0}}$ as the key matrix as well as the value matrix. Then the
representation of $X_{i}$ is calculated as:

$\hat{e}_{i}=\mathrm{MultiAttn}(q_{i},H_{i}^{s_{0}},H_{i}^{s_{0}})$ (5)

where $\hat{e}_{i}\in\mathbb{R}^{1\times d}$.

In this way, we can obtain the representation of all events within $X$, which
we denote as $\hat{E}=\\{\hat{e}_{1},\cdots,\hat{e}_{t+1}\\}$, where
$\hat{E}\in\mathbb{R}^{(t+1)\times d}$ is a matrix. Note that through the
embedding layer of BERT, position information has been injected into the token
representations. Thus $\hat{E}$ carries event order information.

Then the event representation matrix $\hat{E}$ is used for predicting the link
strength between events. Hence, the performance of link strength prediction
can be strongly influenced by the quality of $\hat{E}$. By deriving $\hat{E}$
from the hidden states of BERT, the abundant language knowledge within BERT
can be utilized to obtain the event representations.

###  4.2 Link Strength Prediction

Given the event representation matrix $\hat{E}$ as node features, we employ an
inferer module to predict the connection strength between arbitrary events
within $X$, regardless of whether these events are seen in the training
process. The output is a matrix $\hat{A}\in\mathbb{R}^{(t+1)\times(t+1)}$ ,
where $\hat{A}_{ij}$ models the probability that event $j$ is the subsequent
event of event $i$.

We stack $n$ graph attention (GAT) layers Veličković et al. (2017) for
consolidating event features. For an event $X_{i}$, the GAT layer works on the
neighborhood of $X_{i}$ to aggregate information. Since the connection between
events are unknown a priori, we set the neighborhood set of event $X_{i}$ as
$\mathcal{N}_{i}=\\{X_{j}\\}$, where $X_{j}\in X,j\neq i$.

Therefore, at the $k$th graph attention layer, given the representation of the
$i$th event $\hat{e}^{k}_{i}$, we calculate the attention coefficients between
other events and derive deep event representation as:

$\begin{split}\alpha_{ij}=\mathrm{softmax}_{j,j\in\mathcal{N}_{i}}&(\mathrm{Relu}(u[W_{\alpha}\hat{e}^{k}_{i}||W_{\alpha}\hat{e}^{k}_{j}]))\\\
\hat{e}^{k+1}_{i}=\sigma&(\sum_{j\in\mathcal{N}_{i}}\alpha_{ij}\mathbf{W}_{\alpha}\hat{e}^{k}_{j})\end{split}$
(6)

where $u\in\mathbb{R}^{1\times 2d},W_{\alpha}\in\mathbb{R}^{d\times d}$are
trainable parameters, $\cdot||\cdot$is a concatenation operation. At the first
GAT layer, $\hat{e}^{1}_{i}$is initialized by $\hat{e}_{i}$derived from the
aggregator.

After $n$ graph attention operations, we employ a bilinear map to calculate a
relation strength score between two events within $X$ based on their deep
representations:

$\Gamma_{ij}=\left(\hat{e}^{n}_{i}\ W_{R}\
\mathrm{T}({\hat{e}^{n}_{j}})\right)$ (7)

where $W_{R}\in\mathbb{R}^{d\times d}$ are learnable parameters, $T(\cdot)$ is
the transpose operation. For all $t+1$ events within $X$, the relation
strength score between arbitrary two events forms a matrix
$\Gamma\in\mathbb{R}^{(t+1)\times(t+1)}$, with each element $\Gamma_{ij}$
measuring the relation strength between $X_{i}$ and $X_{j}$.

Then we normalize the relation strength scores using the softmax function:

$\hat{A}_{ij}=\mathrm{softmax}_{j}(\Gamma_{ij})$ (8)

After the layer normalization, $\sum_{j}\hat{A}_{ij}=1$.

Hence, with the aggregator and the inferer, GraphBERT can obtain
representation and connection strengths for arbitrary events, regardless of
whether or not the event is covered by the event graph. Then the predicted
adjacency matrix $\hat{A}$ and event representations $\hat{E}$ can be used for
prediction, and the process is same as the retrieval-based baseline, as
described in Eq.(2)-Eq.(4).

###  4.3 Training of Inferer

In the training process, we employ a tutor module to supervise the prediction
of $\hat{A}$ using the structural information from a prebuilt event graph.
Given an event sequence $X$, the tutor obtains an adjacency matrix $A$ based
on the edge weights of the event graph. Formally, the weights of $A$ are
initialized as:

$A_{ij}=\left\\{\begin{aligned}
&W_{ij},&\mathrm{if}\,V_{i^{{}^{\prime}}}\rightarrow V_{j^{{}^{\prime}}}\in
R,\\\ &0,&\mathrm{others}.\end{aligned}\right.$ (9)

where $V_{i^{{}^{\prime}}}$, $V_{j^{{}^{\prime}}}$are nodes in the event graph
corresponding to the $i$th and the $j$th event of the candidate event
sequence. The same as the predicted event adjacency matrix $\hat{A}$, $A$ is
also a $\mathbb{R}^{(t+1)\times(t+1)}$matrix.

We scale $A$ to make each row sum equals 1. Therefore, each element of $A$
models the probability that the $j$th event is the subsequent event of the
$i$th event in $X$. In the training process, through minimizing the distance
between $\hat{A}$ and $A$, the inferer module is supervised by the tutor to
learn to predict the event connection strength based on the event
representations.

###  4.4 Optimization

The overall loss function is defined as:

$\displaystyle\begin{split}&L=L_{\text{Event Prediction}}+\lambda
L_{\text{Graph Reconstruction}}\\\ \end{split}$ (10)

where $L_{\text{Event Prediction}}$ is a cross-entropy loss measuring the
difference between predicted relatedness score $Y$ and golden label,
$L_{\text{Graph Reconstruction}}$ assess the difference between $A$ and
$\hat{A}$, $\lambda$ is an additional hyperparameter for balancing the
prediction loss with graph reconstruction loss.

For calculating $L_{\text{Graph Reconstruction}}$, we cast both $A$ and
$\hat{A}$ as a set of random variables, and employ the KL divergence to
measure their difference:

$\displaystyle\begin{split}\vspace{-0.2cm}&L_{\text{Graph Reconstruction}}=\\\
\vspace{-0.2cm}&\sum_{i}\mathrm{KL}(\mathrm{MultiNomial}(\hat{A}_{i})||\mathrm{MultiNomial}(A_{i}))\end{split}$
(11)

where $i$ denotes the $i$th row, and $\mathrm{MultiNomial}(\cdot)$ denotes the
multinomial distribution.

##  5 Experiments

We evaluate our approach on two event prediction tasks: Multiple Choice
Narrative Cloze Task (MCNC) Granroth and Clark (2016) and Story Cloze Test
(SCT) Mostafazadeh et al. (2016) by constructing an event graph based on the
training set of MCNC to train the GraphBERT model and then adapts the
GraphBERT model trained on the MCNC dataset to the SCT dataset to evaluate
whether GraphBERT can predict the link strength between unseen events to
enhance the prediction performance.

###  5.1 Dataset

Multiple Choice Narrative Cloze Task The MCNC task requires the prediction
model to choose the most reasonable subsequent event from five candidate
events given an event context Granroth and Clark (2016). In this task, each
event is abstracted to Predicate-GR form Granroth and Clark (2016), which
represents an event in a structure of {subject, predicate, object,
prepositional object}. Following Granroth and Clark (2016), we extract event
chains from the New York Times portion of the Gigaword corpus. The detailed
statistics of the dataset are shown in Table 1.

Story Cloze Test Task The SCT task requires models to select the correct
ending from two candidates given a story context. Compared with MCNC which
focuses on abstract events, the stories in SCT are concrete events and with
much more details. This dataset contains a five-sentence story training set
with 98,162 instances, and 1,871 four-sentence story contexts along with a
right ending and a wrong ending in the dev. and test dataset, respectively.
Because of the absence of wrong ending in the training set, we only use the
development and the test dataset, and split the development set into 1,771
instances for finetuning models and 100 instances for the development purpose.

###  5.2 Construction of Event Graph

| Training | Dev. | Test
---|---|---|---
#Documents | 830,643 | 103,583 | 103,805
#Event Chains | 140,331 | 10,000 | 10,000
#Unique Events | 430,516 | 44,581 | 47,252
#Uncovered Events | 0 | 24,358 | 24,081
Table 1: Statistics of the MCNC dataset.

The event graph is constructed based on the training set of the MCNC dataset.
Each event within the training set of MCNC is taken as a node of the event
graph, and the edge weights are obtained by calculating the event bigram
frequency. Note that, as shown in Table 1, although the events have been
processed into a highly abstracted form to alleviate the sparsity, there are
still nearly half of the events in the development and test set of MCNC
remains uncovered by the event graph. In the test process, for retrieval-based
methods, given a candidate event sequence with length $t+1$, the edge weights
for events not covered by the event graph are all set as $1/(t+1)$.

###  5.3 Experimental Settings

We implement the GraphBERT model using pretrained BERT-base model, which
contains 12 Transformer layers. We aggregate the _token_ representations from
the 7th Transformer layer of BERT, and merge the updated event representations
to the 10th Transformer layer of BERT. The aggregator has a dimension of 768,
and contains 12 attention heads. The inferer contains 1 GAT layer. The balance
coefficient $\lambda$ equals 0.01. During the training and testing process, we
concatenate the elements of the Predicate-GRs to turn the Predicate-GRs into
strings, so that the event sequences can conform to the input format of the
GraphBERT model. More details are provided in the Appendix.

Baselines for MCNC

Event Pair and Event Chain Based Methods

(i) Event-Comp Granroth and Clark (2016) calculates the pair-wise event
relatedness score using a Siamese network. (ii) PairLSTM Wang et al. (2017)
integrates event order information and pairwise event relations to predict the
ending event. (ii) RoBERTa-RF Lv et al. (2020) enhances pretrained language
model RoBERTa with chain-wise event relation knowledge for making prediction.

Event Graph Based Methods

(i) SGNN Li et al. (2018b) constructs a narrative event evolutionary graph
(NEEG) to describe event connections, and propose a scaled graph neural
network to predict the ending event based on structural information
_retrieved_ from the NEEG. (ii) HeterEvent Zheng et al. (2020) encodes events
using BERT, and implicitly models the word-event relationship by an
heterogeneous graph attention mechanism. (iii) GraphTransformer Koncel et al.
(2019) _retrieves_ structural information from event graph and introduces an
additional graph encoder upon BERT to leverage the structural information.

Pretrained Language Model Based Methods

(i) BERT Devlin et al. (2019) refers to the BERT-base model finetuned on the
MCNC dataset. (ii) GraphBERTλ=0 refers the GraphBERT model optimized with the
balance coefficient $\lambda$ set as 0. Hence, the structural information
cannot be incorporated through the graph reconstruction term.

####  5.3.1 Settings for SCT

To test the generality of GraphBERT, we examine whether GraphBERT can utilize
the structural knowledge learned from MCNC-based event graph to guide the SCT
task. To make fair comparisons, we also trained the BERT Devlin et al. (2019),
GraphTransformer Koncel et al. (2019) on the MCNC dataset, then finetuned them
on the SCT dataset. In the following sections, we use the subscript “MCNC” to
denote the model which has been trained on the MCNC dataset.

However, in the finetuning and test process, GraphTransformer still relies on
an event graph to provide structural information. To address this issue, we
abstract each event in the finetuning set and test set of SCT into the
Predicate-GR form, which is the same form with the nodes in the MCNC-based
event graph. As a result, structural information for an event in SCT can be
retrieved from the MCNC-based event graph using its corresponding Predicate-GR
form, once the event is covered by the event graph.

In addition to the above-mentioned methods, on the SCT dataset, we also
compare GraphBERT with the following event-chain-based baselines:

(i) HCM Chaturvedi et al. (2017) trains a logistic regression model based on
contextual semantic features. (ii) ISCK Chen (2019) integrates narrative
sequence and sentimental evolution information to predict the story ending.

####  5.3.2 Overall Results

We list the results on MCNC and SCT in Table 2 and Table 3, respectively. From
the results on MCNC (Table 2), we can observe that:

(1) Compared to event-pair-based EventComp and event-chain-based PairLSTM,
event-graph-based methods (i.e. SGNN, HeterEvent, GraphTransformer, and
GraphBERT) show better performance. In addition, GraphBERT outperforms event-
chain based RoBERTa-RF, though RoBERTa-RF is built upon a much more powerful
language model. This confirms that involving event structural information
could be effective for this task.

(2) Compared to BERT and GraphBERTλ=0, graph enhanced models GraphTransformer
and GraphBERT further improve the accuracy of script event prediction (T-test;
P-Value $<$ 0.01). This shows that linguistic and structural knowledge can
have a complementary effect.

(3) Compared to the retrieval-based method GraphTransformer, GraphBERT shows
efficiency of learning structural information from the event graph (T-test;
P-Value $<$ 0.01). This indicates that GraphBERT is able to learn the
structural information from the event graph in the training process, and
predict the correct structural information for unseen events in the test
process.

Results on the SCT dataset (Table 3) show that:

(1) Comparing GraphBERT with $\text{BERT}_{\text{MCNC}}$,
$\text{GraphBERT}_{\lambda=0,\text{MCNC}}$ shows that the graph information
can also be helpful for the SCT task.

(2) Though incorporated graph information, the performance of GraphTransformer
is close or inferior to BERT on SCT. This could be because of the limited size
of the SCT development set, which contains 1,771 samples and might be
insufficient to adapt GraphTransformer to the SCT problem. However, GraphBERT
shows a 1.3% absolute improvement over BERT, which indicates the efficiency of
GraphBERT in predicting the link strength between unseen events for predicting
the ending event.

Methods | Accuracy(%)
---|---
Random | 20.00**
EventComp Granroth and Clark (2016) | 49.57**
PairLSTM Wang et al. (2017) | 50.83**
SGNN Li et al. (2018b) | 52.45**
BERT Devlin et al. (2019) | 57.35**
GraphTransformer Koncel et al. (2019) | 58.53**
HeterEvent Zheng et al. (2020) | 58.10**
GraphBERTλ=0 | 57.23**
RoBERTa-RF Lv et al. (2020) | 58.66**
GraphBERT | 60.72
Table 2: Performance of GraphBERT and baseline methods on the test set of MCNC. Accuracy marked with * means p-value $<$ 0.05 and ** indicates p-value $<$ 0.01 in T-test. Methods | Accuracy(%)
---|---
HCM Chaturvedi et al. (2017) | 77.6**
ISCK Chen (2019) | 87.6**
BERT Devlin et al. (2019) | 88.1*
$\text{BERT}_{\text{MCNC}}$ | 88.5*
$\text{GraphTransformer}_{\text{MCNC}}$ Koncel et al. (2019) | 88.9
HeterEvent${}_{\text{MCNC}}$ Zheng et al. (2020) | 88.4*
$\text{GraphBERT}_{\lambda=0,\text{MCNC}}$ | 88.3*
$\text{GraphBERT}_{\text{MCNC}}$ | 89.8
Table 3: Model performance on the test set of SCT. Accuracy marked with *
means p-value $<$ 0.05 and ** indicates p-value $<$ 0.01 in T-test.

###  5.4 Influence of the Accuracy of the Predicted Link Strength

Figure 3: (a) The distribution of Pearson correlation coefficients between the
predicted connection strength and connection strength derived from the event
graph. (b) Relationship between correlation coefficient and model performance.

We investigate the relationship between the accuracy of the predicted link
strengths with the model performance. However, for events in the test set, the
golden event graph is unavailable. To address this issue, we split the
original training set of MCNC into a new training and evaluating set,
containing 120,331 and 20,000 instances, respectively. For each sample, we
calculate the Pearson correlation coefficient between the predicted connection
strengths and connection strengths derived from the event graph, as well as
the relationship between such correlation coefficient and model performance.
The results are shown in Figure 3. We observe that, in general, GraphBERT can
predict the connection between arbitrary events with reasonable accuracy.
Also, the model performance improves as the connection prediction accuracy
increases. This confirms that correctly predicting the event connections for
unseen events can be helpful for the event prediction process.

###  5.5 Influence of the Coverage of the Event Graph

Figure 4: The performance of GraphBERT and GraphTransformer under different
proportion of uncovered events.

We conduct experiments to investigate the specific influence of the sparsity
of the event graph on model performance. Based on the original test set of
MCNC, we build new test sets with different proportions of uncovered events,
and compare the performances of the GraphBERT framework with retrieval-based
method GraphTransformer Koncel et al. (2019) on these test sets. As shown in
Figure 4, as the proportion of uncovered events increase from 0 to 1, the
performance of GraphTransformer shows a negative trend in general. This is
because, for retrieval-based methods, with the increase of sparsity, the
availability of structural information decreases. Compared to
GraphTransformer, the performance of GraphBERT is more stable. These results
indicate that predicting the structural information can be useful for
enhancing the performance of event prediction.

Event Context | Candidate Subsequent Event | Model
---|---|---
A: I heard that my school’s campus had been closed. B: The message said there was a bear on the grounds ! C: The police had to come and help get the bear away. D: They gave the bear a tranquilizer. | $\bm{\mathrm{E_{1}}}$: The bear fell asleep. ($\surd$) | GraphBERT
$\bm{\mathrm{E_{2}}}$: The bear became very violent. ($\times$) | GraphTransformer
Table 4: An example of event predictions made by GraphTransformer and
GraphBERT on the SCT dataset.

###  5.6 Case Study

Table 4 provides an example of prediction results from different models on the
test set of SCT. The event context describes a story that a bear appeared in
the campus and policemen came to tranquilize the bear. Given the event
context, GraphBERT is able to choose correct ending $\mathrm{E_{1}}$ _The bear
fell asleep_ , while GraphTransformer chooses the incorrect ending
$\mathrm{E_{2}}$ _The bear became very violent_.

To correctly predict the story ending, a model should understand the
relationship between _gave a tranquilizer_ and _fell asleep_. However, event
_gave a tranquilizer_ is not covered by the event graph. Hence, the retrieval-
based method GraphTransformer is unable to obtain structural information from
the event graph. On the other hand, in the event graph, there is a directed
edge from a node _obj. sedated_ to node _subj. slept_. This indicates that,
GraphBERT can learn the structural knowledge from the MCNC-based event graph,
and predict the connection between _gave a tranquilizer_ and _fell asleep_ for
instances in the SCT dataset.

##  6 Discussion

The GraphBERT model employs a structure variable $\hat{A}$ to capture the
“is_next_event” relationship between events. By introducing more parallel
structural variables $\\{\hat{A}^{1},\dots,\hat{A}^{k}\\}$, it can be extended
to simultaneously learn multiple kinds of event relationships, such as
temporal or causal relationship. Furthermore, previous researches demonstrate
that the graph-structured relationship extensively exist between other
semantic units, such as sentencesYasunaga et al. (2017), or even paragraphs
Sonawane and Kulkarni (2014). However, similar to the situation in event
graph, it would be impractical to construct knowledge graphs that cover all
possible connection relationships between all the sentences or paragraphs.
This restricts the applicable of retrieval-based methods in these situations.
On the contrary, our generative approach suggests a potential solution by
learning the connection relationship from graph-structured knowledge base with
limited size, then generalizing to the unseen cases.

##  7 Related Work

The investigation of scripts dates back to 1970’s Minsky (1974); Schank
(1975). The script event prediction task models the relationships between
abstract events. Previous studies propose to model the pair-wise relationship
Chambers (2008); Jans et al. (2012); Granroth and Clark (2016) or event order
information Pichotta and Mooney (2016); Pichotta (2016); Wang et al. (2017)
for predicting the subsequent event. Li et al. (2018b) and Lv et al. (2019)
propose to leverage the rich connection between events using graph neural
network and attention mechanism, respectively.

Different from script event prediction, the story cloze task Mostafazadeh et
al. (2016) focuses on concrete events. Therefore, it requires prediction
models to learn commonsense knowledge for understanding the story plot and
predicting the ending. To this end, Li et al. (2018a) and Guan (2019) propose
to combine context clues with external knowledge such as KGs. Li et al. (2019)
finetune pretrained language models to solve the task. Compared to their
works, our approach can use both the language knowledge enriched in BERT to
promote the comprehension of event context, and the structural information
from event graph to enhance the modeling of event connections.

A recent line of work has been engaged in combining the strength of
Transformer based models with graph structured data. To integrate KG with
language representation model BERT, Zhang et al. (2019) encode KG with a graph
embedding algorithm TransE Bordes et al. (2013), and takes the representation
of entities in KG as input of their model. However, this line of work only
linearizes KGs to adapt the input of BERT. Graph structure is not
substantially integrated with BERT. Guan (2019) and Koncel et al. (2019)
propose retrieval-based methods to leverage the structural information of KG.
However, in the event prediction task, the diversity of event expression
challenges the coverage of the event graph, and prevents us from simply
retrieving events in the test instances from the event graph. We propose to
integrate the graph structural information with BERT through a predictive
method. Compared to retrieval-based methods, our approach is able to learn the
structural information of the event graph and _generate_ the structural
information of events to avoid the unavailable of structural information in
test instances.

##  8 Conclusion

We devised a graph knowledge enhanced BERT model for the event prediction
task. In addition to the BERT structure, GraphBERT introduces a structured
variable to learn structural information from the event graph, and model the
relationship between the event context and the candidate subsequent event.
Compared to retrieval-based methods, GraphBERT is able to predict the link
strength between all events, thus avoiding the (inevitable) sparsity of event
graph. Experimental results on MCNC and SCT task show that GraphBERT can
improve the event prediction performances compared to state-of-the-art
baseline methods. In addition, GraphBERT could also be adapted to other graph-
structured data, such as knowledge graphs.

##  9 Acknowledgments

We thank the anonymous reviewers for their constructive comments, and
gratefully acknowledge the support of the New Generation Artificial
Intelligence of China (2020AAA0106501), and the National Natural Science
Foundation of China (62176079, 61976073).

## References

  * Bordes et al. (2013) Antoine Bordes, Nicolas Usunier, Alberto Garcia-Duran, Jason Weston, and Oksana Yakhnenko. 2013.  Translating embeddings for modeling multi-relational data.  _Advances in neural information processing systems_ , 26.
  * Chambers (2008) Chambers. 2008.  Unsupervised learning of narrative event chains.  In _Proceedings of the Association for Computational Linguistics-08: HLT_ , pages 789–797.
  * Chaturvedi et al. (2017) Snigdha Chaturvedi, Haoruo Peng, Dan Roth, and nbd. 2017.  Story comprehension for predicting what happens next.  In _Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing_ , pages 1603–1614.
  * Chen (2019) Chen. 2019.  Incorporating structured commonsense knowledge in story completion.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pages 6244–6251.
  * Clark et al. (2019) Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D Manning. 2019.  What does bert look at? an analysis of bert’s attention.  In _Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pages 276–286.
  * Coenen et al. (2019) Andy Coenen, Emily Reif, Ann Yuan, Been Kim, Adam Pearce, Fernanda Viégas, and Martin Wattenberg. 2019.  Visualizing and measuring the geometry of bert.  In _Proceedings of the 33rd International Conference on Neural Information Processing Systems_ , pages 8594–8603.
  * 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 _Proceedings of the 2019 Annual Meeting of the North American Chapter of the Association for Computational Linguistics_ , pages 4171–4186.
  * Granroth and Clark (2016) Mark Granroth and Stephen Clark. 2016.  What happens next? event prediction using a compositional neural network model.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 30.
  * Guan (2019) Guan. 2019.  Story ending generation with incremental encoding and commonsense knowledge.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pages 6473–6480.
  * Jans et al. (2012) Bram Jans, Steven Bethard, Ivan Vulić, and Marie Francine Moens. 2012.  Skip n-grams and ranking functions for predicting script events.  In _Proceedings of the 13th Conference of the European Chapter of the Association for Computational Linguistics_ , pages 336–344.
  * Jawahar et al. (2019) Ganesh Jawahar, Benoît Sagot, Djamé Seddah, Samuel Unicomb, Gerardo Iñiguez, Márton Karsai, Yannick Léo, Márton Karsai, Carlos Sarraute, Éric Fleury, et al. 2019.  What does bert learn about the structure of language?  In _57th Annual Meeting of the Association for Computational Linguistics (ACL), Florence, Italy_.
  * Koncel et al. (2019) Rik Koncel, Dhanush Bekal, Yi Luan, Mirella Lapata, and Hannaneh Hajishirzi. 2019\.  Text generation from knowledge graphs with graph transformers.  In _Proceedings of the 2019 Annual Meeting of the North American Chapter of the Association for Computational Linguistics_ , pages 2284–2293.
  * Li et al. (2018a) Qian Li, Ziwei Li, Jin-Mao Wei, Yanhui Gu, Adam Jatowt, and Zhenglu Yang. 2018a.  A multi-attention based neural network with external knowledge for story ending predicting task.  In _Proceedings of the 27th International Conference on Computational Linguistics_ , pages 1754–1762.
  * Li et al. (2018b) Zhongyang Li, Xiao Ding, Ting Liu, et al. 2018b.  Constructing narrative event evolutionary graph for script event prediction.  In _International Joint Conference on Artificial Intelligence 2018_ , pages 4201–4207. AAAI Press.
  * Li et al. (2019) Zhongyang Li, Xiao Ding, Ting Liu, et al. 2019.  Story ending prediction by transferable bert.  _arXiv preprint arXiv:1905.07504_.
  * Lv et al. (2019) Shangwen Lv, Wanhui Qian, Longtao Huang, Jizhong Han, and Songlin Hu. 2019.  Sam-net: Integrating event-level and chain-level attentions to predict what happens next.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pages 6802–6809.
  * Lv et al. (2020) Shangwen Lv, Fuqing Zhu, and Songlin Hu. 2020.  Integrating external event knowledge for script learning.  In _Proceedings of the 28th International Conference on Computational Linguistics_ , pages 306–315.
  * Minsky (1974) Marvin Minsky. 1974.  A framework for representing knowledge.
  * Mostafazadeh et al. (2016) Nasrin Mostafazadeh, Nathanael Chambers, Xiaodong He, Devi Parikh, Dhruv Batra, Lucy Vanderwende, Pushmeet Kohli, and James Allen. 2016.  A corpus and cloze evaluation for deeper understanding of commonsense stories.  In _Proceedings of the 2016 Annual Meeting of the North American Chapter of the Association for Computational Linguistics_ , pages 839–849.
  * Pichotta (2016) Pichotta. 2016.  Using sentence-level lstm language models for script inference.  In _Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 279–289.
  * Pichotta and Mooney (2016) Karl Pichotta and Raymond Mooney. 2016.  Statistical script learning with recurrent neural networks.  In _Proceedings of the Workshop on Uphill Battles in Language Processing: Scaling Early Achievements to Robust Methods_ , pages 11–16.
  * Schank (1975) Schank. 1975.  Scripts, plans, and knowledge.  In _Proceedings of the 4th international joint conference on Artificial intelligence-Volume 1_ , pages 151–157.
  * Sonawane and Kulkarni (2014) Sheetal S Sonawane and Parag A Kulkarni. 2014.  Graph based representation and analysis of text document: A survey of techniques.  _International Journal of Computer Applications_ , 96(19).
  * 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.
  * Veličković et al. (2017) Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. 2017.  Graph attention networks.  _arXiv preprint arXiv:1710.10903_.
  * Wang et al. (2017) Zhongqing Wang, Yue Zhang, Ching Yun Chang, and nbd. 2017.  Integrating order information and event relation for script event prediction.  In _Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing_ , pages 57–67.
  * Yasunaga et al. (2017) Michihiro Yasunaga, Rui Zhang, Kshitijh Meelu, Ayush Pareek, Krishnan Srinivasan, and Dragomir Radev. 2017.  Graph-based neural multi-document summarization.  In _Proceedings of the 21st Conference on Computational Natural Language Learning (CoNLL 2017)_ , pages 452–462.
  * Zhang et al. (2019) Zhengyan Zhang, Xu Han, Zhiyuan Liu, Xin Jiang, Maosong Sun, and Qun Liu. 2019.  Ernie: Enhanced language representation with informative entities.  _arXiv preprint arXiv:1905.07129_.
  * Zheng et al. (2020) Jianming Zheng, Fei Cai, Yanxiang Ling, and Honghui Chen. 2020.  Heterogeneous graph neural networks to predict what happen next.  In _Proceedings of the 28th International Conference on Computational Linguistics_ , pages 328–338.

##  10 Experimental Settings

###  10.1 Training Details

To conform to the input format of BERT, for an event described in the
Predicate-GR form {subject, predicate, object, prepositional object}, we first
concatenate each element within the predicate-GR into a string “subject
predicate object prepositional object”, so that an event described in a
structured form is turned into a string. Then for satisfying the requirement
of BERT, the candidate event sequence is further preprocessed into the form
of:

[CLS] $e_{1}$ [SEP] $\dots e_{t}$ [SEP] candidate [SEP] (12)

On the MCNC dataset, the GraphBERT model is trained for 3 epochs, with a batch
size of 64, and a learning rate of 2e-5. While during the finetuning process
on SCT, GraphBERT is optimized with a batch size of 16, and a learning rate of
1e-5, with 5 epochs.

###  10.2 Searching for the Balance Coefficient

Figure 5: The performance of model trained with different balance coefficient
$\lambda$.

In this paper, the objective function is composed of two components. Through
minimizing the graph reconstruction loss, model learns to modeling the bigram
event adjacency patterns. While through minimizing the prediction loss, model
is trained to choose the correct ending given an event context. These two
components are balanced with a coefficient $\lambda$.

To investigate the effect of the balance coefficient, we compare the
prediction accuracy of the GraphBERT model trained with different $\lambda$
and show the results in Figure 5. From which we could observe that, the
prediction accuracy increases as the balance coefficient increase from 0 to
0.1. This is because the additional event graph structure information is
helpful for the event prediction task. However, as the $\lambda$ exceeds 0.5,
the model performances start to decrease. This is because the overemphasis of
graph reconstruction loss would in turn decrease the model performance.

(4, 10) | (5, 10) | (6, 10) | (7, 10) | (8, 10) | (9, 10)
---|---|---|---|---|---
58.76 | 60.28 | 60.57 | 60.72 | 60.28 | 60.01
Table 5: Influence of start layer and merge layer on model performance.

###  10.3 Searching of Start and Merge Layer in BERT

Different transformer layers of BERT tend to concentrate on different semantic
and syntactic information Clark et al. (2019); Coenen et al. (2019).
Therefore, which layer is selected in the BERT to start integrating event
graph knowledge, and which layer is selected to merge graph enhanced event
representations can affect the performance of the model. We study such effect
in two ways: first, we fix the start layer and change the merge layer. Second,
we fix the gap between start and merge layer, and change the start layer.
Results are shown in Table 5. The tuple ($n_{1}$, $n_{2}$) denotes the (start,
merge) layer. From which we could observe that, under the same gap between
merge and start layer, employing the 7th transformer layer of BERT as the
start layer can achieve the best result. While setting the merge–start gap as
2 is more efficient than other choices. Interestingly, Jawahar et al. (2019)
find that the syntactic features can be well captured in the middle layers of
BERT, especially in the 7–9 layer. This indicates that the middle layers of
BERT focus more on sentence level information, and implicitly support the
reasonableness that choosing the 7th and 10th transformer layer of BERT as the
start end merge layer.

##  11 Enhancing Different Kinds of Pretrained Transformer-based Pretrained
Language Models with Event Graph Knowledge

Model | Prediction Accuracy (%)
---|---
BERT | 57.35
GraphBERT | 60.72
RoBERTa | 61.19
GraphRoBERTa | 62.81
Table 6: Performance of the event graph knowledge enhanced RoBERTa model
(Graph-RoBERTa) on the MCNC dataset.

In this paper, we propose the GraphBERT framework, which enhances the
transformer-based pertrained language model BERT with event graph knowledge
through an additional structural variable $\hat{A}$. We argue that, using the
structural variable, we can also equip other transformer-based pretrained
language models, such as RoBERTa, with the event graph knowledge, and then
enhance the event prediction process. This could be achieved by adapt the
aggregator, inferer and merger module upon the other transformer-based
frameworks.

Using the above-mentioned manner, we implemented a GraphRoBERTa model and
examined its performance on the MCNC dataset. The results are shown in Table
6. We observe that, compared with BERT, RoBERTa and GraphRoBERTa show better
performance. This is because, during the pretraining process, RoBERTa can
acquire more abundant linguistic knowledge for understanding the events
through the dynamic masked token prediction mechanism. Moreover, the
comparison between GraphBERT with BERT, and between GraphRoBERTa with RoBERTa
show the effectiveness of our approach in incorporating event graph knowledge
with multiple prevailing transformer-based pretrained language models, to
consistently enhancing the performance of event prediction.


# U-MixFormer: UNet-like Transformer with Mix-Attention
for Efficient Semantic Segmentation

Seul-Ki Yeom, Julian von Klitzing Corresponding author.

###### Abstract

Semantic segmentation has witnessed remarkable advancements with the
adaptation of the Transformer architecture. Parallel to the strides made by
the Transformer, CNN-based U-Net has seen significant progress, especially in
high-resolution medical imaging and remote sensing. This dual success inspired
us to merge the strengths of both, leading to the inception of a U-Net-based
vision transformer decoder tailored for efficient contextual encoding. Here,
we propose a novel transformer decoder, U-MixFormer, built upon the U-Net
structure, designed for efficient semantic segmentation. Our approach
distinguishes itself from the previous transformer methods by leveraging
lateral connections between encoder and decoder stages as feature queries for
the attention modules, apart from the traditional reliance on skip
connections. Moreover, we innovatively mix hierarchical feature maps from
various encoder and decoder stages to form a unified representation for keys
and values, giving rise to our unique _mix-attention_ module.

Our approach demonstrates state-of-the-art performance across various
configurations. Extensive experiments show that U-MixFormer outperforms
SegFormer, FeedFormer, and SegNeXt by a large margin. For example,
U-MixFormer-B0 surpasses SegFormer-B0 and FeedFormer-B0 with 3.8% and 2.0%
higher mIoU and 27.3% and 21.8% less computation and outperforms SegNext with
3.3% higher mIoU with MSCAN-T encoder on ADE20K. Code available at
https://github.com/julian-klitzing/u-mixformer.

## Introduction

Figure 1: Performance vs. computational efficiency on ADE20K (single-scale
inference). U-MixFormer outperforms previous methods in all configurations.

Semantic segmentation, a fundamental downstream task in computer vision, has
consistently received increasing attention in industry and academia. The
importance of semantic segmentation is highlighted by its widespread
applications in real-world scenarios such as autonomous driving (Wang et al.
2022a) and medical diagnosis (Chaitanya et al. 2020). Despite these
advancements, achieving precise pixel-wise predictions remains a challenge due
to the need to balance global and local contexts.

The introduction of the fully convolutional network (FCN)(Long, Shelhamer, and
Darrell 2015) popularized the encoder-decoder structure, in which the encoder
extracts high-level semantics and the decoder combines them with spatial
details. While variants(Badrinarayanan, Kendall, and Cipolla 2017; Chen et al.
2018) have improved this approach, traditional CNNs struggle to capture long-
range context. This limitation has prompted interest in vision transformer-
based methods for segmentation.

Transformer (Vaswani et al. 2017), initially designed for natural language
processing, has been adapted for vision tasks with notable success thanks to
the self-attention mechanism to capture global relationships within input
sequences. Inspired by its success, Dosovitskiy et al. (2020) adapted it for
vision tasks, leading to the Vision Transformer (ViT), which interprets images
as sequences of embedded patches and processes them using the Transformer
encoder. This approach yielded remarkable results on ImageNet. Since the
introduction of the ViT, numerous studies have been performed into its
adaptation for semantic segmentation. The primary objectives have been two-
fold: refining the encoder and crafting decoders that adeptly utilize features
from the encoder stage.

There has been a notable shift towards utilizing the transformer’s decoder
structure in vision tasks. DETR by Carion et al. (2020) pioneered this
approach and integrated the transformer encoder-decoder framework into
detection and segmentation. Following DETR, Segmenter (Strudel et al. 2021),
MaskFormer (Cheng, Schwing, and Kirillov 2021), and Mask2Former (Cheng et al.
2022) introduced decoders for mask prediction with global class labels,
emphasizing high-level features. More recently, FeedFormer (Shim et al. 2023)
proposed a decoder design that decodes high-level encoder features using only
the lowest-level encoder feature. Despite the ongoing progress on transformer-
based decoders for segmentation, these methods often rely on computationally
intensive feature configurations within their attention mechanisms. Moreover,
these methods exhibit inefficiencies in the propagation of feature maps across
the decoder stages.

Traditionally, the U-Net architecture (Ronneberger, Fischer, and Brox 2015),
characterized by its symmetric CNN-based encoder-decoder structure, has been a
favored choice for semantic segmentation, particularly in the medical field.
This favor stems from U-Net’s characteristics of effectively capturing and
propagating hierarchical features. Furthermore, its lateral connections play
an important role, facilitating the flow of multi-stage features between the
encoder and decoder. We hypothesize that leveraging these inherent strengths
of the U-Net architecture can lead to efficient refinement of features, which
can then be hierarchically integrated into the transformer decoder phase.

In this paper, we propose _U-MixFormer_ , a novel UNet-like transformer
decoder. Building upon the foundational principles of U-Net, U-MixFormer
adaptively integrates multi-stage features as keys and values within its
dedicated _mix-attention_ module. This module ensures the gradual propagation
of features and successively remixes them across decoder stages, effectively
managing dependencies between these stages to capture context and refine
boundaries. This can emphasize the hierarchical feature representation like
traditional CNNs and enhance the global context understanding capabilities of
Transformers. To the best of our knowledge, it is the first work to synergize
the inherent strengths of U-Net with the transformative capabilities of Vision
Transformers, particularly through a novel attention module for effectively
harmonizing queries, keys, and values for semantic segmentation.

Our contributions are summarized as follows:

  * •

Novel Decoder Architecture with U-Net We propose a novel powerful transformer-
decoder architecture motivated by the U-Net for efficient semantic
segmentation. Capitalizing on U-Net’s proficiency in capturing and propagating
hierarchical features, our design distinctively uses the lateral connections
of a transformer encoder as query features. This approach ensures a harmonious
fusion of high-level semantics and low-level structures.

  * •

Optimized Feature Synthesis for Enhanced Contextual Understanding To improve
the efficiency of our UNet-like transformer architecture, we mix and update
multiple encoder and decoder outputs as integrated features for keys and
values, resulting in our proposed _mix-attention_ mechanism. This approach has
not only rich feature representation for each decoder stage but also boosts
contextual understanding.

  * •

Compatibility with Diverse Encoders We demonstrate the compatibility of
U-MixFormer combined with the existing popular encoders of both transformer-
based (MiT and LVT) and CNN-based (MSCAN) encoders.

  * •

Empirical Benchmarking As shown in Figure 1, U-MixFormer achieves a new state-
of-the-art in terms of computational cost as well as accuracy among semantic
segmentation methods. It consistently outperforms light-weight, middle-weight,
and even heavy-weight encoders. This superiority is demonstrated for the
ADE20K and Cityscapes datasets, with notable performance on the challenging
Cityscape-C dataset.

Figure 2: U-MixFormer architecture: Encoder (left) extracts multi-resolution
feature maps from input image. U-MixFormer decoder (right) fuses lateral
encoder outputs as $\mathbf{X_{q}}$ with previous decoder stage outputs by
incorporating them into $\mathbf{X_{kv}}$ using our mixed-attention mechanism.
Finally, feature maps from all decoder stages are concatenated, and a MLP
predicts the output.

## Related Work

### Encoder Architectures

SETR (Zheng et al. 2021) was the first architecture to adopt ViT as an encoder
for semantic segmentation. Because ViT only divides the input image into
patches, SETR produces single-scaled encoder features. PVT (Wang et al. 2021)
and Swin Transformer (Liu et al. 2021) repeatedly group feature maps into new
non-overlapping patches between encoder stages, thereby hierarchically
generating multi-scale encoder features. Both methods also enhance the
efficiency of the self-attention module by either reducing the spatial
dimensions of keys and values (PVT) or grouping patches with shifted windows
(Swin Transformer). SegFormer (Xie et al. 2021) reuses PVT’s efficiency
strategy while removing positional encodings and embedding feature maps into
overlapping patches. In contrast to the previously mentioned methods, the
encoders of SegNeXt (Guo et al. 2022) and LVT (Yang et al. 2022) employs
convolutional attention mechanisms.

### Decoder Architectures

DETR (Carion et al. 2020) was the first method to deploy a transformer decoder
for semantic segmentation. Subsequent works (Strudel et al. 2021; Cheng,
Schwing, and Kirillov 2021; Cheng et al. 2022) adapted DETR but also rely on
object-learnable queries, which are computationally expensive, particularly
when combined with multi-scale encoder features. In contrast, FeedFormer (Shim
et al. 2023) directly utilizes the features from the encoder stages as feature
queries, leading to enhanced efficiency. FeedFormer decodes the high-level
encoder features (used as features for queries) with the lowest-level encoder
feature (used as features for keys and values). However, this setup processes
the feature maps individually, without incremental propagation of feature maps
between decoder stages, thereby missing the opportunity for more incremental
refinement to improve object boundary detection. Additionally, other recent
MLP or CNN-based decoders (Xie et al. 2021; Guo et al. 2022) also lack
incremental propagation of decoder features.

### UNet-like Transformer

In both the medical and remote sensing domains, efforts have been made to
transition the UNet architecture from a CNN-based framework to a transformer-
based one. TransUNet (Chen et al. 2021) marked the first successful endeavor
to incorporate the Transformer into medical image segmentation, using ViT in
conjunction with their CNN encoder. Other hybrid approaches are presented in
(Petit et al. 2021; Hatamizadeh et al. 2022; Wang et al. 2022b). Cao et al.
(2023) introduced Swin-UNet, the first fully transformer-based UNet-like
architecture. This design features heavyweight Swin Transformer stages for
both the encoder and decoder, preserving the lateral connections between them
as skip connections. _In contrast to Swin-UNet, our approach employs lighter-
weight decoder stages, rendering it suitable for a wider range of downstream
tasks. Furthermore, we interpret the lateral connections as features for
queries instead of as skip connections and incorporate a unique attention
mechanism._

## Proposed Method

This section introduces U-MixFormer, a novel UNet-like transformer decoder
architecture for semantic segmentation. In general, our decoder is composed of
as many stages $i\in\\{1,...,N\\}$ as there are encoder stages. For clarity,
Figure 2 provides a visual overview of this architecture, exemplified with a
four-stage ($N=4$) hierarchical encoder such as MiT, LVT, or MSCAN.

First, the encoder processes an input image with $H\times W\times 3$. The four
stages $i\in\\{1,...,4\\}$ yield hierarchical, multi-resolution features
$\mathbf{E}_{i}$ with $\frac{H}{2^{i+1}}\times\frac{W}{2^{i+1}}\times C_{i}$.
Second, our decoder stages $i$ sequentially generate refined features
$\mathbf{D}_{4-i+1}$ by performing _mix-attention_ where features for queries
$\mathbf{X}_{q}^{i}$ equal the respective lateral encoder feature map. The
features for keys and values $\mathbf{X}_{kv}^{i}$ are given by a mix of
encoder and decoder stages. Notably, our decoder mirrors the dimensions of the
encoder stage outputs. Third, the decoder features are upsampled using
bilinear interpolation to match the height and width of $\mathbf{D}_{1}$.
Finally, the concatenated features are processed by an MLP to predict the
segmentation map with $H/4\times W/4\times 3$.

### Mix-Attention

Attention modules used in transformer blocks compute the scaled dot-product
attention for queries $\mathbf{Q}$, keys $\mathbf{K}$ and values $\mathbf{V}$
as follows:

$Attention(\mathbf{Q},\mathbf{K},\mathbf{V})=Softmax(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{d_{k}}})\mathbf{V}$
(1)

where $d_{k}$ is the embedding dimension of the keys and $\mathbf{Q}$,
$\mathbf{K}$, $\mathbf{V}$ are obtained from linear projections of selected
features.

Central to our method is the selection of features $\mathbf{X_{kv}}$ that are
to be projected to keys and values, which results in our proposed mix-
attention mechanism. A comparison between traditional self-, cross-, and the
novel _mix-attention_ is illustrated in Figure 3.

Figure 3: Comparison of three different atttention module: Self- (left),
cross- (middle) and mix-attention (right).

In self-attention, the features used to generate queries, keys and values are
identical ($\mathbf{X_{qkv}}$) and originate from the same source, i.e. the
same encoder/decoder stage. Cross-attention employs two distinct features,
$\mathbf{X_{q}}$ and $\mathbf{X_{kv}}$, each derived from a _single_ unique
source. In contrast, our mix-attention mechanism utilizes a mixed feature for
$\mathbf{X_{kv}}$ sourced from _multiple_ multi-scale stages. This concept
allows the query to find matches across all different stages, i.e. degrees of
contextual granularity, thereby facilitating enhanced feature refinement. The
efficacy of this method is validated by our experiments in the Ablation
Studies section.

The selection of a feature set $\mathcal{F}^{i}$ for a decoder stage $i$ is
formalized piecewise as follows:

$\mathcal{F}^{i}=\left\\{\begin{array}[]{cc}\\{\mathbf{E}_{j}\\}_{j=1}^{N}&\text{if
}i=1\\\\[5.0pt]
\\{\mathbf{E}_{j}\\}_{j=1}^{N-i+1}\cup\\{\mathbf{D}_{j}\\}_{j=N-i+2}^{N}&\text{otherwise}\end{array}\right.$
(2)

where for the first decoder stage ($i=1$) all encoder features are selected.
For subsequent stages, previously computed decoder stage outputs are
propagated by replacing their lateral encoder counterparts in
$\mathcal{F}^{i}$.

To align the spatial dimensions of the features in $\mathcal{F}^{i}$, we adapt
the _spatial reduction_ procedure introduced by Wang et al. (2021) :

$\displaystyle\mathbf{\hat{F}}_{j}^{i}$
$\displaystyle=AvgPool(pr_{j},pr_{j})(\mathcal{F}^{i}_{j}),\forall
j\in\\{1,...,N-1\\}$ (3) $\displaystyle\mathbf{\hat{F}}_{j}^{i}$
$\displaystyle=Linear(C_{j},C_{j})(\mathbf{\hat{F}}_{j}^{i}),\forall
j\in\\{1,...,N-1\\}$

where $\mathcal{F}^{i}_{j}$ denotes the $j$-th element of feature set
$\mathcal{F}^{i}$, and $pr_{j}$ is the pooling ratio which aligns the size
with the smallest feature map $\mathcal{F}^{i}_{N}$. The operations AvgPool
and Linear are configured as AvgPool($kernelSize$, $stride$)($\cdot$) and
Linear($C_{in}$, $C_{out}$)($\cdot$), respectively.

The spatially aligned features are concatenated along the channel dimension to
form a mixed feature $\mathbf{X}_{kv}^{i}$ for keys and values.

$\mathbf{X}_{kv}^{i}=Concat(\\{\mathbf{\hat{F}}_{j}^{i}\\}_{j=1}^{N-1}\cup\\{\mathcal{F}_{N}^{i}\\})$
(4)

### Decoder Stage

We adapt the traditional transformer decoder block, by discarding the self-
attention module as recommended by Shim et al. (2023). Additionally, we
replace the cross-attention module with our proposed mix-attention module. The
resulting structure is depicted in Figure 4.

Figure 4: Structure of decoder stage employing our mix-attention mechanism and
taking features for queries $\mathbf{X}_{q}$ and a mixed feature for keys and
values $\mathbf{X}_{kv}$ as input.

Using layer normalization (LN) and a feedforward network (FFN), the output for
a $DecoderStage_{i}$ is computed as follows:

$\displaystyle\mathbf{A}_{i}=LN(MixAtt.(LN(\mathbf{X}_{kv}^{i},\mathbf{X}_{q}^{i}))+LN(\mathbf{X}_{q}^{i}))$
(5) $\displaystyle
DecoderStage_{i}=\mathbf{D}_{N-i+1}=FFN(\mathbf{A}_{i})+\mathbf{A}_{i}$

where $MixAtt.$ denotes our proposed mix-attention.

Table 1: Performance comparison with the state-of-the art methods on ADE20K
and Cityscapes.

|  |  |  | ADE20K | Cityscapes
---|---|---|---|---|---
| Method | Encoder | Params (M) $\downarrow$ |  GFLOPs $\downarrow$ |  mIoU $\uparrow$ |  GFLOPs $\downarrow$ |  mIoU $\uparrow$
| FCN (Long, Shelhamer, and Darrell 2015) | MobileNetV2 | 9.8 | 39.6 | 19.7 | 317.1 | 61.5
| PSPNet (Zhao et al. 2017) | MobileNetV2 | 13.7 | 52.9 | 29.6 | 423.4 | 70.2
| DeepLabV3+ (Chen et al. 2018) | MobileNetV2 | 15.4 | 69.4 | 34.0 | 125.5 | 75.2
| SwiftNetRN (Orsic et al. 2019) | ResNet-18 | 11.8 | - | - | 104.0 | 75.4
| Semantic FPN (Li et al. 2021) | ConvMLP-S | 12.8 | 33.8 | 35.8 | - | -
| SegFormer (Xie et al. 2021) | MiT-B0 | 3.8 | 8.4 | 37.4 | 125.5 | 76.2
| FeedFormer (Shim et al. 2023) | MiT-B0 | 4.5 | 7.8 | 39.2 | 107.4 | 77.9
| FeedFormer (Shim et al. 2023) | LVT | 4.6 | 10.0 | 41.0 | 124.6 | 78.6
| SegNeXt (Guo et al. 2022) | MSCAN-T | 4.3 | 6.6 | 41.1 | 50.5 | 79.8
| U-MixFormer (Ours) | MiT-B0 | 6.1 | 6.1 | 41.2 | 101.7 | 79.0
| U-MixFormer (Ours) | LVT | 6.5 | 9.1 | 43.7 | 122.1 | 79.9
Light-weight | U-MixFormer (Ours) | MSCAN-T | 6.7 | 7.6 | 44.4 | 90.0 | 81.0
| CCNet (Huang et al. 2019) | ResNet-101 | 68.9 | 278.4 | 43.7 | 2224.8 | 79.5
| DeepLabV3+ (Chen et al. 2018) | ResNet-101 | 52.7 | 255.1 | 44.1 | 2032.3 | 80.9
| Auto-DeepLab (Liu et al. 2019) | Auto-DeepLab-L | 44.4 | - | - | 695.0 | 80.3
| OCRNet (Yuan, Chen, and Wang 2020) | HRNet-W48 | 70.5 | 164.8 | 45.6 | 1296.8 | 81.1
| Seg-S-/16 (Strudel et al. 2021) | ViT-S | 22.0 | 31.8 | 45.4 | - | -
| SegFormer (Xie et al. 2021) | MiT-B2 | 27.5 | 62.4 | 46.5 | 717.1 | 81.0
| MaskFormer (Cheng, Schwing, and Kirillov 2021) | Swin-T | 42.0 | 55.0 | 46.7 | - | -
| Mask2Former (Cheng et al. 2022) | Swin-T | 47.0 | 74.0 | 47.7 | - | 82.1
| Mask2Former (Cheng et al. 2022) | ResNet-101 | 63.0 | 90.0 | 47.8 | - | -
| SegFormer (Xie et al. 2021) | MiT-B1 | 13.7 | 15.9 | 42.2 | 243.7 | 78.5
| SegNeXt (Guo et al. 2022) | MSCAN-S | 13.9 | 15.9 | 44.3 | 124.6 | 81.3
| U-MixFormer (Ours) | MiT-B1 | 24.0 | 17.8 | 45.2 | 246.8 | 79.9
| SegFormer (Xie et al. 2021) | MiT-B2 | 27.5 | 62.4 | 46.5 | 717.1 | 81.0
| FeedFormer (Shim et al. 2023) | MiT-B2 | 29.1 | 42.7 | 48.0 | 522.7 | 81.5
| U-MixFormer (Ours) | MiT-B2 | 35.8 | 40.0 | 48.2 | 515.0 | 81.7
Middle-weight | U-MixFormer-S (Ours) | MSCAN-S | 24.3 | 20.8 | 48.4 | 154.0 | 81.8

Table 2: Performance comparison for heavyweight encoders on ADE20K.

Method | Encoder |  Params (M) $\downarrow$ |  GFLOPs $\downarrow$ |  mIoU $\uparrow$
---|---|---|---|---
Seg-B-Mask/16 | ViT-B | 102.4 | 129.4 | 48.7
Seg-L-Mask/16 | ViT-L | 333.2 | 399.9 | 51.8
SETR-PUP | ViT-L | 318.3 | 425.9 | 48.6
SETR-MLA | ViT-L | 310.5 | 368.4 | 48.6
MaskFormer | Swin-S | 63.0 | 79.0 | 49.8
SegFormer | MiT-B3 | 47.3 | 79.0 | 49.4
U-MixFormer (Ours) | MiT-B3 | 55.7 | 56.8 | 49.8
SegFormer | MiT-B4 | 64.1 | 95.7 | 50.3
U-MixFormer (Ours) | MiT-B4 | 72.4 | 73.4 | 50.4
SegFormer | MiT-B5 | 84.7 | 183.3 | 51.0
U-MixFormer+ (Ours) | MiT-B4 | 72.4 | 74.5 | 51.2
U-MixFormer (Ours) | MiT-B5 | 93.0 | 149.5 | 51.9
U-MixFormer+ (Ours) | MiT-B5 | 93.0 | 152.0 | 52.0

### Relationship to UNet Architectures

We introduced U-MixFormer as a UNet-like architecture. However, we want to
underline the main differences between our approach and other UNet-like
variants.

  * •

Because we view the lateral connections as features for queries, our decoder
feature maps implicitly increase in spatial resolution without the need for
explicit upsampling between decoder stages.

  * •

Our approach uses all decoder stages to predict the segmentation map, not only
the final one.

  * •

The feature map of the last decoder stage yields a resolution of $H/4$, $W/4$,
whereas others restore the original spatial resolution $H$, $W$.

## Experiments

### Experimental Settings

#### Datasets

Experiments are conducted on two popular benchmark datasets: ADE20K (Zhou et
al. 2017) and Cityscapes (Cordts et al. 2016). ADE20K is a rigorous scene
parsing benchmark highlighting 150 intricate semantic concepts, split into
20,210 images for training and 2,000 for validation. Cityscapes incorporates
19 densely annotated object categories from urban imagery, aggregating 5,000
images of a high resolution of 2048 $\times$ 1024\. It also introduces 19,998
roughly annotated images for enhanced model training.

Figure 5: Qualitative analysis on ADE20K and Cityscapes datasets: We select
SegFormer and FeedFormer as benchmarks since they share the same encoder. Our
observations indicate that U-MixFormer outperforms these methods, particularly
in segmenting complex object details such as the boundaries between the lake
and floor and across human objects.

#### Implementation Details

To evaluate the versatility of U-MixFormer across various encoder
complexities, we incorporated three distinct encoder backbones: the Mix
Transformer (MiT) (Xie et al. 2021), Light Vision Transformer (LVT) (Yang et
al. 2022), and the multi-scale convolutional attention-based encoder (MSCAN)
(Guo et al. 2022). In detail, we utilized MiT-B0, LVT, and MSCAN-T for our
light-weight variants and MiT-B1/2 and MSCAN-S for the middle-weight
architectures, while the heavier variants included MiT-B3/4/5. The embedding
dimensions were 128 in the final MLP phase for light-weight models and 768 for
others. Section A.1 of the supplementary material provides additional
information about the training and evaluation setting.

Table 3: Ablation study showing the effectiveness of mix-attention and U-Net
like architecture.

Methods | Params (M) | GFLOPs | mIoU
---|---|---|---
Baseline - FeedFormer (Cross-Attention) | 4.5 | 7.8 | 39.2
Mix-Attention | 6.0 | 5.7 | 39.9 (+0.7)
Cross-Attention + U-Net | 5.0 | 6.1 | 40.8 (+0.9)
| Mix-Attention + U-Net
---
(proposed method)
6.1 | 6.1 | 41.2 (+0.4)

Table 4: The average mIoU values for both the clean and corrupted variants of
the Cityscapes validation set were computed for three semantic segmentation
methods, all of which utilize the same encoder (MiT). mIoU is averaged across
all applicable severity levels except for the noise corruption category that
takes into account for the first three out of the five severity levels.

Method | Clean | Blur | Noise | Digital | Weather
---|---|---|---|---|---
Motion | Defoc | Glass | Gauss | Gauss | Impul | Shot | Speck | Bright | Contr | Satur | JPEG | Snow | Spat | Fog | Frost
SegFormer | 76.2 | 59.3 | 59.8 | 48.7 | 60.0 | 26.2 | 27.5 | 31.1 | 52.0 | 73.2 | 66.6 | 72.0 | 38.3 | 21.5 | 53.2 | 67.2 | 31.8
FeedFormer | 77.9 | 59.5 | 60.0 | 50.3 | 59.4 | 21.3 | 21.5 | 25.6 | 47.4 | 74.0 | 66.7 | 73.6 | 39.1 | 22.2 | 52.0 | 65.9 | 32.1
U-MixFormer | 79.0 | 62.4 | 61.7 | 51.7 | 62.1 | 32.8 | 34.4 | 38.4 | 56.7 | 76.2 | 67.2 | 74.8 | 42.7 | 27.5 | 56.1 | 68.0 | 33.9

### Experimental Results

We compare our results with existing semantic segmentation methods on the
ADE20K and Cityscapes datasets. Table 1 showcases our results, including the
number of parameters, Floating Point Operations (FLOPs), and mIoU across both
datasets. As shown in Figure 1, we plot the performance-computation curves of
different methods on the Cityscapes and ADE20K validation set.

#### Light- and Middle-weight Models

In Table 1, the upper section shows the performance of the light-weight
models. As shown in the table on ADE20K, our light-weight U-MixFormer-B0
establishes 41.2% mIoU with 6.1M parameters and 6.1 GFLOPs, outperforming all
other light-weight counterparts in terms of FLOPs and mIoU demonstrating a
better trade-off of performance-computation. Notably, compared to SegFormer
and FeedFormer, which employ the same encoder (MiT-B0), U-MixFormer achieved
mIoU improvement of 3.8% and 2.0% while reducing computation by 27.3% and
21.8%. The performance disparity is even more pronounced on Cityscapes, where
our model achieves 79.0% mIoU with only 101.7 GFLOPs, indicating mIoU
increases of 2.8% and 1.1% and computation reductions of 18.9% and 5.3%
compared to SegFormer-B0 and FeedFormer-B0, respectively. When utilizing LVT,
our model’s performance observes further enhancements with 2.7% and 1.3% of
increased mIoU across datasets. Furthermore, our U-MixFormer with MSCAN-T, a
recent encoder from SegNeXt, also delivers standout results: 44.4% and 81.0%
mIoU on ADE20K and Cityscapes, respectively, using 6.7M parameters.

The latter section of Table 1 shifts the focus to middle-weight models, where
our approach continues to demonstrate superior results, maintaining its edge
over competitors.

#### Heavy-weight Models

As detailed in Table 2, U-MixFormer outperforms SegFormer when paired with the
same heavy encoders, specifically MiT-B3/4/5. For instance, on ADE20K,
U-MixFormer-B3 yields 49.8% mIoU with only 56.8 GFLOPs. This shows a 0.4%
improvement in mIoU and a 28.1% reduction in computation compared to
SegFormer-B3. We additionally hypothesized that enlarging the model size
(changing from MiT-B0 to MiT-B5) would allow for richer contextual information
extraction from the encoder stage, potentially boosting performance. For this
reason, we trained and evaluated heavy-weight model variants, MiT-B4 and
MiT-B5, introducing a method to extract additional keys and values from the
encoder’s 3rd stage midpoint where stack lots of attention blocks. We have
termed this enhanced variant, U-MixFormer+. For MiT-B4 and MiT-B5
configurations, we extracted 5 and 6 keys and values, respectively, to
facilitate mix-attention. As a result, we observed a reasonable performance
improvement of 0.8% for MiT-B4 and 0.1% for MiT-B5 when integrating more
contextual data from the encoder, with only a marginal increase in
computational demands.

Figure 6: Qualitative comparison between MLP, Cross attention, and the
proposed mix-attention approach on ADE20K. We can clearly see the border
segment of objects (Top row: wall/bed (blue) and ramp/box/bed (red), bottom
row: building/background (blue) and lake/floor (red)), which can lead to boost
semantic segmentation performance.

### Qualitative Results

Figure 5 presents the qualitative results of U-MixFormer, FeedFormer, and
SegFormer, all utilizing identical encoders on the ADE20K and Cityscapes
datasets. U-MixFormer excels in segmenting intricate object details and
challenging regions more clearly than other approaches. It can significantly
identify semantically relevant areas and object details, underlying its
ability to learn contextual feature representations from the multi-stage
encoder for efficient segmentation.

Table 5: Performance comparison with MSCAN encoder on ADE20K and Cityscapes.

|  |  | ADE20K | Cityscapes
---|---|---|---|---
Method | Encoder | Params (M) $\downarrow$ | GFLOPs $\downarrow$ | mIoU $\uparrow$ | GFLOPs $\downarrow$ | mIoU $\uparrow$
SegNeXt | MSCAN-T | 4.3 | 6.6 | 41.1 | 50.5 | 79.8
SegNeXt | MSCAN-S | 13.9 | 15.9 | 44.3 | 124.6 | 81.3
U-MixFormer | MSCAN-T | 6.7 | 7.6 | 44.4 | 90.0 | 81.0
U-MixFormer | MSCAN-S | 24.3 | 20.8 | 48.4 | 154.0 | 81.8
SegNeXt | MSCAN-B | 27.6 | 34.9 | 48.5 | 275.7 | 82.6

### Ablation Studies

#### Effectiveness of Mix-Attention and U-Net like Architecture

In Table 3, we systematically evaluate model performance for various design
choices. To guarantee fairness in our comparison, all models are trained and
evaluated under a uniform random seed. We designate FeedFormer, which is based
on conventional cross-attention, as our baseline. Integrating contextual
information from multiple encoder stages via the mix-attention module results
in a 0.7% mIoU improvement while reducing computational costs. Adapting a
U-Net transformer decoder without mix-attention boosts mIoU by 0.9% with a
slight FLOPs increase(+0.4). Remarkably, by applying the mix-attention module
for the U-Net architecture, the model’s performance increases to 41.2%,
signifying a substantial enhancement over the traditional cross-attention in
U-Net like configurations.

#### Robustness in Image Corruption

In safety-critical domains such as autonomous driving and intelligent
transportation systems, the robustness of image segmentation is paramount. In
this respect, we assessed the robustness of U-MixFormer against corruption and
disturbances. Following (Kamann and Rother 2020), we introduced Cityscapes-C,
an augmented version of the Cityscapes val, encompassing 16 algorithmic
corruptions spanning noise, blur, weather, and digital categories. We compare
our U-MixFormer with SegFormer and FeedFormer, both sharing the same encoder.
The findings, shown in Table 4 highlight the superior robustness of
U-MixFormer. Notably, it achieved significant improvements across all
corruption categories, with a remarkable margin of up to 20.0% and 33.3%
against shot noise and 21.8% and 19.2% in snowy conditions. These results
demonstrate the robustness of U-MixFormer, making it an ideal choice for
applications where safety and reliability are crucial.

#### Effectiveness of Mix-Attention

To validate the efficacy of the mix-attention module based on U-Net
microscopically, we conducted an ablation study by extracting feature maps
from identical locations compared to the common study (MLP and Cross-
attention) in Figure 6. We observed a distinct enhancement in our approach.
Specifically, the feature map depicted object details with greater precision
and clearly delineated the boundaries between objects. This observation
demonstrates that the proposed method can significantly segment and capture
distinguishable visual details locally and globally.

#### Effectiveness of Decoding head with the same encoder

We aimed to benchmark our methodology against the state-of-the-art using a
consistent encoder, specifically adopting MiT and MSCAN. These multi-stage
design encoders have received significant attention for their efficiency and
innovative design. As highlighted by results in SegNeXt, both MiT and MSCAN
achieved higher mean Intersection over Union (mIoU) scores and reduced
computational overhead, evident from their fewer FLOPs. This analysis is vital
to highlight the advantages of our method compared to these established
encoders. As depicted in Table 1 and 2, our U-MixFormer consistently
outperforms across various model complexities from B0 to B5 in both mIoU and
FLOPs. Table 5 further highlights U-MixFormer’s superiority, achieving a 3.3%
mIoU increase on ADE20K over the heavier SegNeXt-S. Additionally, MSCAN-S
closely rivals the performance of the more intensive MSCAN encoder
(SegNeXt-B). These findings demonstrate U-MixFormer is a promising decoder
architecture in semantic segmentation.

#### Limitation and Future Works

Despite the competitive results regarding the computational cost and mIoU of
our U-MixFormer, certain limitations need to be addressed. We tested the
inference time of a single 2048 $\times$ 1024 image using a single A100 GPU
under the mmsegmentation benchmark setup. As evident from Table 6, U-MixFormer
tends to have a slower inference time than other light-weight models. The
latency can be attributed to the inherent structure of the U-Net, which
necessitates the preservation of information through lateral (or residual)
connections. While essential for capturing hierarchical features, these
connections introduce overheads during the inference phase. To address this
limitation, we aim to explore model compression techniques such as pruning and
knowledge distillation in our future work. These approaches are anticipated to
potentially improve the inference speed while retaining the accuracy benefits
of U-MixFormer.

Table 6: Comparison of inference times with light-weight models

Method | Params (M) | GFLOPs | mIoU | Inf. time (ms)
---|---|---|---|---
PSPNet | 13.7 | 423.4 | 70.2 | 26.7
DeepLabV3+ | 15.4 | 125.5 | 75.2 | 36.0
SegFormer | 3.8 | 125.5 | 76.2 | 44.8
FeedFormer | 4.5 | 107.4 | 77.9 | 54.4
U-MixFormer | 6.1 | 101.7 | 79.0 | 55.4
U-MixFormer | 6.7 | 90.0 | 81.0 | 68.4

## Conclusion

In this paper, we present U-MixFormer, built upon the U-Net structure designed
for semantic segmentation. U-MixFormer starts with the most contextual encoder
feature map and progressively incorporates finer details, building upon
U-Net’s capability to capture and propagate hierarchical features. Our mix-
attention design emphasizes components of merged feature maps, aligning them
with increasingly granular lateral encoder features. This ensures a harmonious
fusion of high-level contextual information with intricate low-level details,
which is pivotal for precise segmentation. We demonstrate the superiority of
our U-MixFormer across diverse encoders on popular benchmark datasets.

## References

  * Badrinarayanan, Kendall, and Cipolla (2017) Badrinarayanan, V.; Kendall, A.; and Cipolla, R. 2017.  SegNet: A Deep Convolutional Encoder-Decoder Architecture for Image Segmentation.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 39(12): 2481–2495.
  * Cao et al. (2023) Cao, H.; Wang, Y.; Chen, J.; Jiang, D.; Zhang, X.; Tian, Q.; and Wang, M. 2023.  Swin-Unet: Unet-Like Pure Transformer for Medical Image Segmentation.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , 205–218.
  * Carion et al. (2020) Carion, N.; Massa, F.; Synnaeve, G.; Usunier, N.; Kirillov, A.; and Zagoruyko, S. 2020.  End-to-End Object Detection with Transformers.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , 213–229.
  * Chaitanya et al. (2020) Chaitanya, K.; Erdil, E.; Karani, N.; and Konukoglu, E. 2020.  Contrastive learning of global and local features for medical image segmentation with limited annotations.  In _Advances in Neural Information Processing Systems_.
  * Chen et al. (2021) Chen, J.; Lu, Y.; Yu, Q.; Luo, X.; Adeli, E.; Wang, Y.; Lu, L.; Yuille, A. L.; and Zhou, Y. 2021.  TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation.  _arXiv preprint arXiv:2102.04306_.
  * Chen et al. (2018) Chen, L.-C.; Papandreou, G.; Kokkinos, I.; Murphy, K.; and Yuille, A. L. 2018.  DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 40(4): 834–848.
  * Cheng et al. (2022) Cheng, B.; Misra, I.; Schwing, A. G.; Kirillov, A.; and Girdhar, R. 2022.  Masked-Attention Mask Transformer for Universal Image Segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 1290–1299.
  * Cheng, Schwing, and Kirillov (2021) Cheng, B.; Schwing, A.; and Kirillov, A. 2021.  Per-Pixel Classification is Not All You Need for Semantic Segmentation.  In _Advances in Neural Information Processing Systems_ , 17864–17875.
  * Cordts et al. (2016) Cordts, M.; Omran, M.; Ramos, S.; Rehfeld, T.; Enzweiler, M.; Benenson, R.; Franke, U.; Roth, S.; and Schiele, B. 2016.  The Cityscapes Dataset for Semantic Urban Scene Understanding.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 3213–3223.
  * Dosovitskiy et al. (2020) Dosovitskiy, A.; Beyer, L.; Kolesnikov, A.; Weissenborn, D.; Zhai, X.; Unterthiner, T.; Dehghani, M.; Minderer, M.; Heigold, G.; Gelly, S.; Uszkoreit, J.; and Houlsby, N. 2020.  An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.  _arXiv preprint arXiv:2010.11929_.
  * Guo et al. (2022) Guo, M.; Lu, C.; Hou, Q.; Liu, Z.; Cheng, M.; and Hu, S. 2022.  SegNeXt: Rethinking Convolutional Attention Design for Semantic Segmentation.  In _Advances in Neural Information Processing Systems_ , 1140–1156.
  * Hatamizadeh et al. (2022) Hatamizadeh, A.; Tang, Y.; Nath, V.; Yang, D.; Myronenko, A.; Landman, B.; Roth, H. R.; and Xu, D. 2022.  UNETR: Transformers for 3D Medical Image Segmentation.  In _IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)_ , 1748–1758.
  * Huang et al. (2019) Huang, Z.; Wang, X.; Huang, L.; Huang, C.; Wei, Y.; and Liu, W. 2019.  CCNet: Criss-Cross Attention for Semantic Segmentation.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)_ , 603–612.
  * Kamann and Rother (2020) Kamann, C.; and Rother, C. 2020.  Benchmarking the Robustness of Semantic Segmentation Models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 8825–8835.
  * Li et al. (2021) Li, J.; Hassani, A.; Walton, S.; and Shi, H. 2021.  ConvMLP: Hierarchical Convolutional MLPs for Vision.  _arXiv preprint arXiv:2109.04454_.
  * Liu et al. (2019) Liu, C.; Chen, L.; Schroff, F.; Adam, H.; Hua, W.; Yuille, A. L.; and Fei-Fei, L. 2019.  Auto-DeepLab: Hierarchical Neural Architecture Search for Semantic Image Segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 82–92.
  * Liu et al. (2021) Liu, Z.; Lin, Y.; Cao, Y.; Hu, H.; Wei, Y.; Zhang, Z.; Lin, S.; and Guo, B. 2021\.  Swin Transformer: Hierarchical Vision Transformer Using Shifted Windows.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)_ , 10012–10022.
  * Long, Shelhamer, and Darrell (2015) Long, J.; Shelhamer, E.; and Darrell, T. 2015.  Fully Convolutional Networks for Semantic Segmentation.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 3431–3440.
  * Orsic et al. (2019) Orsic, M.; Kreso, I.; Bevandic, P.; and Segvic, S. 2019.  In Defense of Pre-Trained ImageNet Architectures for Real-Time Semantic Segmentation of Road-Driving Images.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 12607–12616.
  * Petit et al. (2021) Petit, O.; Thome, N.; Rambour, C.; Themyr, L.; Collins, T.; and Soler, L. 2021.  U-Net Transformer: Self and Cross Attention for Medical Image Segmentation.  In _Machine Learning in Medical Imaging (MLMI)_ , 267–276.
  * Ronneberger, Fischer, and Brox (2015) Ronneberger, O.; Fischer, P.; and Brox, T. 2015.  U-Net: Convolutional Networks for Biomedical Image Segmentation.  In _Medical Image Computing and Computer-Assisted Intervention (MICCAI)_ , 234–241.
  * Shim et al. (2023) Shim, J.; Yu, H.; Kong, K.; and Kang, S. 2023.  FeedFormer: Revisiting Transformer Decoder for Efficient Semantic Segmentation.  In _AAAI Conference on Artificial Intelligence (AAAI)_ , 2263–2271.
  * Strudel et al. (2021) Strudel, R.; Garcia, R.; Laptev, I.; and Schmid, C. 2021.  Segmenter: Transformer for Semantic Segmentation.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)_ , 7262–7272.
  * Vaswani et al. (2017) Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A. N.; Kaiser, L. u.; and Polosukhin, I. 2017.  Attention is All you Need.  In _Advances in Neural Information Processing Systems_ , 5998–6008.
  * Wang et al. (2022a) Wang, L.; Li, D.; Liu, H.; Peng, J.; Tian, L.; and Shan, Y. 2022a.  Cross-Dataset Collaborative Learning for Semantic Segmentation in Autonomous Driving.  In _AAAI Conference on Artificial Intelligence (AAAI)_ , 2487–2494.
  * Wang et al. (2022b) Wang, L.; Li, R.; Zhang, C.; Fang, S.; Duan, C.; Meng, X.; and Atkinson, P. M. 2022b.  UNetFormer: A UNet-like transformer for efficient semantic segmentation of remote sensing urban scene imagery.  _Journal of Photogrammetry and Remote Sensing_ , 190: 196–214.
  * Wang et al. (2021) Wang, W.; Xie, E.; Li, X.; Fan, D.-P.; Song, K.; Liang, D.; Lu, T.; Luo, P.; and Shao, L. 2021.  Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction Without Convolutions.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)_ , 568–578.
  * Xie et al. (2021) Xie, E.; Wang, W.; Yu, Z.; Anandkumar, A.; Alvarez, J. M.; and Luo, P. 2021.  SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers.  In _Advances in Neural Information Processing Systems_ , 12077–12090.
  * Yang et al. (2022) Yang, C.; Wang, Y.; Zhang, J.; Zhang, H.; Wei, Z.; Lin, Z.; and Yuille, A. 2022\.  Lite Vision Transformer With Enhanced Self-Attention.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 11998–12008.
  * Yuan, Chen, and Wang (2020) Yuan, Y.; Chen, X.; and Wang, J. 2020.  Object-Contextual Representations for Semantic Segmentation.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , volume 12351, 173–190.
  * Zhao et al. (2017) Zhao, H.; Shi, J.; Qi, X.; Wang, X.; and Jia, J. 2017.  Pyramid Scene Parsing Network.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 6230–6239.
  * Zheng et al. (2021) Zheng, S.; Lu, J.; Zhao, H.; Zhu, X.; Luo, Z.; Wang, Y.; Fu, Y.; Feng, J.; Xiang, T.; Torr, P. H.; and Zhang, L. 2021.  Rethinking Semantic Segmentation From a Sequence-to-Sequence Perspective With Transformers.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 6881–6890.
  * Zhou et al. (2017) Zhou, B.; Zhao, H.; Puig, X.; Fidler, S.; Barriuso, A.; and Torralba, A. 2017.  Scene Parsing through ADE20K Dataset.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , 5122–5130.