# On the Effects of Fine-tuning Language Models for Text-Based Reinforcement
Learning

Maurício Gruppi
Villanova University
<EMAIL_ADDRESS>
&Soham Dan
IBM Research
<EMAIL_ADDRESS>
Keerthiram Murugesan
IBM Research
<EMAIL_ADDRESS>
&Subhajit Chaudhury
IBM Research
<EMAIL_ADDRESS>
###### Abstract

Text-based reinforcement learning involves an agent interacting with a
fictional environment using observed text and admissible actions in natural
language to complete a task. Previous works have shown that agents can succeed
in text-based interactive environments even in the complete absence of
semantic understanding or other linguistic capabilities. The success of these
agents in playing such games suggests that semantic understanding may not be
important for the task. This raises an important question about the benefits
of LMs in guiding the agents through the game states. In this work, we show
that rich semantic understanding leads to efficient training of text-based RL
agents. Moreover, we describe the occurrence of semantic degeneration as a
consequence of inappropriate fine-tuning of language models in text-based
reinforcement learning (TBRL). Specifically, we describe the shift in the
semantic representation of words in the LM, as well as how it affects the
performance of the agent in tasks that are semantically similar to the
training games. We believe these results may help develop better strategies to
fine-tune agents in text-based RL scenarios.

On the Effects of Fine-tuning Language Models for Text-Based Reinforcement
Learning

Maurício Gruppi Villanova University<EMAIL_ADDRESS>Soham Dan IBM
Research<EMAIL_ADDRESS>

Keerthiram Murugesan IBM Research<EMAIL_ADDRESS>Subhajit
Chaudhury IBM Research<EMAIL_ADDRESS>

##  1 Introduction

Text-based games (TBGs) are a form of interactive fiction where players use
textual information to manipulate the environment. Since information in these
games is shared as text, a successful player must hold a certain degree of
natural language understanding (NLU). TBGs have surfaced as important testbeds
for studying the linguistic potential of reinforcement learning agents along
with partial observability and action generation. TBGs can be modeled as
partially observable Markov decision processes (POMDP) defined by the tuple
$\langle S,A,O,T,E,R\rangle$, where $S$ is the set of states, $A$ the set of
actions, $O$ the observation space, $T$ the set of state transition
probabilities, $E$ is the conditional observation emission probabilities, and
$R:S\times A\rightarrow\mathbb{R}$ the reward function. The goal of a TBG
agent is to reach the end of the game by interacting with the environment
through text, while maximizing the final score.

Figure 1: Semantic degeneration of the terms kitchen and bloody axe in Zork 1.

In TBGs, observations and actions are presented in the form of unstructured
text, therefore, they must be encoded before being passed onto the RL network.
Recent works in text-based RL adopt a strategy where such encoding is learned
from the game, typically by fine-tuning a language model, such as embeddings
or transformers, using the rewards values from the training Yao et al. (2020);
Wang et al. (2022a). We hypothesize that this approach may cause the language
model to overfit the training games, leading to the degeneration of the
semantic relationships learned during pretraining, and, subsequently,
negatively impacting the agent’s training efficiency and transfer learning
capacity. We conduct experiments in two distinct TBG domains: (1) TextWorld
Commonsense (TWC) (Murugesan et al., 2021a), and (2) Jericho (Hausknecht et
al., 2019). The former provides a number of games where the goal is to perform
house cleaning tasks such as taking objects from a location and placing them
in their appropriate places, using commonsense knowledge. The latter provides
a library of classic text-adventure games, such as the _Zork_ (1977), each
having its own unique objectives, characters, and events. Unlike TWC games,
Jericho games may not let the player know a priori what the final goal is.
Instead, the player is expected to explore the game to learn the story and
complete the tasks one-by-one. In both domains, the agent selects an action
from a given list.

Under this framework, we address the following research questions:

  1. 1.

Does fine-tuning the language model to the RL rewards improve the training
efficiency in comparison to fixed pre-trained LMs?

  2. 2.

Does fine-tuning LMs make agents robust to tasks containing out-of-training
vocabulary?

Our goal is to evaluate what are the implications, pros, and cons of fine-
tuning LMs to the RL tasks. Our results indicate fine-tuning LMs to rewards
leads to a decrease in the agent’s performance and hinders its ability to play
versions of the training games where the observations and actions are slightly
reworded, such as through paraphrasing or lexical substitution (synonyms). In
comparison to fixed pre-trained LMs, these fine-tuned agents under-performed
in training and in test settings. We refer to this process as semantic
degeneration, because it leads to loss of relevant semantic information, in
the LM, that would be crucial to produce generalizable representation. For
instance, by learning that the terms “bloody axe” and “kitchen” are related to
each other in the game Zork 1, the agent overfits to this setting and, in
turn, loses relevant information about “kitchen” and “bloody axe” that could
be important to other games. In NLP, semantic generation might be an expected
consequence of fine-tuning Mosbach et al. (2020), however, the vast majority
of text-based RL agents employ LMs that are fully fine-tuned to the game’s
semantics.

##  2 Background

#### Model and Architecture

The general architecture of the agents in this work consist of a state encoder
akin to the DRRN (He et al., 2015) with an actor-critic policy learning (Wang
et al., 2016) and experience replay. The main components of the agent’s
network are (1) a text encoder, (2) a state-action encoder, and (3) an action
scorer. The text encoder module is a language model that converts an
observation $o\in O$ and action $a\in A$ from text form to fixed length
vectors $f(o)$ and $f(a)$. The state-action encoder consists of a GRU Dey and
Salem (2017) that takes as input the encoded state and actions, and predicts
the Q-values for each pair: $Q_{\phi}(o,a)=g(f(o),f(a))$ given parameters
$\phi$. The action predictor is a linear layer that outputs the probabilities
based on the Q-values from the previous layer. The chosen action is drawn
following the computed probability distribution. The agent is trained by
minimizing the temporal differences (TD) loss:
$\mathcal{L}_{TD}=(r+\gamma\max_{a^{\prime}\in
A}Q_{\phi}(o^{\prime},a^{\prime})-Q_{\phi}(o,a))^{2}$ where $o^{\prime}$ and
$a^{\prime}$ are the next observation and next actions sampled from a replay
memory, $\gamma$ is the reward discount factor.

#### Text Encoders

We used three distinct types of encoders in this study:

  * •

Hash \- does not capture semantic information from the text. It utilizes a
hash function to reduce the observation to a random vector. Similarly to Yao
et al. (2021).

  * •

Word Embedding \- pre-trained static GloVe embeddings Pennington et al. (2014)
and a GRU to encode the sequences of tokens.

  * •

