11institutetext: Artificial Intelligence and Learning Systems Laboratory
School of Electrical and Computer Engineering
National Technical University of Athens
11email<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
# Masked Generative Story Transformer with Character Guidance and Caption
Augmentation

Christos Papadimitriou  Giorgos Filandrianos  Maria Lymperaiou  Giorgos Stamou

###### Abstract

Story Visualization (SV) is a challenging generative vision task, that
requires both visual quality and consistency between different frames in
generated image sequences. Previous approaches either employ some kind of
memory mechanism to maintain context throughout an auto-regressive generation
of the image sequence, or model the generation of the characters and their
background separately, to improve the rendering of characters. On the
contrary, we embrace a completely parallel transformer-based approach,
exclusively relying on Cross-Attention with past and future captions to
achieve consistency. Additionally, we propose a Character Guidance technique
to focus on the generation of characters in an implicit manner, by forming a
combination of text-conditional and character-conditional logits in the logit
space. We also employ a caption-augmentation technique, carried out by a Large
Language Model (LLM), to enhance the robustness of our approach. The
combination of these methods culminates into state-of-the-art (SOTA) results
over various metrics in the most prominent SV benchmark (Pororo-SV), attained
with constraint resources while achieving superior computational complexity
compared to previous arts. The validity of our quantitative results is
supported by a human survey111Code: https://github.com/chrispapa2000/MaskGST .

###### Keywords:

Story Visualization Transformers LLM augmentation

##  1 Introduction

The task of Story Visualization (SV), introduced in 2019 by Li et al. [11],
involves generating a sequence of images, each corresponding to a sentence in
a given textual narrative. This task can be regarded as an extension of Text-
to-Image Generation, incorporating a temporal aspect. Similar to Text-to-
Image, concerns in SV include image quality and text-image relevance. However,
the narrative aspect of SV implies that objects appearing in a visual frame
must maintain a consistent appearance in later frames. The most prominent
examples of such objects are the main characters in the stories.

To address these challenges we construct a framework based on a MaskGIT[3]
model. This approach has recently shown competitive results in Image
Generation, both in terms of diversity and fidelity, whilst being
significantly more efficient compared to auto-regressive Transformers and
Diffusers [3, 2]. We enhance it with Cross-Attention sub-layers, to allow for
past and future captions in the story to serve as context when generating an
image. Cross-Attention mechanisms have been extensively adopted as a way to
inform an input according to some context of different modality or origin[17,
34, 20]. Additionally, we employ a simple LLM-based caption augmentation
technique to improve our model’s robustness and attention on important textual
concepts, similarly to [7] that applies such a method on a contrastive
language-image task. Finally, we introduce a novel Character-Guidance
technique, to prompt our model towards the generation of desired characters.
Specifically, at inference, we form text-conditional logits, as well as logits
conditioned on desired characters and logits conditioned on undesired
characters. We combine those using a mechanism similar to the negative
prompting formula from MUSE[2], to push text-conditional logits, towards the
direction of the desired characters and away from the undesired ones. Our
contributions can be summarized as follows:

  * •

We are the first to employ the promising MaskGIT to construct a MaskGIT-style
Transformer enhanced with Cross-Attention sub-layers for SV.

  * •

We successfully employ an LLM for augmentation of the textual training data in
a completely image-agnostic manner, achieving advanced robustness.

  * •

We propose a simple, yet effective Character Guidance technique that
significantly improves the generation of characters.

  * •

We achieve SOTA results in terms of _Char-F1_ , _Char-Acc_ and _BLEU-2/3_
metrics on Pororo-SV. Additionally, our model has the best _FID_ score,
compared to all previous Transformer-based SV architectures. Our results are
strongly supported by a human survey.

  * •

Our method is significantly more time-efficient than previous SOTAs, developed
under tight resource constraints (16GB of vRAM).

##  2 Related Work

###  2.1 Text-to-Image Generation

The field of Text-to-Image Generation had been previously dominated by GANs
[9, 32, 33]. However, GANs were notoriously unstable during training and
usually specialized in a narrow space of visual themes. More recently,
Diffusion Models [19, 21] have shown exceptional results in the task, by
improving the quality of the generated images, whilst being able to generate a
broader range of themes. However this comes at a cost of larger parameter
counts, as well as increased training/inference times, imposing the need for
high-end hardware and budget. At the same time, there have been several works
that employ Transformers for image generation. At first, auto-regressive
Transformers[6, 18], that infer visual tokens in the latent space of a VQ-
VAE[28] one by one were employed. More recently, Chang _et al_. [3] proposed
MaskGIT, a parallel Transformer for image generation, while MUSE [2] built
upon this idea using a two-stage approach, which achieves results comparable
to those of Diffusion Models, whilst being significantly faster, even at a
similar parameter size. However, the MaskGIT architecture still remains
relatively under-explored, especially in the open-source realm, compared to
Diffusers, that have been widely studied and adopted.

###  2.2 Story Visualization

Initial research in the field revolved around GANs, commencing with
StoryGAN[11]. Several works build on top of this pivotal approach, including
CP-CSV[25], DUCO-StoryGAN[13] and VLC-StoryGAN[12]. In terms of Transformers,
two approaches have been published. VP-CSV [4] employs a two-stage approach,
that starts by predicting the visual tokens in the image regions that
correspond to characters and completes the background of the images at the
second stage. CMOTA [1] uses memory modules to improve consistency and
proposes a bidirectional approach (both text-to-image and image-to-text) to
perform online caption augmentation, during training. Other than that, there
has been a collection of recent diffusion-based approaches[15, 8, 24], all of
which modify and fine-tune pre-trained LDM[19] for the task of SV. Finally,
StoryLDM[17] and StoryGPT-V[19] leverage pre-trained LDM[19], as well.
However, they focus on a modified version of SV where repeated character
references in captions are replaced by pronouns (_e.g_. he, she, they).

##  3 Background

###  3.1 MaskGIT

####  3.1.1 Image Tokenization

MaskGIT [3] works in the latent space of a VQ-GAN [6]. VQ-GAN’s operation on
an image $x\in\mathbb{R}^{3\times H\times H}$ is summarized as follows:

$z=\mathcal{Q}(\mathcal{E}(x))\in\mathbb{R}^{D\times\frac{H}{f}\times\frac{H}{f}},\
\hat{x}=\mathcal{D}(z)\in\mathbb{R}^{3\times H\times H}$ (1)

where $\mathcal{E},\mathcal{Q},\mathcal{D}$ are the Encoder, Quantizer and
Decoder, respectively. The latent representation $z$ is discrete, meaning that
each $D$-dimensional vector in the ($\frac{H}{f}\times\frac{H}{f}$)-sized
output of the Encoder is substituted via nearest-neighbor lookup, using a
library of $K$ visual tokens (this operation is carried out by $\mathcal{Q}$).
$H$ is the dimension of the images and $f$ is a compression factor.

####  3.1.2 Text-to-Image Generation

A Transformer is trained to predict the visual tokens of an image by
conditioning on text tokens. The training and inference techniques proposed in
[3] are summarized below.

##### Training

An image is encoded into its discrete, latent representation $z$ using VQ-
GAN’s Encoder and Quantizer. $Y=[y_{i}]_{i=1}^{N}$, with
$N=(\frac{H}{f})^{2}$, is formed by flattening the visual tokens, $z$ into a
vector. $M=[m_{i}]_{i=1}^{N}$ is a random binary mask for all tokens. At each
training step, the tokens in positions ($i$), where $m_{i}=1$ are masked out
(replaced by a special token, [MASK] ). The Mask $M$ is applied over $Y$ to
obtain $Y_{\bar{M}}$. MaskGIT’s training objective is:

$\mathcal{L}_{mask}=-\mathbb{E}[\sum_{\forall i\in[1,N],m_{i}=1}\log
p_{\theta}(y_{i}|Y_{\bar{M}})]$ (2)

where $\theta$ represents the parameters of the Transformer. Predicting
randomly masked tokens during training allows for the utilization of an
efficient inference technique (outlined in the next paragraph), so that
multiple visual tokens to be predicted at each step, in contrast with
traditional Next Token Prediction.

##### Inference

Inference takes place over $T$ iterations. It starts with all visual tokens
masked out; $Y_{\bar{M}}^{(0)}=[\texttt{[MASK]}]_{i=1}^{N}$. In the $t$-th
iteration, the masked tokens ($Y_{\bar{M}}^{(t)}$) are passed through the
Transformer to predict probabilities $p^{(t)}\in\mathbb{R}^{N\times K}$ for
all masked tokens. In every position, a token is sampled:

$y_{i}^{(t)}\sim p_{i}^{(t)},\ p_{i}^{(t)}\in\mathbb{R}^{K},\ \forall
i\in[1,N]$ (3)

The probabilities according to which the tokens were sampled are now dealt
with as confidence scores. The most confident tokens are kept unmasked and the
rest of them are re-masked, to be predicted in a future iteration.

###  3.2 Caption Augmentation using LLMs

The outstanding capabilities of LLMs have been previously leveraged to perform
text augmentation in the context of various tasks [27, 30, 5, 31, 7]. [7]
proposes a method for augmenting captions of text-image pairs that are used to
train a CLIP[16] model. At first, alternative captions are generated for a
small number of text-image pairs, through various methods, including human
annotation and chatbots. Original and generated captions are paired to form
meta-input-output pairs. Subsequently, LLaMA[26] is used to produce
alternative captions for all samples in the training data. The meta-input-
output pairs are used as context for the LLM to better understand the task.

##  4 Method

###  4.1 Preliminaries

Let $S=\\{s_{1},...,s_{n}\\}$ and $X=\\{x_{1},...,x_{n}\\}$ be the story
captions and story images respectively, where $n$ is the number of caption-
image pairs in a story. Additionally, we denote
$z=\mathcal{Q}(\mathcal{E}(x))$ the latent, discrete representation of image
$x$, encoded by VQ-GAN.

###  4.2 MaskGST

We propose MaskGST (Masked Generative Story Transformer), based on a modified
version of MaskGIT[3] for SV. Following MaskGIT, we adopt VQ-GAN [6] for image
tokenization. VQ-GAN’s function is briefly discussed in Sec. 3.1.1.

####  4.2.1 Transformer