Transformers \- pre-trained LMs to encode observations (Devlin et al., 2018;
Liu et al., 2019).

These encoders are often the top performer Murugesan et al. (2021b);
Ammanabrolu and Hausknecht (2020); Wang et al. (2022b); Atzeni et al. (2021);
Yao et al. (2020); Tuyls et al. (2021) in benchmark environments for text-
based reinforcement learning such as Textworld Côté et al. (2018), Jericho
Hausknecht et al. (2019), Scienceworld Wang et al. (2022b), etc.

##  3 Results

We now present our main results. In the TWC environment, agents are trained
for 100 episodes, with a maximum of 50 steps per episode (repeated over 5
runs). In the Jericho environment, agents were trained over 100000 steps with
no limit to the number of episodes (repeated over 3 runs). These settings were
chosen following previous work reference in this manuscript, such as Yao et
al. (2021) and Murugesan et al. (2021b).

We deploy agents of the same architecture as described in Section 2, the only
exception being that the input encoder used by them is different. The encoders
are the _Hash_ encoder, which produces semantic-less vectors, the Word
_Embedding_ which uses pre-trained GloVe embeddings, and the transformer LMs
Albert Lan et al. (2019) and RoBERTa Liu et al. (2019). The transformer
encoders are used in two variations: Fixed, where the LMs weights are frozen;
and Fine-tuned (FT), where the LMs weights are updated according to the
rewards. This allows us to compare the performance of the typical text-based
RL fine-tuning approach to unconventional ones.

###  3.1 Semantic Information from Pre-Training Improves the Overall RL
Performance

We evaluate the use of different LMs to encode the observations and actions
into fixed-length vectors. We begin our analysis with the weights of the
language model-based encoders fixed, i.e., only the RL network parameters
$\phi$ are updated.

Figure 2: Training performance comparing LM-based encoding models and
hash/word embedding-based models. (left) shows the normalized scores for TWC
games and (right) shows the game score achieved in training across 100k steps
in _Zork 1_. Shaded area corresponds to one standard deviation.

#### The rich semantic information of LMs accelerates training:

The results from these experiments show that even an agent without semantic
information can properly learn to play the games. However, an agent leveraging
the semantic representations from language models are able to: (1) converge
more quickly, in training, to a stable score than hash and simple, as shown in
Figure 2; (2) handle out-of-training vocabulary, Table 1 shows the performance
of the models under two settings: games using an in-training vocabulary (ID)
and games using an out-of-training vocabulary (OOD).

Model | ID | OOD
---|---|---
Hash | $0.58\pm 0.06$ | $0.15\pm 0.03$
Embedding | $0.58\pm 0.08$ | $0.43\pm 0.07$
Albert Lan et al. (2019)* | $0.66\pm 0.05$ | $\textbf{0.65}\pm 0.05$
RoBERTa Liu et al. (2019)* | $\textbf{0.70}\pm 0.05$ | $0.53\pm 0.06$
Table 1: Normalized scores for the in-distribution vocabulary (ID) and out-of-
distribution vocabulary (OOD) game sets in TWC’s Medium difficulty games. (*)
Indicates fixed language models.

These results show that the fixed transformer LMs outperform the Hash and
Embedding models in both vocabulary distributions, highlighting the importance
of keeping the semantic information from pre-training intact.

###  3.2 Semantic Degeneration Hurts Learning

In this experiment, we address the first proposed research question: “does
fine-tuning the LM to the RL rewards improve the training efficiency in
comparison to fixed pre-trained LMs?” To that end, we trained the Fixed and
Fine-tuned variations of Albert and RoBERTa encoders on the same games, and
compared their scores during training. Figure 3 shows the outcome of the
experiment. The findings suggest that traditional text-based RL approach of
fine-tuning the LMs lead to substantially lower training scores, which are due
to semantic degeneration. That is, semantic degeneration leads to ineffective
training of the RL agents, whereas the fixed models converge to a higher score
after a relatively small number episodes/steps.

Figure 3: Training curves of fixed/fine-tuned LMs on (left) TWC medium
difficulty games and (right) Zork 1. Due to semantic degeneration, the fine-
tuned models do not exhibit an increasing score converging to a maximum value.
Shaded areas denote one standard deviation.

Semantic degeneration arises from fine-tuning the LMs to the training rewards.
The LM “forgets” its semantic associations it had learning during its pre-
training, such as the masked token prediction in the case of transformers.
This “forgetting” originates from overfitting the model’s weights to the
games’ word distributions. The biggest problem arises from the fact that the
RL network receives the encoded vectors from the LM and updates its weights
based on such initial representations. However, since the LMs are fine-tuned,
the encoding will change between each episode, causing the RL network to
receive a different encoding for the same observation as the training goes on.

A comparison of pre-trained and semantically degenerated word vectors is seen
in Figure 4. A 2D TSNE plot of pre-trained word vectors from a RoBERTa model
is seen in Figure 4a; Figure 4b shows the plot of the word vectors after fine-
tuning the LM to the game Zork 1. Notice the shift of the term “bloody axe”
towards the term “kitchen” from (a) to (b). The shift happens because both
terms appear in a sequence early on in the game, therefore, the association
between their vectors becomes stronger as the LM is fine-tuned. Moreover, the
terms “egg” and “nest” shift away from “chicken”. The first two terms are also
employed in the game in a sequence where the agent receives a positive reward,
whereas the last is never used in the game. Despite being related in-game,
these terms should have their semantic relationships preserved, which is
possible by utilizing fixed LMs.

Figure 4: Shift caused by the semantic degeneration to the contextual word
vectors in the RoBERTa model fine-tuned to _Zork 1_ : (a) shows the word
embeddings from the pre-trained model, (b) shows the word embeddings after
fine-tuning to _Zork 1_. The bold words denote the case where the term “bloody
axe” shifts towards the word “kitchen” as a result of them co-occurring in a
positively rewarded state.

###  3.3 Agents with fine-tuned LMs are less robust to language change

We address the second research question: “does fine-tuning LMs make agents
robust to tasks containing out-of-training vocabulary?”. To test the
robustness of each model, we first train each agent on a particular game.
Then, we evaluate the agents by having them play games where the observations
are transformed in one of the following ways: Paraphrasing, we run the
observations through a paraphrasing model to rephrase the descriptions (using
a Bart-based paraphrase Lewis et al. (2019)); Lexical Substitution, we replace
words in the observations using synonyms and hypernyms from WordNet Fellbaum
(2010). By playing these versions of the games, agents have to perform the
same task as seen in training, but with reworded or slightly modified
observations.

Figure 5: Evaluation of a RoBERTa agent on original (none), paraphrased, and
lexical substitution observations on (left) TWC medium games and (right) _Zork
1_. In both scenarios, fixed LMs exhibit strong robustness to the
perturbations, scoring as much as in the games without perturbations.

Figure 5a shows the fixed LM agent is robust to paraphrasing as it is able to
maintain the original score even in the modified versions. This is due to the
ability of LMs to handle such perturbations in text. This evidence emphasizes
the hypothesis that semantic understanding is important for generalization to
words unseen in training. Figure 5b shows the performance of the three agents
in _Zork 1_. The fine-tuned agent exhibits a decline in performance while
playing the paraphrased and lexical substitution games. This is explained by
the fact that the LM has been adjusted to the semantics of the original game,
thus, tokens are no longer distributed according to semantic similarity. The
hash-based agent is unable to score in either of the modified games due to the
lack of semantic information. The fixed agent, however, exhibits strong
robustness to the perturbations. This shows how semantic degeneration leads to
decrease in performance in unseen or slightly different games.

##  4 Conclusion

In this paper, we have put forth a novel perspective over the occurrence of
semantic degeneration at the intersection of LM fine-tuning and text RL. We
have shown that semantic understanding brings benefits to the training of
agents. Moreover, despite being the typical approach to text-based RL,
learning the semantics from the game may not be the optimal approach to
training agents. Our results corroborate the well known trends of trading-off
general semantics for task-specific representations in NLP tasks; we shine
light on how this affects agents in carrying out tasks that are semantically
similar to the training ones. Our results indicate that using meaningful
semantic representations can be beneficial, and fine-tuning strategies may be
developed to ensure prior semantic information is not lost by the model, while
learning task-specific representations.

## Limitations

Our work focuses on popular TBG environments and also popular choices of LMs.
In future work it would be interesting to study rarer TBG environments,
potentially beyond English. In that context it would also be interesting to
study multilingual LMs as the semantic representation for these games. Since
we use LM representations for game playing, some of the limitations of these
representations (like inability to distinguish between some related concepts,
or certain biases), might carry over. Investigating these in detail is another
interesting avenue to be explored.

## References

  * Ammanabrolu and Hausknecht (2020) Prithviraj Ammanabrolu and Matthew Hausknecht. 2020.  Graph constrained reinforcement learning for natural language action spaces.  _arXiv preprint arXiv:2001.08837_.
  * Atzeni et al. (2021) Mattia Atzeni, Shehzaad Zuzar Dhuliawala, Keerthiram Murugesan, and Mrinmaya Sachan. 2021.  Case-based reasoning for better generalization in textual reinforcement learning.  In _International Conference on Learning Representations_.
  * Côté et al. (2018) Marc-Alexandre Côté, Akos Kádár, Xingdi Yuan, Ben Kybartas, Tavian Barnes, Emery Fine, James Moore, Matthew Hausknecht, Layla El Asri, Mahmoud Adada, et al. 2018.  Textworld: A learning environment for text-based games.  In _Workshop on Computer Games_ , pages 41–75. Springer.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2018.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_.
  * Dey and Salem (2017) Rahul Dey and Fathi M Salem. 2017.  Gate-variants of gated recurrent unit (gru) neural networks.  In _2017 IEEE 60th international midwest symposium on circuits and systems (MWSCAS)_ , pages 1597–1600. IEEE.
  * Fellbaum (2010) Christiane Fellbaum. 2010.  Wordnet.  In _Theory and applications of ontology: computer applications_ , pages 231–243. Springer.
  * Hausknecht et al. (2019) Matthew Hausknecht, Prithviraj Ammanabrolu, Côté Marc-Alexandre, and Yuan Xingdi. 2019.  Interactive fiction games: A colossal adventure.  _CoRR_ , abs/1909.05398.
  * He et al. (2015) Ji He, Jianshu Chen, Xiaodong He, Jianfeng Gao, Lihong Li, Li Deng, and Mari Ostendorf. 2015.  Deep reinforcement learning with a natural language action space.  _arXiv preprint arXiv:1511.04636_.
  * Lan et al. (2019) Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. 2019.  Albert: A lite bert for self-supervised learning of language representations.  _arXiv preprint arXiv:1909.11942_.
  * Lewis et al. (2019) Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. 2019.  Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  _arXiv preprint arXiv:1910.13461_.
  * Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019.  Roberta: A robustly optimized bert pretraining approach.  _arXiv preprint arXiv:1907.11692_.
  * Mosbach et al. (2020) Marius Mosbach, Maksym Andriushchenko, and Dietrich Klakow. 2020.  On the stability of fine-tuning bert: Misconceptions, explanations, and strong baselines.  _arXiv preprint arXiv:2006.04884_.
  * Murugesan et al. (2021a) Keerthiram Murugesan, Mattia Atzeni, Pavan Kapanipathi, Pushkar Shukla, Sadhana Kumaravel, Gerald Tesauro, Kartik Talamadupula, Mrinmaya Sachan, and Murray Campbell. 2021a.  Text-based rl agents with commonsense knowledge: New challenges, environments and baselines.
  * Murugesan et al. (2021b) Keerthiram Murugesan, Mattia Atzeni, Pavan Kapanipathi, Kartik Talamadupula, Mrinmaya Sachan, and Murray Campbell. 2021b.  Efficient text-based reinforcement learning by jointly leveraging state and commonsense graph representations.  In _Acl-Ijcnlp 2021: The 59Th Annual Meeting Of The Association For Computational Linguistics And The 11Th International Joint Conference On Natural Language Processing, Vol 2_ , pages 719–725. ASSOC COMPUTATIONAL LINGUISTICS-ACL.
  * Pennington et al. (2014) Jeffrey Pennington, Richard Socher, and Christopher D Manning. 2014.  Glove: Global vectors for word representation.  In _Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP)_ , pages 1532–1543.
  * Tuyls et al. (2021) Jens Tuyls, Shunyu Yao, Sham M Kakade, and Karthik R Narasimhan. 2021.  Multi-stage episodic control for strategic exploration in text games.  In _International Conference on Learning Representations_.
  * Wang et al. (2022a) Ruoyao Wang, Peter Jansen, Marc-Alexandre Côté, and Prithviraj Ammanabrolu. 2022a.  Behavior cloned transformers are neurosymbolic reasoners.  _arXiv preprint arXiv:2210.07382_.
  * Wang et al. (2022b) Ruoyao Wang, Peter Jansen, Marc-Alexandre Côté, and Prithviraj Ammanabrolu. 2022b.  Scienceworld: Is your agent smarter than a 5th grader?  In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing_ , pages 11279–11298.
  * Wang et al. (2016) Ziyu Wang, Victor Bapst, Nicolas Heess, Volodymyr Mnih, Remi Munos, Koray Kavukcuoglu, and Nando de Freitas. 2016.  Sample efficient actor-critic with experience replay.  _arXiv preprint arXiv:1611.01224_.
  * Yao et al. (2021) Shunyu Yao, Karthik Narasimhan, and Matthew Hausknecht. 2021.  Reading and acting while blindfolded: The need for semantics in text game agents.  _arXiv preprint arXiv:2103.13552_.
  * Yao et al. (2020) Shunyu Yao, Rohan Rao, Matthew Hausknecht, and Karthik Narasimhan. 2020.  Keep calm and explore: Language models for action generation in text-based games.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 8736–8754.