MaskGST’s Transformer (Fig. 1) employs two kinds of layers. _Full-Layers_ have
a Self-Attention sub-layer, followed by a Cross-Attention sub-layer, followed
by a Feed-Forward Network (as in the Decoder of the original Transformer[29]).
_Self-Layers_ have the same structure, except that they omit Cross-Attention
(as in the Encoder of the original Transformer[29]). The Transformer comprises
of two Full-Layers, followed by four Self-Layers.

##### Input

For the image-caption pair $(x_{i},s_{i})$ the Transformer’s input $I_{i}$ is
composed as follows:

$I_{i}=(z_{i}^{FLAT};s_{i}^{BPE})\in\mathbb{R}^{(\frac{H}{f}\cdot\frac{H}{f}+L)\times
d}$ (4)

where $z_{i}^{FLAT}\in\mathbb{R}^{(\frac{H}{f}\cdot\frac{H}{f})\times d}$ is
the flattened version of $z_{i}=\mathcal{Q}(\mathcal{E}(x_{i}))$, projected to
the Transformer’s hidden space. $s_{i}^{BPE}\in\mathbb{R}^{L\times d}$ are the
text embeddings that correspond to the BPE-encoding [22] of text caption
$s_{i}$, with $L$ being the number of text tokens in the representation.
Finally, $d$ is the hidden dimension of the Transformer.

##### Context

The context utilized in the Cross-Attention sub-layers, for the $i$-th image
of a story is:

$\mathit{ctx}_{i}=(s_{1}^{BPE};s_{2}^{BPE};...;s_{n}^{BPE})\in\mathbb{R}^{(n\cdot
L)\times d}$ (5)

That is, when generating the $i$-th image, the model performs Cross-Attention
to all captions in the story, both past and future. This way, in the first two
Full-Layers of the generative process, the model is able to integrate relevant
information from previous and consecutive captions into the visual and textual
tokens of the input sequence. Then, in the following Self-Layers, the visual
tokens of the input sequence are forged through mutual self-attention, as well
as self-attention applied between them and the textual tokens of the input
sequence, to form the final output sequence.

##### Training

For Training we adopt MaskGIT’s Masked Visual Token Modeling algorithm [3]
(Sec. 3.1.2). For each training sample, a random subset of visual tokens is
masked out in the input according to a binary mask $M$. The Transformer’s
final hidden states are projected to the $K$-dimensional space to form:

$O_{i}=(O_{i}^{vis};O_{i}^{text})\in\mathbb{R}^{(\frac{H}{f}\cdot\frac{H}{f}+L)\times
K},\text{where }O_{i}^{vis}\in\mathbb{R}^{(\frac{H}{f}\cdot\frac{H}{f})\times
K}$ (6)

The cross-entropy between the masked positions in $O_{i}^{vis}$ and the
corresponding ground-truth tokens in the VQ-GAN encoding $z$ is calculated.

##### Inference

As in training, we adopt MaskGIT’s iterative inference algorithm.

Figure 1: MaskGST’s Transformer model

###  4.3 Caption Augmentation Using an LLM

We use ChatGPT to augment the training captions, similarly to [7]. Instead of
providing pairs of original and generated captions as context, we provide
ChatGPT a description of its role as a caption augmentation assistant, along
with information regarding the main characters in the dataset. Then, we prompt
it with the captions of a story $S=\\{s_{1},...,s_{n}\\}$ in the following
form:

$1.\\{s_{1}\\}\ 2.\\{s_{2}\\}\ ...\ n.\\{s_{n}\\}$ (7)

ChatGPT returns an alternative version for each one of the captions:

$1.\\{\hat{s}_{1}\\}\ 2.\\{\hat{s}_{2}\\}\ ...\ n.\\{\hat{s}_{n}\\}$ (8)

At training time, we randomly pick either the original or the generated
caption for each training sample at each epoch. Using different text
descriptions for the same image is expected to shield the generative process
against over-fitting and help the Transformer distinguish between important
and insignificant textual concepts when generating visual tokens. At
inference, we only use the original captions of the dataset. Notably, this is
a completely image-agnostic caption augmentation method, _i.e_. the LLM
provides alternative captions without access to the image it is describing. In
contrast to [7], where the augmented data are used to train a contrastive
model, we use the alternative captions to train a generative model, where the
accuracy of the descriptions is of greater importance. An example of a caption
augmentation prompt is provided in Appendix 0.A.

###  4.4 Character Guidance

We propose Character Guidance, an auxiliary method to improve the generation
of characters in the images. An additional library of $2n_{c}$ Character
Embeddings (CE) is introduced to the Transformer, where $n_{c}$ is the number
of main, recurring characters in the dataset. For each character, we have a
positive embedding (the character is referenced in the caption) and a negative
one (the character is not referenced in the caption). When using this
technique, we concatenate $n_{c}$ embeddings to the input of the Transformer,
one for each character (positive for characters that are referenced in the
description and negative for the rest). The input now becomes:

$I_{i}=(z_{i}^{FLAT};s_{i}^{BPE};emb^{c}_{i})\in\mathbb{R}^{(\frac{H}{f}\cdot\frac{H}{f}+L+n_{c})\times
d}$ (9)

where $emb^{c}_{i}\in\mathbb{R}^{n_{c}\times d}$ is the concatenation of
positive CE, for characters that are _present_ in the current caption
($s_{i}$) and negative CE for characters that are _absent_ from the caption.

####  4.4.1 Training

In order to reinforce the model’s focus on the CE, we completely drop the text
conditioning for a percentage of training samples per batch, by replacing all
text embeddings with a [NULL] embedding and keep only the CE. Other than that,
the training process remains intact.

####  4.4.2 Inference

During inference, we compute three sets of logits222The term logits refers to
the unnormalized outputs of the Transformer when generating an image,
corresponding to the prompts described below.

##### Text-Conditional Prompt

The first set of logits ($\ell_{tc}\in\mathbb{R}^{N\times
K}$)333$N=(\frac{H}{f})^{2}$, while K is the size of the VQ-GAN’s token
library is computed by conditioning the generative process on text
descriptions.

##### Positive Character Prompt

The second set of logits ($\ell_{char}\in\mathbb{R}^{N\times K}$) is computed
using only Character Embeddings and completely dropping the captions.

##### Negative Character Prompt

We compute a third set of logits $\ell_{\overline{char}}$. To compute
$\ell_{\overline{char}}$ we completely drop text descriptions from the
transformer’s input, as we do when computing $\ell_{char}$. However, instead
of using positive embeddings for characters present in the description and
negative ones for absent characters, we do the opposite. Negative embeddings
are used for characters _present_ in the description and positive embeddings
for characters _absent_ in the description. In a sense, $p_{\overline{char}}$
is computed using the "complement" of the input that is used to generate
$p_{char}$. We employ Negative Prompting, inspired by the way it is used in
[2]. However, [2] uses this technique to push the logits away from a negative
text prompt, that needs to be provided by the user, while, we perform Negative
Prompting using a character prompt, that can be automatically formed, by
observing which characters are absent from the given text prompt.

##### Computing the Final Logits

The final logits are now computed as follows:

$\ell=(1-\lambda)\ell_{tc}+2\lambda\ell_{char}-\lambda\ell_{\overline{char}},\
\lambda\in[0,1)$ (10)

This combination of logits is formed at every step of the iterative inference
process. Softmax is applied to obtain probabilities at each one of the N
positions. These probabilities are used to sample tokens $Y$. $\ell_{tc}$
encode the specific information regarding the generation of an image from its
description, while $\ell_{char}$ encapsulate information solely regarding the
presence of desired characters in the image. Similarly,
$\ell_{\overline{char}}$ encode information about undesired characters. By
adding $\ell_{char}$ and subtracting $\ell_{\overline{char}}$ from $\ell_{tc}$
we attempt to actively push the text-conditional logits towards the generation
of the desired characters and away from the generation of undesired
characters.

##  5 Experiments and results

###  5.1 Experimental Setup

####  5.1.1 Software and Hardware

Our entire codebase is developed in PyTorch. For image tokenization we use VQ-
GAN’s [6] original implementation444 https://github.com/CompVis/taming-
transformers . Our Transformer model is adapted from an open-source
implementation of MUSE[2].555 https://github.com/lucidrains/muse-maskgit-
pytorch We perform all our experiments - both training and inference - on a
single NVIDIA V100 GPU (16GB). We consider this to be a strong indicator for
the resource-friendliness of our approach.

####  5.1.2 Dataset

Figure 2: Main characters featured in Pororo-SV

We train and test our approach on Pororo-SV [11], which is the most widely
adopted benchmark in previous SV works. We adopt the split proposed in [13],
which comprises of 10191/2334/2208 train/validate/test stories. Each story
contains 5 images ($64\times 64$ pixels) and the corresponding captions. There
are 9 recurring characters in the dataset, shown in Fig. 2. We choose Pororo-
SV because there are publicly available evaluation models for it, that have
been widely adopted in the past and make it easier to compare our results with
previous baselines.

####  5.1.3 Evaluation metrics

Following previous works, we adopt _FID_ to assess image quality. We also
employ _Char-F1_ and _Char-Acc_ , proposed in [13], to evaluate character
generation. Finally, we report BLEU scores based on video redescription[13] to
evaluate global semantic alignment. More details on the metrics are given in
Appendix 0.E.

####  5.1.4 Baselines

We compare our approach with several previous arts, using the reported results
of the respective papers. Among GANs we only compare it with VLC-StoryGAN
[12], since it has been shown to be the best of them. In terms of
Transformers, there are two baselines: VP-CSV [4] and CMOTA [1].

It would be unfair to include recent diffusion-based approaches[15, 8, 24] in
our main comparison, since they leverage pre-trained Latent Diffusion Models
(LDM)[19], trained on a vast dataset with massive compute. They also require
excessively expensive hardware to run on. For reference, AR-LDM[15] reports
using 8 NVIDIA A100 (80GB), which is equivalent to [$40\times$] the vRAM we
use. Besides, they only report FID scores, which we deem insufficient for SV
on their own. However, we acknowledge that in terms of FID they are superior
to our model. We elaborate more on this in Appendix 0.C. [17] and [23] are
applied on a modified version of the task, therefore a direct comparison is
not applicable.

###  5.2 Experiments