##  Appendix A Appendix

###  A.1 TextWorld Commonsense

This section contains information about the games (Table 2) in TextWorld
Commonsense as well as an example of an observation and plausible actions
(Figure 6).

The goal of TWC games are to complete a series of household tasks, such as
“picking up an apple and putting it in an appropriate location”. The agent is
provided with the description of a scene and a list of plausible actions. They
must then decide which action to be taken in the current game state. If the
action performed is good, the agent is rewarded with points.

TWC games are split into easy, medium and hard difficulties. As the difficulty
increases, the number of target objects and rooms to cleanup increases.
Details can be seen in Table 2.

| Objects | Targets | Rooms
---|---|---|---
Easy | 1 | 1 | 1
Medium | 2–3 | 1–3 | 1
Hard | 6–7 | 5–7 | 1–2
Table 2: No. of objects, target objects and rooms in TWC games per difficulty
level.

Observation You’ve entered a kitchen. Look over there! A dishwasher. You can
see a closed cutlery drawer. You see a ladderback chair. On the ladderback
chair you can make out a dirty whisk. Plausible Actions Open dishwasher Open
cutlery drawer Take dirty whisk from ladderback chair

Figure 6: Example of an observation from a TextWorld Commonsense game.

###  A.2 Model comparison in TWC

Figure 7 shows the comparison between all language models in all three
difficulties of TWC in terms of normalized score and number of movements.

These results show how agents using fixed LMs converge earlier to a stable
score (Figures 7 a, b, c) and to stable number of movements (Figures 7 d, e,
f). Higher scores are better. Lower number of movements are better because it
means the agent can complete the task while taking fewer actions, avoiding
unnecessary moves.

Easy

Medium

Hard

(a)

(b)

(c)

(d)

(e)

(f)

Figure 7: Comparison of the performance across several language encoding
models. Figures a, b, c show the normalized score for easy, medium and hard
games, respectively. Figures d, e, f show the number of movements needed by
the agent to complete the task (lower values are better). Shaded region
corresponds one standard deviation.

###  A.3 Complete Table of TWC Results

Tables 3 and 4 show the results for all difficulties in TWC in the in-
distribution set and out-of-distribution set.

We can see that fixed LMs consistently perform better when applied to both in-
distribution and out-of-distribution tasks. This is due to the fact that they
can keep rich semantic information and not suffering from semantic
degeneration.

###  A.4 Complete results for perturbation experiments in TWC

Figure 8 shows the results for the perturbation experiments in TWC
difficulties.

The result show how that a fixed LM model (RoBERTa) can maintain a relatively
similar performance to the original observations when playing noisy versions
of the game.

(a) TWC Easy

(b) TWC Medium

(c) TWC Hard

Figure 8: Evaluation of an RoBERTa agent on original, paraphrased, and lexical
substitution observations on (a) Easy, (b) Medium and (c) Hard games.

###  A.5 Comparison of Fine-tuned/Fixed LMs on Jericho games

Figure 9 shows the fine-tuning/fixed LM comparison on additional games from
the Jericho library: detective, pentari, inhumane, and enchanter. The models
show a consistent trend in which the fixed LMs outperform the fine-tuned
models.

(a) inhumane

(b) enchanter

(c) detective

(d) pentari

Figure 9: Comparison of fine-tuned/fixed LMs on various Jericho games.

###  A.6 Text perturbations

This sections presents a description of the perturbations applied to the game
texts.

A perturbation is a modification of an original piece of text in the game to
produce an “out-of-training” example. Perturbations are applied to the
observations, actions and inventories.

The types of perturbations are:

  * •

Lexical substitution - we use WordNet synsets to find replacements for words
in the text

  * •

Paraphrasing - we use a sequence-to-sequence BART paraphraser to rephrase the
original text

##  Appendix B Reproducibility

The code needed used to implement the methods described in this manuscript are
submitted along with the supplementary material. The code is anonymous and
contains the instructions to set up the environments, download the game data,
and train the agents.

| Easy | Medium | Hard
---|---|---|---
Model | Score | Moves | Score | Moves | Score | Moves
DRRN | $0.88\pm 0.04$ | $24\pm 2$ | $0.60\pm 0.02$ | $44\pm 1$ | $0.30\pm 0.02$ | $50\pm 0$
TPC | $0.89\pm 0.06$ | $21\pm 5$ | $0.62\pm 0.03$ | $43\pm 1$ | $0.32\pm 0.04$ | $48\pm 1$
KG-A2C | $0.86\pm 0.06$ | $22\pm 3$ | $0.62\pm 0.03$ | $42\pm 0$ | $0.32\pm 0.00$ | $48\pm 1$
BiKE | $0.94\pm 0.00$ | $18\pm 1$ | $0.64\pm 0.02$ | $39\pm 1$ | $0.34\pm 0.00$ | $47\pm 1$
BiKE + CBR | $0.95\pm 0.04$ | $16\pm 1$ | $0.67\pm 0.03$ | $\textbf{35}\pm 1$ | $\textbf{0.42}\pm 0.04$ | $\textbf{45}\pm 1$
Hash | $0.31\pm 0.07$ | $43\pm 2$ | $0.58\pm 0.06$ | $43\pm 2$ | $0.22\pm 0.03$ | $50\pm 0$
Simple | $0.83\pm 0.08$ | $26\pm 4$ | $0.58\pm 0.08$ | $43\pm 2$ | $0.35\pm 0.05$ | $49\pm 0$
Albert* | $0.96\pm 0.02$ | $10\pm 2$ | $0.66\pm 0.05$ | $38\pm 2$ | $0.41\pm 0.05$ | $49\pm 0$
MPNet* | $0.85\pm 0.04$ | $19\pm 3$ | $0.66\pm 0.06$ | $38\pm 2$ | $0.36\pm 0.04$ | $49\pm 0$
RoBERTa* | $0.94\pm 0.03$ | $12\pm 2$ | $\textbf{0.70}\pm 0.05$ | $38\pm 2$ | $0.40\pm 0.04$ | $49\pm 0$
XLNet* | $1.00\pm 0.00$ | $\textbf{6}\pm 1$ | $0.65\pm 0.08$ | $36\pm 3$ | $0.37\pm 0.07$ | $48\pm 1$
Table 3: Results for the in-distribution (valid) sets in TWC. (*) Indicates agents with fixed LM encoders. | Easy | Medium | Hard
---|---|---|---
Model | Score | Moves | Score | Moves | Score | Moves
DRRN | $0.78\pm 0.02$ | $30\pm 3$ | $0.55\pm 0.01$ | $46\pm 0$ | $0.20\pm 0.02$ | $50\pm 0$
TPC | $0.78\pm 0.07$ | $28\pm 4$ | $0.58\pm 0.01$ | $45\pm 2$ | $0.19\pm 0.03$ | $50\pm 0$
KG-A2C | $0.80\pm 0.07$ | $28\pm 4$ | $0.59\pm 0.01$ | $43\pm 3$ | $0.21\pm 0.00$ | $50\pm 0$
BiKE | $0.83\pm 0.01$ | $26\pm 2$ | $0.61\pm 0.01$ | $41\pm 2$ | $0.23\pm 0.02$ | $50\pm 0$
BiKE + CBR | $0.93\pm 0.03$ | $17\pm 1$ | $0.67\pm 0.03$ | $35\pm 1$ | $\textbf{0.40}\pm 0.03$ | $\textbf{46}\pm 1$
Simple | $0.50\pm 0.12$ | $39\pm 4$ | $0.43\pm 0.07$ | $43\pm 2$ | $0.26\pm 0.04$ | $50\pm 0$
Hash | $0.19\pm 0.06$ | $44\pm 2$ | $0.15\pm 0.03$ | $50\pm 0$ | $0.09\pm 0.02$ | $50\pm 0$
Albert* | $0.64\pm 0.05$ | $33\pm 3$ | $0.65\pm 0.05$ | $\textbf{38}\pm 2$ | $0.16\pm 0.02$ | $50\pm 0$
MPNet* | $0.85\pm 0.05$ | $23\pm 2$ | $0.58\pm 0.06$ | $42\pm 2$ | $0.14\pm 0.02$ | $50\pm 0$
RoBERTa* | $0.90\pm 0.04$ | $19\pm 2$ | $0.53\pm 0.06$ | $44\pm 1$ | $0.19\pm 0.03$ | $50\pm 0$
XLNet* | $0.64\pm 0.05$ | $30\pm 3$ | $0.42\pm 0.07$ | $47\pm 1$ | $0.17\pm 0.03$ | $50\pm 0$
Table 4: Results for the out-of-distribution (test) sets in TWC. (*) Indicates
agents with fixed LM encoders.


11institutetext: Biomedical Image Analysis Group, Imperial College London,
London, UK
22institutetext: NIHR Barts BRC, Queen Mary University of London, London, UK
33institutetext: Data Science Institute, Imperial College London, London, UK
44institutetext: Department of Medicine, Imperial College London, London, UK
44email<EMAIL_ADDRESS>
# Learning shape priors for robust cardiac MR segmentation from multi-view
images

Chen Chen(✉) 11 Carlo Biffi 11 Giacomo Tarroni 11 Steffen Petersen 22
Wenjia Bai 3344 Daniel Rueckert 11

###### Abstract

Cardiac MR image segmentation is essential for the morphological and
functional analysis of the heart. Inspired by how experienced clinicians
assess the cardiac morphology and function across multiple standard views
(i.e. long- and short-axis views), we propose a novel approach which learns
anatomical shape priors across different 2D standard views and leverages these
priors to segment the left ventricular (LV) myocardium from short-axis MR
image stacks. The proposed segmentation method has the advantage of being a 2D
network but at the same time incorporates spatial context from multiple,
complementary views that span a 3D space. Our method achieves accurate and
robust segmentation of the myocardium across different short-axis slices (from
apex to base), outperforming baseline models (e.g. 2D U-Net, 3D U-Net) while
achieving higher data efficiency. Compared to the 2D U-Net, the proposed
method reduces the mean Hausdorff distance ($mm$) from 3.24 to 2.49 on the
apical slices, from 2.34 to 2.09 on the middle slices and from 3.62 to 2.76 on
the basal slices on the test set, when only 10% of the training data was used.

##  1 Introduction

Accurate segmentation of cardiac magnetic resonance (CMR) images is
fundamental for assessing cardiac morphology and diagnosing heart conditions
[1]. Manual segmentation of the anatomical structures is tedious, time-
consuming and prone to subjective errors, which is not suitable for large-
scale studies such as UK Biobank111https://imaging.ukbiobank.ac.uk/ [2].
Therefore, it is essential to develop automated, fast and accurate CMR
segmentation techniques.

Recently, convolutional neural network (CNN) based methods have achieved very
good performance for cardiac image segmentation in terms of both speed and
accuracy [2, 3, 4]. However, they may still produce sub-optimal segmentation
results in some circumstances. For example, in the Automatic Cardiac Diagnosis
Challenge (ACDC) [3], the top segmentation methods (all CNN-based) achieve
high overall segmentation scores for mid-ventricular short-axis slices.
However, they sometimes produce poor results or even fail to locate the
myocardium in basal slices (due to its more complex shape) and apical slices
(due to its small size). This problem is not uncommon and has been reported in
the related literature [5, 3, 6]. Methods based on 2D networks, trained in a
slice-by-slice fashion, are particularly affected by this problem since they
do not incorporate spatial context from neighboring SA images or long-axis
(LA) views. On the other hand, 3D networks are capable of incorporating 3D
spatial information to perform the segmentation task. Yet the 3D spatial
context can be affected by potential inter-slice motion artefacts [7] and the
low through-plane spatial resolution in cardiac SA stacks, thus limiting their
segmentation performance. Compared to 2D ones, 3D networks usually contain
more parameter and are prone to over-fitting especially when the training set
is limited in size since they use 3D volumes rather than 2D slices as input,
significantly reducing the number of training samples.

Experienced clinicians are able to assess the cardiac morphology and function
from multiple standard views, using both SA and LA images to form an
understanding of the cardiac anatomy. Inspired by this, we propose a method
which learns the anatomical prior knowledge across four standard views and
leverages this to perform segmentation on 2D SA images. The intuition behind
our work is that the representation learnt from multiple standard views is
beneficial for the segmentation task on the SA slices as different views
should share the same representation of the 3D anatomy if they are from the
same subject.