All of our models are trained from scratch. First, we train a VQ-GAN, with a
downsampling factor $f=8$, which corresponds to a $8\times 8$ resolution for
the latent image representations (since the image resolution is $64\times
64$). We use a library of $K=128$ latent embeddings. After completing VQ-GAN’s
training, we train MaskGST’s Transformer. The Transformer has 6 layers; 2 full
layers, followed by 4 Self-Layers, as shown in Fig. 1. We use Scaled Dot-
Product Attention [29], with the number of heads set to $n_{head}=8$. Based on
the results reported in [3], a cosine function is adopted for mask scheduling.
We train for 250 Epochs, with a learning rate $lr=1e-3$. Caption Augmentation
(Sec. 4.3) and Character Guidance (Sec. 4.4) are employed during training. To
reinforce Character Guidance, we drop the text descriptions in $20\%$ of
training samples, at each Epoch. Two variants of the Transformer are trained;
one with hidden dimension $d=1024$ and one with $d=2048$. At inference we use
both Positive and Negative Prompting, in terms of Character Guidance. To that
end, we form logits as shown in Eq. 10, with $\lambda=0.2$. We perform
inference over $T=20$ timesteps, for all our models.

###  5.3 Quantitative Results

We gather our results, as well as results from previous works in Tab. 1. In
the names of our models CG± refers to Character Guidance (positive and
negative), while w/ aug. captions is used for models trained with caption
augmentation.

Table 1: Results from our models ($d=1024$ and $d=2048$) and previous baselines, for the test set of Pororo-SV. We report scores for _FID_ (lower is better), as well as _Char-F1_ , _Char-Acc_ , _BLEU-2/3_ (higher is better).  Model Family | Model |  FID($\downarrow$) | Char-F1 | Char-Acc |  BLEU-2/3($\uparrow$)
---|---|---|---|---|---
GAN | VLC-StoryGAN[12] | 84.96 | 43.02 | 17.36 | 3.80/1.44
Auto-Regressive Transformer | VP-CSV[4] | 65.51 | 56.84 | 25.87 | 4.45/1.80
CMOTA[1] | 52.13 | 53.25 | 24.72 | 4.58/1.90
MaskGST-CG± w/ aug. captions | d=1024 | 51.65 | 60.46 | 33.62 | 4.82/2.00
d=2048 | 42.86 | 65.10 | 37.50 | 5.13/2.26

####  5.3.1 MaskGST-CG± w/ aug. captions ($d=1024$)

Our model with $d=1024$ performs better than all previous GAN and Transformer
architectures, in all metrics (lowest FID and highest Char-F1, Char-Acc,
BLEU-2/3). Especially Character metrics are raised significantly.
Specifically, in Char-F1 there is a $3.6$ point improvement, while in Char-Acc
we achieve an improvement of $7.7$ points, compared to the previous best (VP-
CSV). We largely attribute our models superiority in terms of character
generation, to our Character Guidance mechanism.

####  5.3.2 MaskGST-CG± w/ aug. captions ($d=2048$)

Doubling the hidden dimension to $d=2048$ improves our results across all
metrics. Most notably, there is a major improvement in terms of FID, which
decreases by 8.8 points compared to the $d=1024$ version of our model. We
assume that doubling the hidden dimension leaves more room to the model to
learn more complex and fine mappings between words and visual features, which
result in higher quality images with more details. The model’s ability to
learn more complex representations and produce more detailed images can
account for the improvements in the other metrics as well. Specifically, high
quality images will depict improved versions of the characters and produce
better BLEU scores through video captioning. In terms of Char-F1, Char-Acc and
BLEU-2/3, this model achieves the current SOTA, compared to all previous
methods.

###  5.4 Qualitative Results

In Fig. 3 we provide four examples of image-stories from the Pororo-SV test
set. For each story, we provide its captions, the ground-truth images
(Original), the image sequence generated by CMOTA [1] and the one generated by
our best model (MaskGST-CG± w/ aug. captions ($d=2048$)). To that end, we used
the pre-trained CMOTA model that was released by its
authors666https://github.com/yonseivnl/cmota. CMOTA is the most recent
Transformer model applied to the task of SV. We carry out a qualitative
comparison based on these examples.

Figure 3: Qualitative Comparison between our model (MaskGST-CG± w/ aug.
captions) and CMOTA[1] across 4 story examples.

Images generated by our model are of higher Visual Quality, with clearer and
more accurate features, compared to CMOTA’s that tend to be blurry. In terms
of Temporal Consistency, our model performs better as well. This is showcased
by the fact that backgrounds and characters are well maintained across the
sequences, contrary to CMOTA that is prone to switching between different
backgrounds (indoor and outdoor). Concerning Semantic Relevance (caption-image
alignment) our model is again superior. This is evident by the fact that it
generates the exact set of characters mentioned in captions, with remarkable
accuracy. On the contrary CMOTA struggles, especially in instances where
multiple characters are referenced. A more detailed qualitative comparison is
provided in Appendix 0.B.

###  5.5 Human Evaluation

In order to further investigate the qualitative results of Sec. 5.4, we
conducted a human survey across these three criteria which were also adopted
by previous works [13, 12, 1], comparing our model with CMOTA [1]. The
evaluation is done over 100 stories from the test set of Pororo-SV. Each story
is evaluated by 2 distinct annotators. The results of the study (Tab. 2)
indicate that our model is superior across all 3 criteria, thus supporting our
quantitative results. More details on the human study can be found in the
Appendix 0.F.

Table 2: Results of our human survey. We compare the results of our model MaskGST-CG± w/ aug. caption (d=2048) (Ours) against CMOTA’s, across three criteria. Our(%) and CMOTA(%) indicate the percentage of cases where each model was chosen by both annotators, whereas Tie(%) accounts for the remaining cases.  Criterion | Ours (%) | CMOTA (%) | Tie(%)
---|---|---|---
Visual Quality | 78% | 3% | 19%
Temporal Consistency | 66% | 8% | 26%
Semantic Relevance | 64% | 9% | 27%

###  5.6 Ablation Study

We perform an ablation study in order to determine the specific effects of the
different components of our approach. We consider MaskGST (Sec. 4.2.1) to be
our basic approach. Using _Caption Augmentation_ (Sec. 4.3) and employing
_Character Guidance_ (Sec. 4.4) are the two major components that build on top
of it, which we evaluate separately. Tab. 3 shows the results for these
experiments, on the test set of Pororo-SV. We use Transformers with $d=1024$.

####  5.6.1 Caption Augmentation

We observe that _Caption Augmentation_ (Tab. 3, MaskGST w/ aug. captions)
brings substantial imporovements across all metrics, with the exception of
BLEU scores, that are slightly lower, compared to the baseline. We believe
that providing alternative captions for the same image, during training is
beneficial for two reasons. On the one hand, it shields the model against
over-fitting. Additionally, it guides the model to focus on more important
textual concepts (_e.g_. character names), when generating the visual tokens,
since such concepts will probably appear in both versions of a caption.

####  5.6.2 Character Guidance

_Character Guidance_ (Positive and Negative) (Tab. 3, MaskGST-CG±) has a
dramatic effect on all metrics, compared to the baseline. This observation
confirms that Character Guidance indeed succeeds in directly impacting the
generation of characters by providing separate logits that guide the model
towards the production of the correct subset of them. This is evident by the
increase in Char-F1 and Char-Acc, by 9 and 7.5 points each. The improvement in
character generation enhances the overall quality of the images, which is
reflected in the additional improvements in FID and BLEU scores.

####  5.6.3 Caption Augmentation & Character Guidance

Combining the two methods (Tab. 3, MaskGST-CG± w/ aug. captions) seems to be
beneficial. FID and Char-F1 are further improved, with Char-Acc staying the
same, as in MaskGST-CG±, while BLEU-2/3 are slightly lower, yet still improved
compared to the baseline.

Table 3: Ablation Study for our method _MaskGST-CG ± w/ aug. captions_. We conduct experiments with hidden dimension $d=1024$ . We report scores for _FID_ (lower is better), as well as _Char-F1_ , _Char-Acc_ , _BLEU-2/3_ (higher is better).  Component | Model |  FID($\downarrow$) | Char-F1 | Char-Acc |  BLEU-2/3($\uparrow$)
---|---|---|---|---|---
baseline |  MaskGST | 66.12 | 50.48 | 26.12 | 4.68/2.01
\+ aug. captions |  {MaskGST w/ aug. captions} | 59.91 | 54.64 | 28.67 | 4.45/1.81
\+ char. guid. | MaskGST-CG± | 54.95 | 59.55 | 33.64 | 4.96/2.10
+{aug. captions, char. guid.} |  {MaskGST-CG± w/ aug. captions} | 51.65 | 60.46 | 33.62 | 4.82/2.00

###  5.7 Study on Character Guidance Factor ($\lambda$)

Figure 4: Comparison of Character Guidance Factor ($\lambda$) values for
different evaluation metrics.

Using our best model, MaskGST-CG± w/ aug.captions (d=2048), we perform an
analysis on the Character Guidance factor $\lambda$ (Eq. 10). Our results are
summarized in Fig. 4. Character metric curves increase with the increase in
$\lambda$ (with a slight decrease for $\lambda=0.8$). This is to be expected,
since by increasing $\lambda$ we pay more attention to the character logits vs
the text-conditional logits. FID, on other hand decreases (improves) until
$\lambda=0.4$ and then it increases for the next two experiments. This is
explained by the fact that improving the generation of Characters, also
improves the overall quality of the images, to a certain extend, since
Characters are a significant part of them. Once $\lambda$ becomes too large,
there is an adverse effect on overall image quality, since aspects of images
other that characters start to be neglected, because character logits are
weighed disproportionately highly. For $\lambda=0.4$ we get the optimal
combination of metrics (68.32, 41.40 and 42.49 for Char-F1, Char-Acc and FID,
respectively). However, in practice, we find that even $\lambda=0.4$ makes the
model focus too much on character generation and has an adverse effect on
inter-image coherence. We empirically find $\lambda=0.2$ to achieve the best
qualitative trade-off between Character Generation and overall quality of
image-stories. We elaborate more on our empirical findings in Appendix 0.D.

###  5.8 Resource Usage Analysis and Training/Inference Times

We compare our best model with previous Transformer architectures in terms of
training resources as well as Training/Inference Times. In terms of Training,
we compare with VP-CSV[4], that reports such information (Tab. 4). CMOTA[1]
does not specify this information, however since we perform inference with it,
we can carry out a comparison with our model, in terms of inference times, on
the same GPU (Tab. 5). We report GPU usage in (GB$\cdot$hours), _i.e_. the GBs
of vRAM in the GPU/GPUs, times the hours it was in use. It is evident that
both of our models are significantly more efficient compared to previous
Transformers applied to the task, while achieving better performance across
all metrics. More details on our calculations can be found in the Appendix
0.G.