The main contributions of this paper are the following: a) we developed a
novel autoencoder architecture (Shape MAE) which learns latent representation
of cardiac shapes from multiple standard views; b) we developed a segmentation
network (multi-view U-Net, adapted from [8]), which is capable of
incorporating the anatomical shape priors learned from multi-view images to
guide the segmentation on SA images; c) we assessed the segmentation accuracy
and the data efficiency of the proposed segmentation method against common 2D
and 3D segmentation baselines by limiting the number of training images,
demonstrating that the proposed method is more robust, and less dependent on
the size of training data.

Related literature. A large number of methods have been developed to improve
the robustness of the cardiac segmentation. One approach is to learn an
ensemble model where the predictions of a 2D and a 3D network are combined
[9]. This method is capable of producing accurate results, but has a
relatively high computational cost and requires an extra post-processing step
to merge the predictions from the two networks. Another approach is to
incorporate cardiac anatomical prior knowledge into segmentation networks [10,
11]. In [10], the learned representation of the 3D cardiac shape is employed
to constrain the segmentation model to predict anatomically plausible shapes.
The main bottleneck of this method is the requirement of fully annotated 3D
high-resolution CMR images which are free from inter-slice motion artefacts
and have high through-plane spatial resolution. However, compared to the
standard 2D imaging protocol, the 3D one requires the subjects to hold their
breath for a relatively long time and therefore is often not feasible for
patients with cardiovascular diseases. Instead of using 3D images, we exploit
_routinely acquired_ 2D standard views to learn the shape representation of
the cardiac structures. The learned representation is then injected into a
segmentation network to improve its performance on SA CMR images. Of note, the
approach in [12] also injects shape priors produced from an autoencoder into a
segmentation network. However, the aim of that approach is to generate
multiple segmentation hypotheses for ambiguous images, and cannot be readily
employed to learn shape priors from different views to enhance cardiac
segmentation.

##  2 Methods

The proposed method consists of two novel architectures: 1) A shape-aware
multi-view autoencoder (Shape MAE) which aims at learning anatomical shape
priors from standard cardiac acquisition planes incl. short-axis and long-axis
views and 2) a multi-view U-Net which performs cardiac short-axis image
segmentation by incorporating anatomical priors learned by Shape MAE into a
modified U-Net architecture.

Figure 1: (a) Overview of Shape MAE. (b) Detailed architectures of each
encoder and each decoder. Each rectangle represents one or a series of
convolutional (Conv) or transposed convolutional (Deconv) layers where the
number in the square box represents the number of filters for each layer. A
‘Res_block’ (pink rectangles) consists of two convolutional layers ($3\times
3$) with a residual connection which adds its input to the features from the
second layer. Instance normalisation and leaky ReLU activation are applied
throughout the network. A sigmoid function is applied to the latent code z to
bound its range.

Shape MAE: Shape-aware multi-view autoencoder. As illustrated in Fig. 1, we
first present a novel architecture named shape-aware multi-view autoencoder
(Shape MAE) which learns anatomical shape priors from standard cardiac views
through multi-task learning. Given a source view $X_{i}$, the network learns
the low-dimensional representation $z_{i}$ of $X_{i}$ that best reconstructs
all the $j$ target views segmentations $Y_{j}$. In this work, we employ four
source views $X_{i}\;(i=1,\dots,4)$ which are three LA images - the two-
chamber view (LA1), three-chamber view (LA2), the four-chamber view (LA3) -
and one mid-ventricular slice (Mid-V) from the SA view. The target
segmentations views $Y_{j}$ ($j=1,\dots,6$) correspond to the four previous
views plus two SA slices: the apical one and the basal one. All encoders
$E_{i}:z_{i}=E_{i}(X_{i})$ and all decoders $D_{j}:Y_{j}=D_{j}(z_{i})$ in the
Shape MAE share the same architecture (see Fig. 1 b).

The loss function $\mathcal{L}_{\text{Shape~{}MAE}}$ for the whole network is
defined as follows:

$\mathcal{L}_{\text{Shape~{}MAE}}=\mathcal{L}_{intra}+\alpha\mathcal{L}_{inter}+\beta\mathcal{L}_{reg}$
(1)

The first two terms of Eq. 1 are defined as the cross entropy loss
$\mathcal{F}$ between the predicted myocardium segmentation
$\hat{Y}_{i\rightarrow j}=D_{j}(E_{i}(X_{i}))$ for the target view $j$ given a
source image $X_{i}$ of the same subject and its ground truth segmentation
$Y_{j}$. $\mathcal{L}_{intra}$ denotes the segmentation loss when the source
view $X_{i}$ and the target view $Y_{j}$ correspond to the same view:
$\mathcal{L}_{intra}=\sum_{i=1,i=j}^{4}\mathcal{F}(Y_{j},\hat{Y}_{i\rightarrow
j})$, whereas the second term $\mathcal{L}_{inter}$ denotes the loss when two
views are different: $\mathcal{L}_{inter}=\sum_{i=1}^{4}\sum_{j=1,i\neq
j}^{6}{\mathcal{F}}(Y_{j},\hat{Y}_{i\rightarrow j})$. The third term is a
regularisation term on the latent representations $z_{i},z_{i}\in Z$:
$\mathcal{L}_{reg}=\frac{1}{|Z|}\sum_{i=1}^{4}{\left|\left|z_{i}-\bar{z}\right|\right|^{2}}$,
which penalises the L2 distance between $z_{i}$ and $\bar{z}$, with
$\bar{z}=\frac{1}{|Z|}\sum_{i=1}^{4}{z_{i}}$ being the average $z$ for a
subject. Although the latent shape codes from different views of the same
subject are not directly shared, this regularisation term forces them to be
close to each other. We use coefficients $\alpha$ and $\beta$ to control the
relative importance of $\mathcal{L}_{inter}$ and $\mathcal{L}_{reg}$.

The principle behind the proposed network is that different views require
_independent_ functions to map them to the latent space that describes global
shape characteristics; whereas translating this latent space to another view
or plane also requires a _specific_ projection function. Predicting the shape
of the myocardium based on the six target views instead of a single view
encourages the network to learn and exploit correlations between different
views, resulting in a global, view-invariant shape representation rather than
a local representation for a particular view. All the encoders and the
decoders in this framework are trained jointly in a multi-task learning
fashion, with the benefit of avoiding over-fitting and encouraging model
generalisation [13].

Figure 2: (a) Overview of the proposed MV U-Net. (b) Architecture of the ‘Fuse
Block’. The number of shown feature map blocks of the U-Net is reduced for
clarity of presentation. Batch normalisation and ReLU activations are applied
throughout the network. For each subject, the shape code of each view is
reshaped to $1\times 8\times 8\times 8$ and then concatenated with the other
three along the second axis to form an input of $1\times 32\times 8\times 8$
to the Fuse Block.

MV U-Net: Multi-view U-Net. As shown in Fig. 2, we propose a segmentation
network called multi-view U-Net (MV U-Net) based on the original U-Net [8] for
cardiac SA image segmentation. The proposed network is capable of
incorporating the anatomical shape priors learned by Shape MAE. Similar to the
original architecture, the proposed architecture comprises 4 down-sampling
blocks and 4 up-sampling blocks to learn multi-scale features. Differently
from the original U-Net, we reduced the number of filters at each level by
four times to account for the fact that cardiac segmentation is simpler than
the lesion segmentation (with multiple candidates) which was the task that the
original U-Net was applied to. In addition, a module called ‘Fuse Block’ is
introduced in the bottleneck of the network (see Fig. 2 b) to inject the
latent codes into the segmentation network. This fusing approach is different
from that in [12] where the latent codes are simply concatenated with U-Net
activations. The proposed module consists of two convolutional (Conv) kernels
($3\times 3$) and a residual connection to combine the shape representations
from different views through learnable weights. Thanks to this module, given
an arbitrary short-axis image slice $I^{p}$ from a subject $p$ and its
correspondent shape representations
${z_{1}^{p},z_{2}^{p},z_{3}^{p},z_{4}^{p}}$ obtained by Shape MAE (one for
each of the four standard views), the network can predict a segmentation
$S^{p}=f_{\text{MV}\operatorname{U-Net}}(I^{p},z_{1}^{p},z_{2}^{p},z_{3}^{p},z_{4}^{p};\theta)$
by distilling the prior knowledge to the high-level features of the network,
allowing it to efficiently refine the segmentations through multi-view
information. The network is trained using standard training procedure with a
cross entropy loss to optimise the parameters $\theta$ of the MV U-Net.

##  3 Experiments and Results

Cardiac multi-view image dataset. Experiments were performed on a dataset
acquired from 734 subjects. For each subject, a stack of 2D SA slices and
three orthogonal 2D LA images are available. All the LV myocardium were
annotated on the SA images as well as the LA images at the end-diastolic (ED)
frame using an automated method followed by manual quality control. All the
images were acquired using one scanner. The spatial resolution of the images
is $1.8\times 1.8\times 10~{}mm$.

In our experiments, the dataset was randomly split into two subsets: a
training set (570 cases), a test set (164 cases). All LA images were
registered to a template subject using rigid transformation with MIRTK
toolkit222https://mirtk.github.io/. All 2D SA slices have been cropped to the
size of $128\times 128$ pixels where the left ventricle is roughly in the
center of every image. Benefiting from the view planning (which is a standard
step during the cardiac image acquisition), we simply use the intersection
point of the three orthogonal LA images on every SA slice to determine its
center of the interest region. All the networks were trained for 200 epochs on
an NVIDIA${}^{\tiny{\text{\textregistered}}}$
GeForce${}^{\tiny{\text{\textregistered}}}$ 2080 Ti, using an Adam optimizer
with a batch size of 10. The learning rate for Shape MAE was set to 0.0001
whereas the learning rate for the segmentation network was set to 0.001. In
our experiments, $\alpha$ was empirically set to 0.5 and $\beta$ to 0.001 in
the $\mathcal{L}_{\text{{Shape~{}MAE}}}$. The proposed algorithm was
implemented in Pytorch.

Segmentation results. To evaluate the segmentation accuracy, we use two
measurements: the Dice score and the Hausdorff distance (HD). The proposed
method is compared against: a 2D U-Net [8], a state-of-the-art 2D FCN for
cardiac MR image segmentation [2], and a 3D U-Net [14]. For fairness and ease
of comparison, all models were set with the same number of filters at each
level (starting with 16 filters in the first layer) and trained with the same
pre-processing and training schedule. For the 3D network, we resampled SA
images to a voxel size of $1.8\times 1.8\times 1.8~{}mm$ and cropped each to a
size of $128\times 128\times 64$ during pre-processing. We trained MV U-Net
and the baseline networks with two settings: in one case we used 10% of the
training set, while in the other one we used 100%. Of note, in each setting,
we first trained a Shape MAE and then trained a MV U-Net where shape priors of
four standard views were obtained using corresponding encoders in the Shape
MAE.

Results on the test set are shown in Table 1. From the table, it can be
observed that the proposed method outperforms the baseline models in both the
low-data setting and the high-data setting, with improved Dice scores at the
apex, middle, and base of the left ventricular myocardium. In particular, when
only 10% data was used, the proposed method reduces the mean HD from 3.24 to
2.49 $mm$ on the apical slices, from 2.34 to 2.09 on the middle slices and
from 3.62 to 2.76 on the basal slices, compared to the 2D U-Net. Fig. 3 shows
examples of the segmentation results from all the networks where the proposed
method not only produces more robust segmentation across slices compared to
the results from the 2D networks, but also achieves more anatomically
plausible results in comparison to the 3D one (see the red arrows in this
figure). Visualization results of the segmentation networks trained in the
high-data setting and Shape MAE are provided in the supplementary material.

Table 1: Comparison of the myocardium segmentation accuracy of the baseline
models and the proposed method in terms of the mean and the standard deviation
of Dice score and HD distance (mm) obtained on the test set (n=164). The
comparison has been carried out separately for apical, mid-ventricular, and
basal slices.

Method | # Training subjects | Dice | HD
---|---|---|---
Apex | Middle | Base | Apex | Middle | Base
2D U-Net | 57 (10%) | 0.898 (0.090) | 0.932 (0.035) | 0.923 (0.077) | 3.239 (6.918) | 2.337 (2.913) | 3.617 (9.058)
2D FCN | 57 (10%) | 0.873 (0.113) | 0.926 (0.041) | 0.919 (0.069) | 3.088 (3.882) | 2.317 (1.440) | 2.948 (2.691)
3D U-Net | 57 (10%) | 0.890 (0.083) | 0.923 (0.043) | 0.923 (0.043) | 2.839 (3.980) | 3.573 (9.05) | 4.469 (10.02)
MV U-Net | 57 (10%) | 0.905 (0.076) | 0.932 (0.025) | 0.926 (0.088) | 2.487 (3.022) | 2.093 (0.577) | 2.758 (3.697)
2D U-Net | 570 (100%) | 0.937 (0.029) | 0.955 (0.016) | 0.948 (0.071) | 1.917 (0.294) | 1.888 (0.178) | 2.327 (2.566)
2D FCN | 570 (100%) | 0.934 (0.032) | 0.958 (0.015) | 0.949 (0.078) | 1.913 (0.297) | 1.890 (0.347) | 2.161 (1.068)
3D U-Net | 570 (100%) | 0.913 (0.112) | 0.945 (0.078) | 0.933 (0.093) | 2.104 (1.24) | 1.957 (0.68) | 2.722 (3.57)
MV U-Net | 570 (100%) | 0.938 (0.027) | 0.958 (0.013) | 0.952 (0.079) | 1.903 (0.345) | 1.874 (0.142) | 2.146 (1.004)

  * •