Table 4: Training times/resources comparison between our models (d=1024 and d=2048) with VP-CSV. We use the stats reported by the authors of VP-CSV[4]. Metric | Ours (d=1024) | Ours (d=2048) | VP-CSV
---|---|---|---
Available GPUs |  $[1\times]$ V100 (16GB) |  $[1\times]$ V100 (16GB) |  $[4\times]$ A100 (40GB)
GPU usage | 576 (GB $\cdot$ hours) | 1712 (GB $\cdot$ hours) | 1920 (GB $\cdot$ hours)
Table 5: Comparison of inference times between our models (d=1024 and d=2048) with CMOTA[1]. We compare based on inference runs on the same NVIDIA V100 (16GB). Metric | Ours (d=1024) | Ours (d=2048) | CMOTA
---|---|---|---
GPU usage | 9.07 (GB $\cdot$ hours) | 25.07 (GB $\cdot$ hours) | 60.8 (GB $\cdot$ hours)

##  6 Limitations and Impacts

We acknowledge that our models are only evaluated on a Cartoon dataset, which
may be limiting for real-world applications. Additionally, despite our models’
merit, they still suffer in terms of FID, where large pre-trained models[15,
8, 24] remain unparalleled. Regarding the impact of our models, we cannot
foresee any direct missuse of them, since they are trained to produce
cartoons. Having said that, we are strongly opposed to any negative use of our
work, that hurts individuals or violates community guidelines and best
practices in any way.

##  7 Conclusion

In this paper, we adopt a MaskGIT model and enhance it with Cross-Attention
sub-layers for the task of Story Visualization. In addition, we propose a
novel Character Guidance method that improves the generation of characters, by
combining text-conditional and character-conditional outputs in the logit
space. We also employ an image-agnostic, LLM-driven caption augmentation
technique and show that it can be successfully used for generative tasks.
Using this combined approach, we achieved SOTA results over multiple metrics
on Pororo-SV, under a tight computational budget. We believe that our work
should encourage further research in the field of MaskGIT architectures, for
generative vision tasks. On the other hand, our Character Guidance method
could be explored more extensively in the context of SV and be paired with
large, possibly pre-trained models. Otherwise, it could be extended to other
generative tasks, where there is a particular interest for a specific set of
concepts, like characters in SV.

## References

  * [1] Ahn, D., Kim, D., Song, G., Kim, S.H., Lee, H., Kang, D., Choi, J.: Story visualization by online text augmentation with context memory. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 3125–3135 (2023)
  * [2] Chang, H., Zhang, H., Barber, J., Maschinot, A., Lezama, J., Jiang, L., Yang, M.H., Murphy, K., Freeman, W.T., Rubinstein, M., Li, Y., Krishnan, D.: Muse: Text-to-image generation via masked generative transformers (2023)
  * [3] Chang, H., Zhang, H., Jiang, L., Liu, C., Freeman, W.T.: Maskgit: Masked generative image transformer. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 11315–11325 (2022)
  * [4] Chen, H., Han, R., Wu, T.L., Nakayama, H., Peng, N.: Character-centric story visualization via visual planning and token alignment. In: Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing. pp. 8259–8272 (2022)
  * [5] Dai, H., Liu, Z., Liao, W., Huang, X., Cao, Y., Wu, Z., Zhao, L., Xu, S., Liu, W., Liu, N., Li, S., Zhu, D., Cai, H., Sun, L., Li, Q., Shen, D., Liu, T., Li, X.: Auggpt: Leveraging chatgpt for text data augmentation (2023)
  * [6] Esser, P., Rombach, R., Ommer, B.: Taming transformers for high-resolution image synthesis. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 12873–12883 (2021)
  * [7] Fan, L., Krishnan, D., Isola, P., Katabi, D., Tian, Y.: Improving clip training with language rewrites. Advances in Neural Information Processing Systems 36 (2024)
  * [8] Feng, Z., Ren, Y., Yu, X., Feng, X., Tang, D., Shi, S., Qin, B.: Improved visual story generation with adaptive context modeling. In: Rogers, A., Boyd-Graber, J., Okazaki, N. (eds.) Findings of the Association for Computational Linguistics: ACL 2023. pp. 4939–4955. Association for Computational Linguistics, Toronto, Canada (Jul 2023). https://doi.org/10.18653/v1/2023.findings-acl.305, https://aclanthology.org/2023.findings-acl.305
  * [9] Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., Bengio, Y.: Generative adversarial nets. Advances in neural information processing systems 27 (2014)
  * [10] Li, J., Li, D., Xiong, C., Hoi, S.: Blip: Bootstrapping language-image pre-training for unified vision-language understanding and generation. In: International Conference on Machine Learning. pp. 12888–12900. PMLR (2022)
  * [11] Li, Y., Gan, Z., Shen, Y., Liu, J., Cheng, Y., Wu, Y., Carin, L., Carlson, D.E., Gao, J.: Storygan: A sequential conditional gan for story visualization. 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) pp. 6322–6331 (2018), https://api.semanticscholar.org/CorpusID:54457433
  * [12] Maharana, A., Bansal, M.: Integrating visuospatial, linguistic, and commonsense structure into story visualization. In: Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. pp. 6772–6786 (2021)
  * [13] Maharana, A., Hannan, D., Bansal, M.: Improving generation and evaluation of visual stories via semantic consistency. In: Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. pp. 2427–2442 (2021)
  * [14] Maharana, A., Hannan, D., Bansal, M.: Storydall-e: Adapting pretrained text-to-image transformers for story continuation (2022)
  * [15] Pan, X., Qin, P., Li, Y., Xue, H., Chen, W.: Synthesizing coherent story with auto-regressive latent diffusion models. In: Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. pp. 2920–2930 (2024)
  * [16] Radford, A., Kim, J.W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al.: Learning transferable visual models from natural language supervision. In: International conference on machine learning. pp. 8748–8763. PMLR (2021)
  * [17] Rahman, T., Lee, H.Y., Ren, J., Tulyakov, S., Mahajan, S., Sigal, L.: Make-a-story: Visual memory conditioned consistent story generation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 2493–2502 (2023)
  * [18] Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Radford, A., Chen, M., Sutskever, I.: Zero-shot text-to-image generation. In: International Conference on Machine Learning. pp. 8821–8831. PMLR (2021)
  * [19] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., Ommer, B.: High-resolution image synthesis with latent diffusion models. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 10684–10695 (2022)
  * [20] Ruan, L., Ma, Y., Yang, H., He, H., Liu, B., Fu, J., Yuan, N.J., Jin, Q., Guo, B.: Mm-diffusion: Learning multi-modal diffusion models for joint audio and video generation. In: 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). pp. 10219–10228 (2023). https://doi.org/10.1109/CVPR52729.2023.00985
  * [21] Saharia, C., Chan, W., Saxena, S., Li, L., Whang, J., Denton, E.L., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Salimans, T., et al.: Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems 35, 36479–36494 (2022)
  * [22] Sennrich, R., Haddow, B., Birch, A.: Neural machine translation of rare words with subword units (2016)
  * [23] Shen, X., Elhoseiny, M.: Storygpt-v: Large language models as consistent story visualizers (2023)
  * [24] Song, T., Cao, J., Wang, K., Liu, B., Zhang, X.: Causal-story: Local causal attention utilizing parameter-efficient tuning for visual story synthesis (2023)
  * [25] Song, Y.Z., Tam, Z.R., Chen, H.J., Lu, H.H., Shuai, H.H.: Character-preserving coherent story visualization. In: Proceedings of the European Conference on Computer Vision (ECCV) (2020)
  * [26] Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., Lample, G.: Llama: Open and efficient foundation language models (2023)
  * [27] Ubani, S., Polat, S.O., Nielsen, R.: Zeroshotdataaug: Generating and augmenting training data with chatgpt (2023)
  * [28] Van Den Oord, A., Vinyals, O., et al.: Neural discrete representation learning. Advances in neural information processing systems 30 (2017)
  * [29] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L., Polosukhin, I.: Attention is all you need. In: Proceedings of the 31st International Conference on Neural Information Processing Systems. p. 6000–6010. NIPS’17, Curran Associates Inc., Red Hook, NY, USA (2017)
  * [30] Whitehouse, C., Choudhury, M., Aji, A.F.: Llm-powered data augmentation for enhanced cross-lingual performance (2023)
  * [31] Wozniak, S., Kocon, J.: From big to small without losing it all: Text augmentation with chatgpt for efficient sentiment analysis. In: 2023 IEEE International Conference on Data Mining Workshops (ICDMW). pp. 799–808. IEEE Computer Society, Los Alamitos, CA, USA (dec 2023). https://doi.org/10.1109/ICDMW60847.2023.00108, https://doi.ieeecomputersociety.org/10.1109/ICDMW60847.2023.00108
  * [32] Xu, T., Zhang, P., Huang, Q., Zhang, H., Gan, Z., Huang, X., He, X.: Attngan: Fine-grained text to image generation with attentional generative adversarial networks. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 1316–1324 (2018)
  * [33] Zhang, H., Xu, T., Li, H., Zhang, S., Wang, X., Huang, X., Metaxas, D.N.: Stackgan: Text to photo-realistic image synthesis with stacked generative adversarial networks. In: Proceedings of the IEEE international conference on computer vision. pp. 5907–5915 (2017)
  * [34] Zheng, J., Liu, H., Feng, Y., Xu, J., Zhao, L.: Casf-net: Cross-attention and cross-scale fusion network for medical image segmentation. Comput. Methods Prog. Biomed. 229(C) (feb 2023). https://doi.org/10.1016/j.cmpb.2022.107307, https://doi.org/10.1016/j.cmpb.2022.107307

##  Appendix 0.A Details on Caption Augmentation via ChatGPT

For every five-caption story in the training set, we carry out a caption
augmentation conversation with ChatGPT. The conversation starts with a system
message that describes its role as a caption augmentation assistant.
Additionally, this message contains basic information about the characters in
the dataset. We presume that the characters in Pororo-SV constitute a very
specific domain of knowledge. Therefore we consider this information to be
necessary domain-specific context for the model to provide suitable
alternative captions. The next message, in the conversation contains the
captions of a story, numbered from 1 to 5. ChatGPT responds with a message
containing the alternative captions, numbered in the same way. We use
ChatGPT3.5 Turbo, through the API. An example can be seen in Figure 5.