Approx. # of conv weights (million)

  * •

2D U-Net: 0.8

  * •

2D FCN: 1.0

  * •

3D U-Net: 2.5

  * •

MV U-Net: 1.2

Figure 3: Visualisation of the predicted segmentations and correspondent
ground truth (GT) from the baseline models and MV U-Net (all trained with 10%
training subjects) on an apical, a mid-ventricular and a basal slice from one
patient. Compared to the baseline models, MV U-Net produces more accurate
segmentation with stronger spatial coherence.

##  4 Discussion and Conclusion

In this work, we presented a shape-aware multi-view autoencoder, a neural
network capable of learning anatomical shape priors from multiple standard
views, as well as a multi-view U-Net, a modification of the original U-Net
architecture that incorporates the learned shape priors to improve the
robustness of cardiac segmentation. In contrast to existing works which treat
long-axis CMR segmentation and short-axis CMR segmentation as two separate
tasks [2, 15], our approach, to the best of our knowledge, is the first that
exploits the spatial context from the long-axis images to guide the
segmentation on the short-axis images. The reported experimental results show
that the proposed segmentation method not only demonstrates superior
segmentation accuracy over state-of-the-art 2D baseline methods [2, 8], but
also outperforms a 3D U-Net [14]. This improvement is particularly evident on
the basal and apical slices in the low-data setting, as expected. When
training data is limited, segmenting these challenging slices particularly
benefits from the additional anatomical information extracted from the LA
views and injected into the segmentation network. Of note, our approach does
not require a dedicated acquisition protocol, since LA images are routinely
acquired in most CMR imaging schemes. Moreover, the proposed MV U-Net
maintains the computational advantage of a 2D network, using fewer parameters
($\sim 1.2$ million weights) than the 3D U-Net ($\sim 2.5$ million weights)
during training. This advantage also contributes to the data efficiency of our
method, achieving high segmentation performance with limited training data.
Importantly, our method could be extended in the future to multi-structure
cardiac segmentation. The proposed approach could also be potentially adopted
to other medical image segmentation tasks.

###  4.0.1 Acknowledgements.

This work was supported by the SmartHeart EPSRC Programme Grant
(EP/P001009/1). Steffen Petersen acknowledges support from the National
Institute for Health Research Barts Biomedical Research Centre. The cardiac
multi-view image dataset has been provided under UK Biobank Access Application
18545.

## References

  * [1] Steffen E. Petersen et al.  Reference ranges for cardiac structure and function using cardiovascular magnetic resonance (CMR) in caucasians from the uk biobank population cohort.  JCMR, 19(1):18, 2017.
  * [2] Wenjia Bai et al.  Automated cardiovascular magnetic resonance image analysis with fully convolutional networks.  JCMR, 20(1):65, 2018.
  * [3] Olivier Bernard et al.  Deep Learning Techniques for Automatic MRI Cardiac Multi-structures Segmentation and Diagnosis: Is the Problem Solved?  IEEE TMI, 2018.
  * [4] Qian Tao et al.  Deep learning–based method for fully automatic quantification of left ventricle function from cine MR images: a multivendor, multicenter study.  Radiology, 290(1):81–88, 2018.
  * [5] Q. Zheng et al.  3-D Consistent and Robust Segmentation of Cardiac Images by Deep Learning With Spatial Propagation.  IEEE TMI, 37(9):2137–2148, 2018.
  * [6] Mahendra Khened et al.  Fully convolutional multi-scale residual densenets for cardiac segmentation and automated cardiac diagnosis using ensemble of classifiers.  MedIA, 51:21 – 45, 2019.
  * [7] Giacomo Tarroni et al.  A Comprehensive Approach for Learning-Based Fully-Automated Inter-slice Motion Correction for Short-Axis Cine Cardiac MR Image Stacks.  In MICCAI, pages 268–276. Springer, 2018.
  * [8] Olaf Ronneberger et al.  U-net: Convolutional networks for biomedical image segmentation.  In MICCAI, pages 234–241. Springer, 2015.
  * [9] F. Isensee et al.  Automatic cardiac disease assessment on cine-MRI via time-series segmentation and domain specific features.  Lecture Notes in Computer Science, 10663 LNCS:120–129, 2018.
  * [10] Ozan Oktay et al.  Anatomically constrained neural networks (ACNNs): application to cardiac image enhancement and segmentation.  IEEE TMI, 37(2):384–395, 2018.
  * [11] Jinming Duan et al.  Automatic 3D bi-ventricular segmentation of cardiac images by a shape-refined multi-task deep learning approach.  IEEE TMI, 2019.
  * [12] Simon Kohl et al.  A probabilistic U-Net for segmentation of ambiguous images.  In NeuralIPS, pages 6965–6975, 2018.
  * [13] Rich Caruana.  Multitask learning.  Machine learning, 28(1):41–75, 1997.
  * [14] Özgün Çiçek, Ahmed Abdulkadir, Lienkamp, et al.  3D U-net: Learning dense volumetric segmentation from sparse annotation.  In Lecture Notes in Computer Science, volume 9901 LNCS, pages 424–432, 2016.
  * [15] Davis M. Vigneault et al.  Omega-Net: Fully automatic, multi-view cardiac MR detection, orientation, and segmentation with deep neural networks.  MedIA, 48:95 – 106, 2018.

Supplementary Materials

(a) Trained with 10% data

(b) Trained with 10% data

(c) Trained with 100% data

(d) Trained with 100% data

Figure 4: Exemplar results of the proposed shape-aware multi-view autoencoder
(Shape MAE) trained with 10% and 100% training data, respectively. Given only
one source view (the first column) as input, the proposed Shape MAE is able to
predict the myocardium shapes on the six target views (column 2 to column 7).
This indicates that the proposed approach has the potential to encode the
global shape characteristics of the myocardium in the latent space instead of
a local embedding for a particular view of a subject.

Figure 5: Example results of the proposed segmentation method (MV U-Net) and
the baseline models (all trained with 100% training data) together with the
ground truth (GT) on a stack of short-axis slices. Representative improvements
for cardiac image segmentation can be observed when using the proposed method.
For example, in contrast to the baseline models which produce poor results
when there are unexpected artefacts on the image (see the region inside the
cyan ellipse), the proposed method is able to properly identify the correct
contours.