Figure 5: Example of caption augmentation using ChatGPT.

##  Appendix 0.B Detailed Qualitative Comparison

In this section we provide a more thorough qualitative comparison with
CMOTA[1], based on the examples that are provided in the main paper. The
following comments refer to Fig. 3.

####  0.B.0.1 Image Quality

In terms of image quality, it is evident from the examples that our model is
superior to CMOTA. For example, in the top-left panel, all of our images
contain discernible objects. Additionally, the characters are depicted with
high quality (_e.g_. the dinosaur’s eyes and the penguin’s beak). On the
contrary, most of CMOTA’s images in the same example are blurry, while some of
them are completely incomprehensible.

####  0.B.0.2 Temporal Consistency

Our image-stories are more coherent than CMOTA’s across all examples.
Especially, in the top-right panel, our model achieves remarkable results with
the room remaining similar in all images and the characters holding similar
positions, as well.CMOTA, on the other hand, struggles to maintain these
features, as evidenced by the contrast between the first and final images,
which depict outdoor scenes, and the three middle ones, which are from an
indoor space.

####  0.B.0.3 Semantic Relevance

The term Semantic Relevance refers to whether the generated image is relevant
to the corresponding caption. In regards to this, our model shows remarkable
capabilities, especially in terms of producing the mentioned characters. For
instance, in the bottom-right panel, it manages to produce the correct subset
of characters in all images, although that changes constantly. On the
contrary, CMOTA struggles to produce multiple characters. This is most evident
in the second image of the botton-right example, where the attempt to produce
a big number of characters collapses into a blurry result.

##  Appendix 0.C Comparison with Large Pre-Trained Models

As we have stated in the main paper, there has been a series of recent
approaches[15, 8, 24] that modify and fine-tune Latent Diffusion Models
(LDMs)[19] for the task of Story Visualization. A direct comparison with them
is not fair, due to the extensive pre-training of LDM, as well as the
significantly superior hardware that is used to develop these models, compared
to ours. However, for the sake of completeness, we devote this section to a
quantitative comparison.

AR-LDM[15] models the image sequence generation in an auto-regressive manner.
Specifically, frames are generated by a Latent Diffusion Model[19] one by one,
from the first one to the last one. The Diffusion process is conditioned on
the current caption, as well as a multi-modal representation of each previous
caption and generated image pair. The current caption is encoded using
CLIP[16], while multi-modal features from previous caption-image pairs are
extracted via BLIP[10]. ACM-VSG[8] similarly models frame generation auto-
regressively, whilst conditioning on past multi-modal context. However, it
also introduces an adaptive guidance mechanism that aims to push frames of
similar captions, to be similar as well. Causal-Story[24] improves AR-LDM by
introducing a local, causal attention mask that limits the size of historical
context tokens, to eliminate confusion, caused by interfering captions.

###  0.C.1 Story Visualization

In Table 6 we provide a comprehensive historical overview of results on the
task of SV. We include large-scale diffusion models[15, 8, 24], along with all
other previous works, as well as our models with $d=1024$ and $d=2048$. As it
is evident from the table, all diffusion models are evaluated solely on FID,
where they are by far better than all other approaches, including ours. Among
other approaches, that are trained from scratch, ours holds the best FID, as
well as the overall SOTA in terms of all other metrics.

Table 6: Story Visualization results for all previous works, including recent diffusion models, on the test set of Pororo-SV. We report scores for _FID_ (lower is better), as well as _Char-F1_ , _Char-Acc_ , _BLEU-2/3_ (higher is better). Parameter counts for [13], [12] and [1] are taken from [1]. Parameters counts for [15] and [8] are takens from [8].  Model Family | Model | #param |  FID($\downarrow$) | Char-F1 | Char-Acc |  BLEU-2/3($\uparrow$)
---|---|---|---|---|---|---
GAN | StoryGAN[11] | - | 158.06 | 18.59 | 9.34 | 3.24 / 1.22
CP-CSV[25] | - | 140.24 | 21.78 | 10.03 | 3.25 / 1.22
DUCO-SG[13] | 101M | 96.51 | 38.01 | 13.97 | 3.68 / 1.34
VLC-SG[12] | 100M | 84.96 | 43.02 | 17.36 | 3.80 / 1.44
Auto-Regressive Transformer | VP-CSV[4] | - | 65.51 | 56.84 | 25.87 | 4.45 / 1.80
CMOTA[1] | 96.6M | 52.13 | 53.25 | 24.72 | 4.58 / 1.90
MaskGST-CG± w/ aug. captions | d=1024 | 105M | 51.65 | 60.46 | 33.62 | 4.82 / 2.00
d=2048 | 276M | 42.86 | 65.10 | 37.50 | 5.13 / 2.26
Diffusion Model | AR-LDM[15] | 1.5B | 16.59 | - | - | -
ACM-VSG[8] | 850M | 15.48 | - | - | -
Causal-Story[24] | - | 16.28 | - | - | -

###  0.C.2 Story Continuation

[15, 8, 24] also report results for the task of Story Continuation[14]. Story
Continuation (SC) is similar to SV. However, in SC the first frame is
considered to be given as input and the rest of the frames need to be
generated. In Table 7 we report SC results for our model ($d=2048$), as well
as AR-LDM[15], ACM-VSG[8] and Causal-Story[24]. Additionally, we include
results for StoryDALLE (fine-tune)[14] and Mega-StoryDALLE[14] that are both
based on pre-trained auto-regressive transformers. These where the first
models applied to the task of SC. We observe that although our model has the
lowest parameter count, by far, it outperforms all the other large,
extensively pre-trained models in terms of Char-F1 and Char-Acc. This speaks
of the remarkable merit of our Character Guidance method. In terms of FID,
large pre-trained models remain unparalleled, as in SV. Since our model is
trained for SV, in order to evaluate it for SC, we just discard the first
generated frame and evaluate it using the remaining four frames.

Table 7: Story Continuation results on the test set of Pororo-SV, for large pre-trained models, as well as MaskGST-CG± w/ aug. captions (d=2048) . We report scores for _FID_ (lower is better), as well as _Char-F1_ , _Char-Acc_ (higher is better). Model | #param |  FID($\downarrow$) | Char-F1 | Char-Acc
---|---|---|---|---
StoryDALLE(fine-tune)[14] | 1.3B | 25.90 | 36.97 | 17.26
Mega-StoryDALLE[14] | 2.8B | 23.48 | 39.91 | 18.01
AR-LDM[15] | 1.5B | 17.40 | - | -
ACM-VSG[8] | 850M | 15.36 | 45.71 | 22.62
Causal-Story[24] | - | 16.98 | - | -
Ours (d=2048) | 276M | 43.31 | 65.32 | 37.38

##  Appendix 0.D Empirical Results for the Character Guidance Factor
($\lambda$)

In the main paper, we present a study on the Character Guidance Factor
($\lambda$). As we point out there, despite the fact that $\lambda=0.4$ yields
better quantitative results, we empirically find that $\lambda=0.2$ achieves
the best trade-off between character generation and overall quality of
generated stories. Specifically, for $\lambda=0.4$, the model becomes more
prone to inconsistency regarding the background of different story frames. In
Figure 6 we provide three examples that demonstrate this observation. In all
three image-stories generated with $\lambda=0.4$, there is an interleaving of
indoor and outdoor backgrounds, whereas for $\lambda=0.2$, the background is
held more consistent.

Figure 6: Three examples of generated image sequences that highlight the fact
that a value of $\lambda=0.4$ hinders the coherence between images, compared
to the results for $\lambda=0.2$.

##  Appendix 0.E Details on Reported Metrics

Following previous works, we adopt _FID_ to assess the quality of the
generated images. In order to evaluate the presence of main characters in the
images we use _Char-F1_ and _Char-Acc_ , proposed in [13]. These metrics are
calculated using a multi-label classifier that recognizes the presence of
character in the generated images. The classifier’s predictions are compared
to the ground-truth character references in the captions to calculate F1-score
and Accuracy. We use the fine-tuned Inception-V3 from the original paper,
available here777https://github.com/adymaharana/StoryViz. Finally, we follow
[13] in using a video captioning model that produces a single caption for each
generated image-story. The generated captions are compared to the ground truth
ones to calculate _BLEU_ scores. We use the fine-tuned video captioner,
available at the same link as the classifier.

##  Appendix 0.F Details on Human Evaluation Survey

We conducted a human survey to evaluate the stories generated by our model vs
CMOTA[1]. The evaluation is done across 100 image-stories. Each story is
evaluated across three criteria:

  * •

_Visual Quality_ refers to whether the images are visually appealing, rather
than blurry and difficult to understand.

  * •

_Temporal Consistency_ measures whether the images are consistent with each
other, having a common topic and naturally forming a story, rather than
looking like 5 independent images.

  * •

Semantic Relevance refers to whether the images accurately reflect the
captions and the characters mentioned in them.

A screenshot from our survey can be seen in Figure 7. In order to eliminate
bias, for half of the examples we assign the label Model 1 to the images
generated by our model and Model 2 to the images generated by CMOTA, while for
the other half, the labels are assigned the other way around. Every pair of
image-stories is annotated by two distinct annotators. We make sure not to
store any personal information about the users in our survey.

Figure 7: A screenshot from our human evaluation survey.

##  Appendix 0.G Training Times and Resource Usage

####  0.G.0.1 Training

We approximately spend 36 and 107 hours to train our Transformers with
$d=1024$ and $d=2048$, respectively, on a single NVIDIA V100 (16GB). This is
equivalent to (36 hours) $\cdot$ (16GB) = 576 (GB $\cdot$ hours) and (107
hours)$\cdot$ (16GB) = 1712 (GB $\cdot$ hours) of GPU usage, respectively. For
reference, VP-CSV [4] reportedly uses 4 NVIDIA A100 (40GB) for 12 hours. This
is equivalent to (12 hours) $\cdot$ (4$\cdot$40 GB) = 1920 (GB $\cdot$ hours)
of GPU usage, without taking into account that A100 is a more modern GPU than
V100. CMOTA [1] does not report training resource usage.

####  0.G.0.2 Inference

Since we performed inference for our models, as well as CMOTA, on the same
GPU, we can carry out a fair comparison. Our models with $d=1024$ and $d=2048$
need 34 minutes and 94 minutes, respectively to perform inference for the 2208
stories of the test set. This is equivalent to 9.07 (GB $\cdot$ hours) and
25.07 (GB $\cdot$ hours). For the same task, CMOTA spent 228 minutes, which
translates to 60.8 (GB $\cdot$ hours). It is evident that our method is
significantly more time-efficient compared to CMOTA. Even the larger version
of our model, is more than $[2\times]$ more efficient, during inference
compared to it. This can be largely attributed to the inference scheme of
MaskGIT-style transformers, that produce multiple visual tokens per step,
compared to auto-regressive transformers, like CMOTA, that infer visual tokens
one at a time.

##  Appendix 0.H More Qualitative Examples

Figure 8 and Figure 9 show more Story Generation examples using our _MaskGST-
CG ± /w aug. captions_. For each example we provide the images generated by
our model, the input captions and the ground-truth images (original) that
correspond to these captions.

Figure 8: More Story Generation Examples using our model _MaskGST-CG ± /w aug.
captions_. Figure 9: More Story Generation Examples using our model _MaskGST-
CG ± /w aug. captions_.


# RecycleNet: Latent Feature Recycling Leads to Iterative Decision Refinement

Gregor Koehler1,2, Tassilo Wald1,3, Constantin Ulrich1.4, David Zimmerer1,3,
Paul F. Jaeger3,5, Jörg K.H. Franke7, Simon Kohl8, Fabian Isensee1,3,6 and
Klaus H. Maier-Hein1,3,4,9

1German Cancer Research Center (DKFZ) Heidelberg, Division of Medical Image
Computing, Germany
2Helmholtz Information and Data Science School for Health,
Karlsruhe/Heidelberg, Germany
3Helmholtz Imaging, DKFZ
4National Center for Tumor Diseases (NCT), NCT Heidelberg,
a partnership between DKFZ and University Medical Center Heidelberg
5Interactive Machine Learning Group, DKFZ
6Applied Computer Vision Lab, DKFZ
7Machine Learning Lab, University of Freiburg, Freiburg, Germany
8Latent Labs (latentlabs.com), London, UK
9Pattern Analysis and Learning Group, Heidelberg University Hospital,
Heidelberg, Germany
<EMAIL_ADDRESS>
###### Abstract

Despite the remarkable success of deep learning systems over the last decade,
a key difference still remains between neural network and human decision-
making: As humans, we can not only form a decision on the spot, but also
ponder, revisiting an initial guess from different angles, distilling relevant
information, arriving at a better decision. Here, we propose RecycleNet, a
latent feature recycling method, instilling the pondering capability for
neural networks to refine initial decisions over a number of recycling steps,
where outputs are fed back into earlier network layers in an iterative
fashion. This approach makes minimal assumptions about the neural network
architecture and thus can be implemented in a wide variety of contexts. Using
medical image segmentation as the evaluation environment, we show that latent
feature recycling enables the network to iteratively refine initial
predictions even beyond the iterations seen during training, converging
towards an improved decision. We evaluate this across a variety of
segmentation benchmarks and show consistent improvements even compared with
top-performing segmentation methods. This allows trading increased computation
time for improved performance, which can be beneficial, especially for safety-
critical applications.

##  1 Introduction

Over the past decade, the field of computer vision has witnessed an
unprecedented paradigm shift due to the advent and proliferation of deep
learning algorithms. Neural networks have become the de facto standard for a
variety of tasks, excelling in their ability to solve previously impossible
tasks across many domains and modalities. One of the most intriguing
distinctions between human cognition and neural networks, however, is the
former’s capacity for iterative decision-making - a skill that is still
lacking from most recent artificial systems. Humans exhibit the innate ability
to dynamically revisit and revise their initial decisions, evaluating their
options from multiple perspectives, and continuously improving their decision
based on new information. This process underlies an inherent characteristic of
human decision-making – that of iterative refinement. It allows humans to
evolve their understanding over time, improving the quality of decisions,
especially in complex, non-deterministic scenarios. By stark contrast,
conventional deep learning architectures have typically operated in a one-
shot, feed-forward manner, lacking the property of iterative revision.
In this work, we seek to bridge this gap, to bring neural networks a step
closer to the iterative decision-making process that characterizes human
cognition. We introduce RecycleNet, a simple approach to deep learning that
leverages the concept of latent feature recycling, enabling neural networks to
refine their initial predictions over a series of iterative steps. The
advantage of RecycleNet lies in its universal applicability - it makes minimal
assumptions about the underlying architecture and is easily adaptable across a
wide array of contexts. Medical image segmentation, with its critical
implications for diagnosis and treatment in healthcare, serves as an excellent
testing ground for our approach. The task’s complexity and the inherent noise
in medical imaging data pose formidable challenges that demand robust,
reliable, and refined predictions.

Figure 1: Schematic overview of the proposed U-Net feature recycling. $n$
depicts the number of recycling cycles where the features close to the
network’s decoder are fed back into early encoder features. The letters I, R
and O refer to the input projection, recycling module and output projection,
as described in Algorithm 1.

Through our evaluation, we demonstrate that RecycleNet exhibits the remarkable
property of refining its predictions iteratively, even beyond the iterations
witnessed during training. The results clearly outperform the state-of-the-art
segmentation methods across a range of segmentation benchmarks, demonstrating
the promise of our approach. At the heart of RecycleNet lies a trade-off - the
opportunity to exchange increased computational time for a significant
improvement in performance. For safety-critical applications, where the stakes
are high, and the margin for error is virtually non-existent, this trade-off
could be an important step in enhancing the reliability and precision of
decision-making processes in neural networks.

##  2 Related Work

Neural networks’ inability to refine initial predictions has been addressed in
approaches which introduce varying degrees of additional complexity.

#### Refinement modules:

One straight-forward way to refine initial network outputs makes use of
additional modules. In the context of image segmentation, these refinement
modules typically act on segmentation maps [15] or features close to the
segmentation layers [4] and make use of additional layers to refine a main
network’s outputs. This naturally introduces complexity by introducing
additional parameters to the original network. In contrast, our suggested
technique of latent feature recycling operates without requiring any
additional parameters, thereby ensuring a more seamless integration in
situations where the cost of extra parameters is prohibitive.

#### Recurrent Learning:

Another natural approach to refining initial network outputs is to cast
refinement as a recurrent learning problem, with refinement steps as the
temporal axis. To render typical computer vision network architectures as
Recurrent Neural Networks, either the whole network [19] or key parts of the
network [18] are adapted. Alternatively, the recurrent learning can also be
done on just the segmentation outputs [12]. While closely related to the
proposed latent feature recycling, these approaches interfere strongly with
the network architecture and due to their recurrent network formulation, come
with substantial memory costs during training, which quickly become
prohibitive for example in the context of medical image segmentation, where
large 3D receptive fields are required to capture all relevant context for the
task.

#### Multi-stage approaches:

To alleviate the memory costs connected to Recurrent Neural Networks, multi-
stage approaches can be employed to refine previous stage segmentations [16]
or learn based on the error feedback coming from the previous iteration [3].
While this does not come at the cost of substantially increased memory demands
during training, multi-stage approaches often require multiple complete
training cycles. Comparatively, our proposed method requires less training
time since the additional forward passes are performed on a single sample.

#### Additional loss terms:

The ability to refine predictions has also been explored by repeatedly
applying a given stateful network architecture and learning a stopping
criterion using an additional loss term together with the task loss [2]. This,
however, requires a careful balancing between loss terms which is not present
in the proposed method.

#### Output recycling:

Recently, Jumper et al. [9] have proposed a technique that uses output
structures in the context of protein structure prediction to refine initial
guesses without additional modules. They re-use outputs from certain
transformer blocks [17] of the architecture over multiple iterations to refine
the predicted structures. While this is closely related to our proposed
technique, we show that recycling not only a part of the network architecture,
but a whole convolutional segmentation network leads to refined predictions.
Additionally, recycling features instead of outputs allows integrating this
mechanism in a flexible way without special requirements w.r.t. the network
architecture. Crucially, we also introduce a robust training schedule for
recycling and demonstrate its importance for both reliable results across
datasets, as well as an emergent convergence property where segmentation
performance increases monotonously when increasing the number of cycles during
inference.

##  3 RecycleNet

Our proposed method, referred to as RecycleNet, relies on increasing the
number of forward passes through a large part of the network (referred to as
cycles), both during training and inference. To instill the capability to
refine initial decisions over a number of such cycles, features close to the
output are fed back into early layers of the neural network via a simple
addition operation. Figure 1 depicts a schematic of the proposed latent
feature recycling process. A given neural network architecture (here the U-Net
[14]), can be partitioned into three disjoint parts: The input projection I,
the recycling module R and the output projection O (see Fig. 1). The recycling
module R is not an additional module, but rather refers to the part of the
architecture where features should be recycled. The recycling process, where
recycled features are summed onto earlier feature representations, is repeated
$N_{c}$ times during training, where $N_{c}$ is sampled uniformly from a
predefined range (see Section 3.1). After the recycling process, the output
projection produces the final output. However, each individual cycle can also
be projected to a meaningful prediction, allowing for introspection and
ensembling. The recycling mechanism is described in Algorithm 1. We note that
while it is, in general, possible to use gradients accumulated for more than
one cycle, we find this to be impractical due to the memory demands involved.
So instead, we only use gradients for the last iteration.

Input: Maximum number of cycles $N_{max}$, model input $x$, input projection
I, recycling module R, output projection O

1 Project input into recycling feature space: $z=I(x)$

2 Sample number of cycles: $N_{c}=RandInt(1,N_{max})$

3 Initialize recycling features as zeros: $r=0$

4 for  _all cycles $i\in[1,...,N_{c}]$_ do

5 if  _$i <N_{c}$_ then

6 r = r.detach() $\\#$ gradients only for last cycle

7 r = R(z, r) $\\#$ 1 cycle

8 end for

9Project to output: $\hat{y}=O(r)$

return $loss(\hat{y},y)$

Algorithm 1 Latent Feature Recycling: Training

Reusing the recycling features $r$ can be achieved in various ways. We propose
a simple addition of normalized recycling features to the input projection,
similar to the standard practice of adding position encodings in the context
of Language Models [17]:

$R(z,r)=R(z+norm(r))$ (1)

This approach requires the feature dimensions after the input projection to
match the dimensions after the recycling module. This property is typically
fulfilled for attention-based transformer architectures, as well as U-Nets
[17]. In architectures, where this is not the case, other conditioning
mechanisms, e.g. using projection layers, can be applied. We propose the
integration of latent feature recycling in the context of the U-Net [14], a
network architecture which is ubiquitous in medical image segmentation.
As depicted schematically in Figure 1, we propose U-Net feature recycling by
reusing features close to the output of the network’s decoder at earlier
layers, e.g. after the encoder’s first convolutional block. This enables the
network to revisit features based on which previous predictions would be
computed, thus instilling the capability to iteratively refine early decision
hypotheses over a number of cycles.

###  3.1 Robust Training Schedule

Feature recycling introduces a single new hyperparameter to an existing neural
network training, the number of cycles $N_{c}$. This hyperparameter determines
how many shots the network gets to refine initial predictions during training.
At the beginning of the training, the initial predictions might be unreliable
to the extent that there is little value in refining them step by step. To
combat this, we introduce a robust recycling schedule, where during an initial
warm-up phase, only a single cycle (no recycling) is used, therefore
defaulting to standard network training during this period. Over time, we
incrementally increase the range of possible cycles to allow for more and more
refinement steps. As the number of cycles during training is not
deterministic, the network is incentivized to distill useful information for
each next recycling step in the recycled features. This helps learning an
iterative refinement, even when only using gradients from the last cycle.

##  4 Experiments and Results

In the following, we present the experiments and results based on the proposed
U-Net feature recycling in the context of challenging medical image
segmentation datasets. To reliably compare the proposed method with a strong
baseline, we compare to the widely used nnU-Net and implement the proposed
method in the same well-tested data pipeline and training framework [6].

###  4.1 Datasets and Evaluation

To demonstrate the general effectiveness of the proposed U-Net latent feature
recycling for medical image segmentation, we test the proposed method on a
range of established segmentation tasks covering various dataset sizes,
segmentation targets and task difficulties. These datasets include the Kidney
Tumor Segmentation (KiTS 2019) dataset [5], the Liver Tumor Segmentation task
of the Medical Segmentation Decathlon [1], the Multi-Atlas Labeling Beyond the
Cranial Vault (BTCV) challenge [11] and the large-scale Abdominal Multi-Organ
Benchmark (AMOS) for versatile medical image segmentation [8]. To also test
the proposed method in the context of MRI tasks, we additionally include two
tasks of the Combined Healthy Abdominal Organ Segmentation (CHAOS) challenge
[10].
To assess segmentation performance, we use the average foreground Dice
coefficient (DSC). It measures the overlap between predicted segmentation and
ground truth, popular in medical image segmentation for its capacity at
handling imbalanced datasets.

###  4.2 Baseline Models

To ensure a fair comparison, we implement all models in the nnU-Net framework
[6], which is regarded as the state-of-the-art framework for medical image
segmentation and serves as the basis for numerous successful segmentation
challenge contributions [7]. Naturally, the default nnU-Net is also included
as a baseline. We also test against the best-performing model proposed in
Wang, Yu et al. [18] which still has a manageable memory demand when applied
in the medical image segmentation domain, referred to as DRU. As the DRU was
not originally tested in the context of 3D segmentation, we base our
implementation on the authors’ code and implement the model within the nnU-Net
framework. This allows to make fair comparisons w.r.t. performance, memory
demands and training time requirements.

###  4.3 RecycleNet

Model | KiTS | BTCV | CHAOS | AMOS | Liver T.
---|---|---|---|---|---
CV | Test | CV | Test | CV(T5) | T3 | T5 | CV | Test | CV
nnU-Net | 89.29 | 89.04 | 82.96 | 87.21 | 94.77 | 93.49 | 91.47 | 88.58 | 90.68 | 78,84
DRU [18] | 89.58 |  | 82.99 |  | 91.62 |  |  | 88.62 |  | 80.36
RecycleNet | 90.26 | 89.11 | 83.75 | 87.80 | 94.92 | 93.48 | 91.85 | 88.77 | 90.82 | 79.88
Table 1: Average DSC scores for 5-fold cross-validation (CV) and public
leaderboard held-out test sets (Test). We compare the vanilla nnU-Net [6] with
the proposed method and don’t use post-processing on either model’s
predictions for a fair comparison. T3 and T5 refer to the test sets of two MRI
tasks of the CHAOS challenge, selected to test the proposed method on a
different modality than CT.

We implement the proposed RecycleNet as an adaptation of nnU-Net, making use
of the same, well-tested, 3D full resolution U-Net architectures. For each
dataset, we adopt the same feature recycling strategy, recycling the feature
maps before the last convolutional layer and feed them back into the encoder
after the first convolutional stage, see Section 3. Through this, we can keep
the automatic configuration property offered by nnU-Net and roll out the same
proposed recycling method for any given segmentation task. For a fair
comparison between the proposed method and the baselines, we use the
preprocessing pipeline of nnU-Net [6] for all experiments and evaluate on
public leaderboards when possible. Where not possible, we employ an identical
5-fold cross validation. We evaluate single full resolution 3D models without
ensembling or postprocessing. As the recycling training schedule, we make use
of the schedule proposed in Section 3.1, starting with 1 cycle for the first
200 epochs, then gradually increasing the range from which the number of
cycles is sampled by 1 every 200 epochs, to a maximum of 3. During inference,
we increase the number of cycles to a maximum of 7 cycles to benefit from the
property discussed in Section 4.4 and always report metrics using the maximum
number of cycles.

Table 1 shows the results comparing RecycleNet to the nnU-Net and DRU
baselines on 5-fold cross validation as well as public leaderboard held-out
test sets, where possible. We show clear performance improvements compared to
both baselines across all evaluation datasets but the liver tumor segmentation
dataset, where the DRU [18] shows the best cross validation score. We note
that although the differences between the individual methods seem small, they
can be regarded as substantial improvements in the saturated performance
domain that is medical image segmentation. We use a fixed recycling schedule
and a fixed number of cycles during inference (determined based on the BTCV
cross-validation, employed on all datasets), leaving all nnU-Net training
hyperparameters untouched. We suspect that further improvements are possible
by fine-tuning the training and recycling schedule on a given target dataset.
From these results, we conclude that the proposed latent feature recycling
represents an effective way to instill iterative refinement capabilities for
even strong image segmentation models. The cost for this performance surplus
is an increased training memory demand, as well as increased training and
inference time (see Section 4.5).

###  4.4 Prediction Convergence of RecycleNet

To test the iterative refinement capabilities of the proposed method, we
investigate predictions over the number of cycles, both quantitatively and
qualitatively. Figure 2 shows the performance in terms of average Dice score
in a 5-fold cross validation on the BTCV dataset.

Figure 2: Evaluation of the 5-fold cross validation performance of the
proposed method on the BTCV dataset over number cycles in inference. The
dotted red line represents the default nnU-Net performance, the yellow to
orange bars represent the Recycling nnU-Net’s performance for different
numbers of cycles during inference.

We observe a monotonous performance increase when increasing the number of
recycling cycles during inference, asymptotically converging towards a
saturation DSC value. This suggests an interesting property arising from the
introduction of feature recycling. We refer to this property as prediction
convergence, where the model naturally converges towards a refined prediction.
Surprisingly, even a single cycle can yield improvements compared with the
baseline, as shown in Figure 2. However, this observation is not consistent
across datasets. We also note that the network learns to refine predictions
even beyond the number of cycles seen during training, while still improving
on prior predictions. This property is also reinforced when qualitatively
inspecting the segmentation predictions the proposed method creates for single
samples, as shown for samples of two segmentation tasks in Figure 3. We see an
increased prediction quality beyond the 3 cycles seen during training for both
samples.

Figure 3: Qualitative showcase of refined segmentations for liver tumor label
in the Liver Tumor Segmentation Task [1] (left) and the gallbladder label in
the BTCV dataset [11] (right). The initial prediction (light yellow) is
iteratively refined (shades of yellow), converging towards an improved
segmentation (brown). The predictions move closer to the ground truth label
(blue) when increasing the number of cycles. Intermediate cycles omitted for
visual clarity.

###  4.5 Memory and Run Time

Figure 4: Memory consumption during training on the BTCV dataset. We show the
memory consumption for the proposed RecycleNet, as well as the Recurrent U-Net
implemented in the nnU-Net framework. The memory consumption during training
is displayed w.r.t. the number of iterations (recycling cycles or recurrent
timesteps). While the proposed RecycleNet (blue) does come at the cost of a
memory surplus, this surplus stays constant irrespective of the number of
recycling cycles. Due to the recurrent nature of the Recurrent U-Net [18]
(light brown), the memory costs quickly become prohibitively large when
increasing the number of recurrent timesteps during training.

In this section, we analyze the memory and training epoch time surplus of
RecycleNet and the Recurrent U-Net baseline [18]. Since we implemented both in
the nnU-Net framework [6], a fair comparison with an identical base network
architecture, epoch definition, and underlying hardware can be ensured.

#### Training memory:

Figure 4 shows the memory consumption during training on the BTCV dataset
w.r.t. the number of iterations. We notice a steep increase in memory
consumption when training a Recurrent U-Net on medical image segmentation
tasks. This increase stems from two sources: the additional recurrent module
employed in the bottleneck of the network architecture and the recurrent
training with backpropagation through time. In the context of 3D medical image
segmentation, where the network architectures rely on large 3D inputs which
lead to a large number of feature activations stored in memory for
backpropagation, the latter contributes significantly to the total memory
consumption. This leads to more than doubled memory requirements compared to a
standard nnU-Net on this task. From this, we conclude that a recurrent U-Net
does not efficiently scale to a larger number of recurrent time steps during
training.
The proposed RecycleNet also shows a memory surplus during training, but due
to the formulation described in Algorithm 1, the memory demands saturate and
don’t grow with recycling cycles larger than two. In the experiments on the
BTCV dataset, we see an increase in memory consumption of roughly 30%. This
surplus varies depending on the dataset at hand. We notice that future
implementations should in principle be able to further reduce this memory
surplus.

#### Training epoch time:

Considering the average training epoch times, we again report numbers on the
BTCV dataset as a benchmark. Due to the 3D nature of medical image
segmentation tasks, we note that data loading and augmentation can often
become bottlenecks during training. However, the nnU-Net framework, along with
current hardware enables alleviating this bottleneck, thus allowing for an
unbiased measure of realistic epoch time changes stemming from the compared
methods. While a typical nnU-Net epoch (defined as 250 mini-batch updates with
2 samples each) on the BTCV dataset takes 74 seconds, the Recurrent U-Net
trained using 2 recurrent timesteps requires 165 seconds on average for each
epoch. This is more than a 100% increase in training time.

The proposed RecycleNet, however, comes with a much lower training time
increase. While the later stage of the proposed recycling schedule leads to an
epoch time increase of roughly 30% during training (up to 96 seconds compared
to nnU-Net’s 74 seconds), this epoch time increase is non-existent during the
first 200 epochs, where the recycling schedule only uses 1 cycle. Between
epochs 200 and 400 of the proposed schedule, we measure an epoch time increase
of roughly 15%. When training for 1000 epochs, this leads to a total training
time increase of about 20% with an average epoch time of 89.4 seconds.

In terms of memory and training time surplus, we see a non-negligible increase
for the proposed RecycleNet. This increase, however, is much lower compared to
the Recurrent U-Net baseline, rendering the RecycleNet much more accessible in
the context of 3D medical image segmentation. We note that the inference time
is also increased, since for each sample, $N_{c}$ forward passes have to be
computed. However, in the context of time-consuming preprocessing and
resampling, these additional forward passes only result in a minor inference
time increase (roughly by a factor of 2 when increasing the number of cycles
by a factor of 7).

###  4.6 Ablation of Training Schedules

Figure 5: Ablation of different recycling training schedules on the BTCV
dataset. The difference in DSC (compared with nnU-Net) is shown w.r.t. the
number of cycles used during inference. A static training schedule of 2 cycles
is shown in blue, while an increasing training schedule (without sampling) is
shown in green. The proposed schedule (see Section 3.1), is marked and
displayed in red.

In this section, we investigate how different recycling training schedules
affect the performance in the case of multi-organ segmentation on the BTCV
dataset. Figure 5 shows differences in DSC score when compared with nnU-Net
[6] in a 5-fold cross-validation. We show that while even using a fixed number
of cycles during training (blue) can increase segmentation performance w.r.t.
the baseline, this approach does not generalize well to a larger number of
cycles during inference. Incrementally increasing the number of cycles over
the course of the training process (green) does lead to a similar convergence
property as discussed in Section 4.4 when increasing the number of inference
cycles beyond the maximum training cycles. However, due to the missing
stochasticity in the number of cycles seen during training, the network
forgets how to handle lower recycling numbers and performs poorly in those
cases during inference. To ensure the observed convergence property described
in Section 4.4, both the sampling and the incremental schedule components are
important to achieve a reliable and strong prediction accuracy out of the box
when confronted with unseen segmentation tasks. This quality is crucial for a
wide applicability in the context of medical image segmentation.

##  5 Conclusion

In this work, we propose a novel method for instilling iterative decision
refinement capabilities into neural networks in a process we call latent
feature recycling. This approach relies on minimal assumptions w.r.t. the
network architecture and naturally leads to performance improvements, showing
a convergence towards refined decisions. We demonstrate these capabilities
using medical image segmentation as a showcase. In this context, latent
feature recycling can improve on even strong models with a simple and robust
recycling schedule. We observe that this schedule leads to iterative
refinement capabilities, allowing to trade inference time for improved
performance. These refinement capabilities even extend to a larger number of
cycles than seen during training. We leave it for future work to explore the
limits of this phenomenon for applications where even the smallest
improvements in performance are worthwhile. As this approach is not limited to
U-Nets or medical image segmentation, we expect further work to leverage
latent feature recycling on a variety of neural network architectures and
tasks. We note, however, that the proposed method comes with an additional
computational cost through increased numbers of forward passes during training
and inference. Compared with other iterative refinement approaches, we observe
a much lower memory and training time surplus, making the RecycleNet a
promising candidate e.g. in safety-critical applications where the additional
memory and time requirements are not the limiting factor. However, this
increased inference time potentially limits the application in time-critical
applications. The proposed method also bears a resemblance to latent diffusion
models [13], in that an iterative refinement process takes place in latent
space, leading to a convergent refinement property during inference. Contrary
to latent diffusion models, latent feature recycling covers multiple
refinement steps also during training, whereas latent diffusion models don’t
allow multistep feedback during training. We leave it for future work to
explore whether such multistep feedback is also beneficial in the denoising
diffusion process.

#### Acknowledgment:

This work was in part supported by the Helmholtz Association under the joint
research school HIDSS4Health - Helmholtz Information and Data Science School
for Health.
Part of this work was funded by Helmholtz Imaging, a platform of the Helmholtz
Incubator on Information and Data Science.

## References

  * [1] Michela Antonelli, Annika Reinke, Spyridon Bakas, Keyvan Farahani, AnnetteKopp-Schneider, Bennett A Landman, Geert Litjens, Bjoern Menze, Olaf Ronneberger, Ronald M Summers, Bram van Ginneken, Michel Bilello, Patrick Bilic, Patrick F Christ, Richard K G Do, Marc J Gollub, Stephan H Heckers, Henkjan Huisman, William R Jarnagin, Maureen K McHugo, Sandy Napel, Jennifer S Goli Pernicka, Kawal Rhode, Catalina Tobon-Gomez, Eugene Vorontsov, Henkjan Huisman, James A Meakin, Sebastien Ourselin, Manuel Wiesenfarth, Pablo Arbelaez, Byeonguk Bae, Sihong Chen, Laura Daza, Jianjiang Feng, Baochun He, Fabian Isensee, Yuanfeng Ji, Fucang Jia, Namkug Kim, Ildoo Kim, Dorit Merhof, Akshay Pai, Beomhee Park, Mathias Perslev, Ramin Rezaiifar, Oliver Rippel, Ignacio Sarasua, Wei Shen, Jaemin Son, Christian Wachinger, Liansheng Wang, Yan Wang, Yingda Xia, Daguang Xu, Zhanwei Xu, Yefeng Zheng, Amber L Simpson, Lena Maier-Hein, and M Jorge Cardoso.  The Medical Segmentation Decathlon.  Nature Communications, 13(1):4128, 2022.
  * [2] Andrea Banino, Jan Balaguer, and Charles Blundell.  Pondernet: Learning to ponder.  arXiv preprint arXiv:2107.05407, 2021.
  * [3] Joao Carreira, Pulkit Agrawal, Katerina Fragkiadaki, and Jitendra Malik.  Human pose estimation with iterative error feedback.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4733–4742, 2016.
  * [4] Wen Gu and Yi Xu.  Retinal vessel segmentation via adversarial learning and iterative refinement.  Journal of Shanghai Jiaotong University (Science), pages 1–8, 2022.
  * [5] Nicholas Heller, Niranjan Sathianathen, Arveen Kalapara, Edward Walczak, Keenan Moore, Heather Kaluzniak, Joel Rosenberg, Paul Blake, Zachary Rengel, Makinna Oestreich, Joshua Dean, Michael Tradewell, Aneri Shah, Resha Tejpaul, Zachary Edgerton, Matthew Peterson, Shaneabbas Raza, Subodh Regmi, Nikolaos Papanikolopoulos, and Christopher Weight.  The KiTS19 Challenge Data: 300 Kidney Tumor Cases with Clinical Context, CT Semantic Segmentations, and Surgical Outcomes.  arXiv, 2019.
  * [6] Fabian Isensee, Paul F Jaeger, Simon AA Kohl, Jens Petersen, and Klaus H Maier-Hein.  nnu-net: a self-configuring method for deep learning-based biomedical image segmentation.  Nature methods, 18(2):203–211, 2021.
  * [7] Fabian Isensee, Constantin Ulrich, Tassilo Wald, and Klaus H Maier-Hein.  Extending nnu-net is all you need.  arXiv e-prints, pages arXiv–2208, 2022.
  * [8] Yuanfeng Ji, Haotian Bai, Jie Yang, Chongjian Ge, Ye Zhu, Ruimao Zhang, Zhen Li, Lingyan Zhang, Wanling Ma, Xiang Wan, and Ping Luo.  AMOS: A Large-Scale Abdominal Multi-Organ Benchmark for Versatile Medical Image Segmentation.  arXiv, 2022.
  * [9] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al.  Highly accurate protein structure prediction with alphafold.  Nature, 596(7873):583–589, 2021.
  * [10] A. Emre Kavur, N. Sinem Gezer, Mustafa Barış, Sinem Aslan, Pierre-Henri Conze, Vladimir Groza, Duc Duy Pham, Soumick Chatterjee, Philipp Ernst, Savaş Özkan, Bora Baydar, Dmitry Lachinov, Shuo Han, Josef Pauli, Fabian Isensee, Matthias Perkonigg, Rachana Sathish, Ronnie Rajan, Debdoot Sheet, Gurbandurdy Dovletov, Oliver Speck, Andreas Nürnberger, Klaus H. Maier-Hein, Gözde Bozdağı Akar, Gözde Ünal, Oğuz Dicle, and M. Alper Selver.  CHAOS Challenge - combined (CT-MR) healthy abdominal organ segmentation.  Medical Image Analysis, 69:101950, 2021.
  * [11] Bennett Landman, Zhoubing Xu, J Igelsias, Martin Styner, T Langerak, and Arno Klein.  Miccai multi-atlas labeling beyond the cranial vault–workshop and challenge.  In Proc. MICCAI Multi-Atlas Labeling Beyond Cranial Vault—Workshop Challenge, volume 5, page 12, 2015.
  * [12] Jun Li, Xiaozhu Lin, Hui Che, Hao Li, and Xiaohua Qian.  Pancreas segmentation with probabilistic map guided bi-directional recurrent unet.  Physics in Medicine & Biology, 66(11):115010, 2021.
  * [13] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684–10695, 2022.
  * [14] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In International Conference on Medical image computing and computer-assisted intervention, pages 234–241. Springer, 2015.
  * [15] Hao Tang, Xingwei Liu, Shanlin Sun, Xiangyi Yan, and Xiaohui Xie.  Recurrent mask refinement for few-shot medical image segmentation.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 3918–3928, 2021.
  * [16] Zhuowen Tu and Xiang Bai.  Auto-context and its application to high-level vision tasks and 3d brain image segmentation.  IEEE transactions on pattern analysis and machine intelligence, 32(10):1744–1757, 2009.
  * [17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [18] Wei Wang, Kaicheng Yu, Joachim Hugonot, Pascal Fua, and Mathieu Salzmann.  Recurrent u-net for resource-constrained segmentation.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 2142–2151, 2019.
  * [19] Amir R Zamir, Te-Lin Wu, Lin Sun, William B Shen, Bertram E Shi, Jitendra Malik, and Silvio Savarese.  Feedback networks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1308–1317, 2017.