# Learning Transductions to Test Systematic Compositionality

Josef Valvoda${}^{\text{{@}}}$ Naomi Saphra${}^{\text{{B}}}$ Jonathan
Rawski${}^{\text{{C}}}$
Ryan Cotterell${}^{\text{{D}}}$ Adina Williams${}^{\text{{E}}}$
${}^{\text{{@}}}$University of Cambridge ${}^{\text{{B}}}$New York University
${}^{\text{{C}}}$San Jose State University ${}^{\text{{D}}}$ETH Zürich
${}^{\text{{E}}}$FAIR
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS>

###### Abstract

Recombining known primitive concepts into larger novel combinations is a
quintessentially human cognitive capability. Whether large neural models in
NLP acquire this ability while learning from data is an open question. In this
paper, we look at this problem from the perspective of formal languages. We
use deterministic finite-state transducers to make an unbounded number of
datasets with controllable properties governing compositionality. By randomly
sampling over many transducers, we explore which of their properties (number
of states, alphabet size, number of transitions etc.) contribute to
learnability of a compositional relation by a neural network. In general, we
find that the models either learn the relations completely or not at all. The
key is transition coverage, setting a soft learnability limit at $400$
examples per transition.

##  1 Introduction

Systematic compositionality is a hallmark of human language Hockett (1959);
Chomsky (1957); Montague (1970); Partee (1995). It is arguably a requirement
for any model to call itself a model of language or to achieve human-like
natural language understanding. For instance, English speakers reading the
Jabberwocky Carroll (1871) comprehend that the adjective slithy composes with
the plural noun toves to form the noun phrase slithy toves, without having a
clear understanding of what either the word slithy or toves—let alone their
composition—means. This ability allows humans to describe a complex
environment using a limited number of primitives (in the case of language,
words or morphemes). In cognitive science, whether neural networks' possess
the ability to acquire such behavior has been debated for over 30 years Fodor
and Pylyshyn (1988); Marcus (1998).

In recent work, researchers have analyzed sequence-to-sequence model
performance on small, controlled datasets in order to draw conclusions about
the inherent limitations of neural models to acquire systematic
compositionality (Lake and Baroni, 2018; Hewitt et al., 2020; Hupkes et al.,
2020; Dankers et al., 2021). For instance, Lake and Baroni (2018) cast doubt
on neural models' ability to learn systematic behavior encoded in their toy
SCAN dataset, but shortly thereafter, Bastings et al. (2018) demonstrate that
an out-of-the-box sequence-to-sequence model could indeed master the task.

Rather than introducing another hand-crafted dataset, we take lessons from the
field of grammatical inference, which studies the general behavior of learning
algorithms on entire _classes_ of formal languages de la Higuera (2010);
Rawski and Heinz (2019). We follow Montague's assertion that there is no
important theoretical difference between natural and artificial languages
Montague (1970). This leads us to believe that testing neural sequence-to-
sequence models' ability to learn systematic compositionality will bear more
fruit if it focuses on evaluating a large number of automatically generated
artificial languages sampled from particular classes of functions, as is
standard practice at grammatical inference competitions Balle et al. (2017).
In this paper, we demonstrate this by studying the class of string functions
encoded by subsequential finite-state transducers (SFST), a restricted class
of general finite-state transducers, which encode rational relations Mohri
(1997). We sample arbitrary SFSTs to generate many different string-to-string
datasets, and evaluate the behavior of neural sequence-to-sequence models when
learning them. By controlling the properties of the SFSTs we sample from, we
are able to make precise statements about the learnability of systematic
phenomena.

Empirically, we find that neural sequence-to-sequence models are capable of
learning SFSTs from finite data. Moreover, we observe an interesting tendency
for neural models to either generalize correctly or to fail, with little
middle ground. Our analysis suggests that an explanation for this behavior
lies in whether or not the training data has sufficient coverage, i.e., every
transition in a given transducer is crossed in a minimum number of training
examples ($\approx 400$ on average in our experiments). We then turn to an
existing hand-crafted dataset, namely the SCAN language, but approach it
through the lens of the SFST that encodes it. Via a comparison with SCAN, we
are able to identify that repeated sub-graphs in the transducer can have a
positive effect on its learnability, relaxing the transition coverage
requirement and thereby enabling learnability of larger SFSTs. Our results
suggest that focusing on formal language theory, which allows for a controlled
setting, could help future research on evaluating systematic compositionality.

10 19 16 28 19 7 27 20 | $\Rightarrow$ | 33 31 31 30 57 56 44
---|---|---
11 10 18 2 5 19 7 25 4 1 4 | $\Rightarrow$ | 48 33 38 45 40 30 57 44 32
1 18 23 2 2 22 25 3 18 11 23 1 0 1 | $\Rightarrow$ | 53 39 45 37 46 44 36 53 48 39 38 41
1 21 20 11 19 2 5 11 11 10 13 8 | $\Rightarrow$ | 40 43 37 30 37 32 48 48 33 43 40
Table 1: Example of our inputs on the left and outputs on the right. Using
numbers instead of words ensures the relationship between the input and output
sequences requires knowing the underlining SFST and can't be a function of
other factors, such as word similarity.

##  2 Learning Finite-State Transductions

We argue that determining whether neural networks exhibit systematic
compositionality is fundamentally equivalent to asking whether neural networks
are capable of learning rule-based behavior.

###  2.1 Montague's Compositionality

Montague famously gave a mathematical definition of what it means for a
function to be compositional—specifically, he asserted that a function is
compositional if it is a homomorphism Andreas (2019). In the context of
string-to-string maps, a function is a homomorphism if it preserves
concatenation. To make this notion more formal, let $\Sigma$ and $\Gamma$, be
a pair of alphabets. Making use of string concatenation $\circ$, we can
construct monoids ($\Sigma^{*},\circ)$ and ($\Gamma^{*},\circ)$ where the
identity element is the empty string $\varepsilon$. In this context, a
function $f:\Sigma^{*}\rightarrow\Gamma^{*}$ is a homomorphism if it obeys the
following two conditions:

  1. 1.

$f(x\circ y)=f(x)\circ f(y)$ for all $x,y$ in $\Sigma^{*}$

  2. 2.

$f(\varepsilon)=\varepsilon$

author=ryan,color=violet!40,size=,fancyline,caption=,]You might want to bold
strings

As we will see, deterministic transducers naturally encode a homomorphism
between two alphabets. Therefore, we test compositional ability of a neural
network by training it to learn string-to-string maps encoded by a transducer.

###  2.2 Why Learn Finte-State Transductions?

Our study focuses on learning a particular kind of transductions.
Specifically, we focus on restricted classes of regular relations, which are
computed by finite-state transducers. We believe this is a natural starting
point, since this class of grammars is mathematically well-studied, has
provable learning guarantees, and has a long use history in linguistics and
NLP Mohri (1997).

Finite-state transducers also encompass most previous work on systematic
compositionality: many datasets, e.g., SCAN Lake and Baroni (2018) and gSCAN
Ruis et al. (2020), describe _finite_ string relations and are, therefore,
finite-state by definition. These handcrafted datasets have many advantages,
like easy interpretability and domain specificity, since they directly encode
particular relevant relationships, such as movement over a grid or specific
linguistic phenomena. However, this realism pays the price of diminished
robustness of any findings over such datasets Rogers and Pullum (2011). By
removing the ability to simply adjust properties of the underlying function
class, and the transducers which compute it, one loses the possibility to
experiment more robustly over a whole class of functions, rather than just one
instantiation of it.

Rather than manually designing individual datasets ourselves, we generate
unboundedly many new datasets via randomly sampled SFSTs. This offers a
principled view of the problem of learning artificial languages possessing
certain properties, by simply varying properties of the class of transducers
that generate them. Furthermore, as we will see in § 5, one may view existing
compositionality tasks as SFSTs which enables a deeper understanding of
modeling results. Both specific artificial languages, such as the
compositionality datasets mentioned above, and those randomly sampled from a
particular function class such as the work presented in this paper, are worth
studying. However, our approach has been missing from the systematic
compositionality discourse.

###  2.3 Subsequential Finite-State Transducers

A subsequential finite-state transducer is a transducer that is deterministic
with respect to the input tape. By construction, subsequential transducers are
functional, i.e., the string-to-string relations they encode are functions. An
alphabet $\Sigma$ is a finite, non-empty set of symbols; let $\Sigma^{*}$ be
all strings over $\Sigma$, including $\lambda$, the empty string. A
subsequential finite-state transducer (SFST) for an input alphabet $\Sigma$
and an input–output alphabet $\Gamma$ is a tuple $\mathcal{T}=\left\langle
Q,q_{0},Q_{f},\delta,o,\omega\right\rangle,$ where $Q$ is the set of states,
$q_{0}\in Q$ is the (unique) start state, $Q_{f}\subseteq Q$ is the set of
final states, $\delta:Q\times\Sigma\rightarrow Q$ is the transition function,
$o:Q\times\Sigma\rightarrow\Gamma^{*}$ is the output function, and
$\omega:Q_{f}\rightarrow\Gamma^{*}$ is the final function. We define the
reflexive, transitive closure of $\delta^{*}$ and $o^{*}$ as
$\delta^{*}:Q\times\Sigma^{*}\rightarrow Q$ and
$o^{*}:Q\times\Sigma^{*}\rightarrow\Gamma^{*}$.

The semantics of a SFST is a transduction $t(\mathcal{T})$ defined as follows;
let $\hat{t}=$ $t(\mathcal{T}).$ For $w\in\Sigma^{*},\hat{t}(w)=uv$ where
$o\left(q_{0},w\right)=u,$ and $\omega\left(q_{f}\right)=v$ if
$\delta^{*}\left(q_{0},w\right)=q_{f}$ for some $q_{f}\in Q_{f};\hat{t}(w)$ is
undefined otherwise. The class of subsequential functions are those
describable with SFSTs, among other algebraic and logical characterisations
(Bhaskar et al., 2020; Oncina et al., 1993). They are a subclass of the
Regular relations, and a superclass of the finite relations.

Informally, for subsequential functions, any of the infinitely many possible
input strings is classified as belonging to exactly one of finitely many
regular languages. For any input element $x$ of that string, the output $u$
only depends on $x$ and the regular stringset to which its preceding input
string belongs. The key compositionality property of subsequential functions
is finite lookahead; the transducer can only remember finitely much
information about the input for any output symbol.

#### SFSTs as Homomorphisms.

Now we discuss the exact sense, in which SFSTs are homomorphisms. Given a path
in a SFST from $q\cdots q^{\prime\prime}$, for any state $q^{\prime}$ on this
path, we have that $o(q\cdots q^{\prime\prime})=o(q\cdots q^{\prime})\circ
o(q^{\prime}\cdots q^{\prime\prime})$, i.e., we can simply take the
concatenation of the yield of the output path of the two subpaths. This will
give us the same string as if we were to take the yield of the full path. By
definition, an SFST therefore encodes a
homomorphism.author=ryan,color=violet!40,size=,fancyline,caption=,]We
overloaded $o$ to take a yield.

###  2.4 Generating Random SFSTs

We follow these steps in generating the SFSTs used to create our datasets. We
first generate random directed graphs from the following stochastic process.
For a given a set of states $Q$ with $|Q|=N$ and an input alphabet $\Sigma$,
we sample a matrix $\mathbf{B}^{(\sigma)}\in\mathbb{B}^{N\times N}$ for each
$\sigma\in\Sigma$ where $\mathbb{B}=\\{0,1\\}$. During sampling, we
additionally enforce the constraint that there be _at most_ one non-zero entry
in every row vector $\mathbf{b}^{(\sigma)}_{i}$. This constraint ensures that
the resulting SFST is subsequential (or input deterministic) by construction.
We sample uniform at random over this space of matrices that satisfy this
constraint. In terms of the semantics, if the entry $b^{(\sigma)}_{ij}=1$, it
means our generated SFST has a transition from state
$q_{i}\overset{\sigma}{\longrightarrow}q_{j}$ with input symbol $\sigma$.
Then, for every transition $q_{i}\overset{\sigma}{\longrightarrow}q_{j}$ in
our generated SFST, we sample its output symbol $\gamma$ from a uniform
distribution over the output alphabet $\Gamma$. This results in a transition
$q_{i}\overset{\sigma/\gamma}{\longrightarrow}q_{j}$. Finally, to get a
canonical representation for particular SFSTs, we minimize it Choffrut (2003).
Minimization also ensures that we have created thousands of _different and
unique_ SFSTs—i.e., that we are not doubly sampling.

A state in a finite-state transducer is accessible if there exists a path from
the initial state to that state in the machine. A state is co-accessible if
there is a path from it to a final state. To create interesting SFSTs, we
additionally enforce both accessibility and co-accessibility. Testing for
accessibility and co-accessiblity is easily achieved with a depth-first
search. In order to ensure our sampled SFSTs satisfy these two constraints, we
reject those samples that violate this constraint.

##  3 Experimental Methods

We evaluate whether neural sequence-to-sequence models can learn our generated
SFSTs, and how they compare to the traditional symbolic learners.

Figure 1: Average minimum frequency of a transition appearing in training data
for learning a SFST is $400$.

###  3.1 Data

Here, we discuss the creation of the datasets we use to investigate the
learnability of the regular relations. Our input alphabet consisted of a
finite set of numbers; we set its size to a constant of $10$. Half the time,
we also add a special empty emission token $\lambda$ (in our Python
implementation, $-1$) to the output alphabets to see if it affects the neural
models learnability of the SFST.

We first generate $100$ unique datasets using SFSTs where the number of states
ranges from $10$ to $100$ in increments of $10$. We further generate $100$
unique datasets using SFSTs with states ranging from $21$ to $39$. In total,
we experiment with $2000$ unique datasets. However, using the sampling method,
one can generate an arbitrary number of SFSTs and, concurrently, datasets, to
various specifications as desired.

To generate the input–output pairs for our experiments, we randomly walk
through the SFST. We define a parameter for the chance that a given state will
emit a final symbol, and set it to $10\%$. This provides some diversity to the
sequence length within computational limits of the GPU. Following Lake and
Baroni (2018), the maximum number of steps is set to $50$. We continue to walk
the SFST until we collect $20,000$ unique input-output pairs. We further
double the dataset size to $40,000$ for SFSTs with $21$ to $39$ states, since
we observe the accuracy drop-off in this region. All datasets are randomly
split $80$–$20$ into training and test sets.

All particular states and transitions (i.e., single connections between
precisely two states) are seen by the model during training. This means that
coverage of both the states and transitions for all datasets is $100\%$.
Recall though, that all sequences are unique, because the path through the
states via the transitions is determined by random walk. Table 1 shows some
example transductions.

Figure 2: As the number of states in SFST increases, accuracy drops to nearly
zero percent between 20 and 40 states.

###  3.2 Neural Sequence-to-Sequence Models

We train neural encoder-decoder models in the style of Bahdanau et al. (2015),
using bi-directional LSTMs with a $\tanh$ activation function and a global
attention mechanism. Attention allows the decoder to selectively ``attend'' to
the encoder's hidden states by learning a set of weights. These weights map
the decoder's current state to another set of weights for timesteps in the
input. Attention then concatenates the current decoder hidden state, i.e. the
weighted combination of all encoder hidden states, to yield a new current
decoder state. This way, the decoder can pull information directly from the
encoder by learning an alignment between the output and input representations.

We use the following hyperparameters; $200$ dimensional encoder and decoder
with $2$ layers, maximum gradient clipping normalization of $5$, dropout of
$0.5$ and batch size $64$. Our alphabet tokens are embedded as $100$
dimensional vectors. Just like Lake and Baroni (2018), we run the model for
$100,000$ steps. Unlike them though, we rely on the neutral transducer code
base from Wu and Cotterell (2019) which uses the Adam optimizer Kingma and Ba
(2015) with the default learning rate of $0.001$. When we report results for
larger LSTM models, this is the effect of changing only the dimensions of the
encoder and decoder, both to $300$ dimensions, and keeping all the other
hyperparameters the same as the smaller models.

###  3.3 OSTIA

The onward subsequential transducer inference algorithm (OSTIA; Oncina et al.
1993) inductively learns the class of subsequential functions from positive
presentations of input-output strings, OSTIA works by first building a prefix-
tree transducer of the training data, which is then transformed through a
series of state-merging operations into the SFST encoding the function the
data is drawn from. If a characteristic sample is contained in the learning
data, OSTIA finds a correct transducer in polynomial (cubic) time. Since OSTIA
is designed specifically with SFST in mind, we wanted to see how it compares
to the neural models above and set a baseline with it. Unfortunately, OSTIA's
cubic nature makes learning from large datasets like ours very time intensive.
To speed up the process, we limit the number of samples we provide to OSTIA to
1000, only 10% to 5% of what the neural transducers obtain. This keeps OSTIA's
run time roughly equivalent to its neural counterparts.

##  4 Results

Training over many datasets enables us to make the following observations
about how different SFST properties affect the learnability of a compositional
language.

#### Minimum Transition Coverage.

Figure 1 reveals a threshold on the number of samples per transition required
to comfortably learn the transducer: It is about $400$ samples. This may seem
unsurprising given neural networks' ``notorious thirst'' for data (Lake and
Baroni, 2018). In the vast majority of cases, SFSTs that do not meet this
threshold learn nothing. In the few cases where they manage to have non-zero
accuracy, we observe that the earlier transitions in the SFST have attained
reasonable coverage, leading us to believe that the coverage is the linchpin
of learning SFSTs.

#### Number of States.

Models' ability to learn datasets depends on the number of states in the SFST
(see Figure 2). There is a sharp drop-off around $20$ states, decreasing
exponentially as states grow linearly, until $40$ states. From there on,
models only achieve low accuracy. Our findings offer an explanation of why
training on shorter sequences impedes the learnability of a dataset. For
example, in Lake and Baroni (2018)'s and Ruis et al. (2020)'s length based
experiments, where the task is to generalize to commands with longer action
sequence, we hypothesize that shorter sequences correlate with lower coverage.
We find that, indeed, splitting the datasets along length drops the coverage
compared to a random split. This suggests that splitting datasets by length of
sequences has little to do with systematicity, but a lot to do with coverage.
We discuss this in further detail in § 5.

Figure 3: As SFST state size increases, average accuracy (blue line)
decreases. Increasing the neural model's size improves accuracy (purple line).
Empty string $\lambda$ in output alphabet further reduces accuracy (yellow
line). Similar behavior appears when plotting number of transitions instead of
states.

#### Number of LSTM Hidden Units.

We find coverage is only a soft limit. When we increase the memory size of the
LSTM from $200$ to $300$ dimensions, the models can learn even those datasets
that seemed previously unlearnable (see Figure 3, where the purple line is the
average accuracy of $300$ dimensional LSTM). Except the higher state coverage
threshold, the trend these models follow is the same as for their lower
dimensional counterparts.

#### Empty Transition.

We explore another way of making a more difficult finite function: we allow
the SFSTs to contain empty transitions $\lambda$ and test whether languages
generated by SFSTs with empty transitions are harder by measuring the average
accuracy over SFSTs with and without them. We find a 12% performance drop on
the test set: SFSTs without empty transitions achieve 47% average accuracy
versus 35% for SFSTs with an empty string. We therefore conclude that the
empty transition does make inducing a SFST empirically harder for the models
(see Figure 3, where the blue line is average accuracy without empty strings,
and the yellow line is average accuracy with them).

#### OSTIA is Slow.

In terms of wallclock time, we find that a open-source implementation of
OSTIA111github.com/alenaks/OSTIA does not scale to dataset sizes above $1000$,
for which it takes roughly $5$ hours on Nvidia Tesla V$100$ GPU to complete a
run on one dataset, about the same time it takes to train our neural models.
With 1000 samples, OSTIA fails to get accuracy above 1%. In contrast, the
neural sequence-to-sequence models reach almost 100%. On the one hand,
reducing the size of the training dataset disadvantages OSTIA. On the other
hand, providing OSTIA with the full $20,000$ samples is unworkable, because it
will run without converging for over $3$ days on a single dataset. Given
sufficient time and the function's characteristic sample, OSTIA _is_
guaranteed to learn the SFSTs perfectly, but in practice, it is too slow to
use.

0$1$$2$$3$$10$$4$$5$$6$$7$$8$$9$jump/$\lambda$right/$\lambda$left/$\lambda$$\lambda$/RTURNJUMP$\lambda$/LTURNJUMP$\lambda$/JUMParound/$\lambda$opposite/$\lambda$right/$\lambda$left/$\lambda$$\lambda$/RTURN(x4)JUMP$\lambda$/LTURN(x4)JUMPright/$\lambda$left/$\lambda$$\lambda$/RTURN(x2)JUMP$\lambda$/LTURN(x2)JUMP
Figure 4: Core SFST block of SCAN encoding a primitive like jump. 0 is the
start state and 10 is the accept state. Moving through states generates simple
_jump_ commands in SCAN dataset.

##  5 Discussion: The Learnability of SCAN

The empirical learnability of SFSTs appears to hinge on the transition
coverage. In this section we discuss the SCAN language and the experiments
conducted on it. First, we argue the transition coverage can explain some of
the results observed over SCAN. Second, we cast SCAN as a SFST to reason why
it is learnable despite the parameters of the SFST that encodes it would
suggest otherwise.

Lake and Baroni (2018) find that training models on a random split of the SCAN
dataset dramatically outperforms models trained only on short commands, i.e.
the lengthwise split. While the best random split model achieves $100$%
accuracy, the best model trained on the lengthwise split achieves only $21$%
accuracy. Furthermore, the accuracy drops from 80% to _zero_ between the
$24^{th}$ and $26^{th}$ token of the models output. A situation similar to
what we observe in Fig. 3.

Lake and Baroni (2018) explore two possible sources of this error,
specifically decoder search failures and early terminations. However, they
admit that neither seems to fully explain the performance chasm. In light of
our work, the likely reason for the performance drop-off comes down to the
lack of sufficient transition coverage necessary to learn the particular
transitions enabling longer commands. Since the longer action sequences
require mastering repetition inducing commands: twice, thrice, and, and after,
we hypothesise that a providing sufficient coverage of these commands would
dissolve the difference in accuracy between the lenghtwise and random splits.
In other words, we suspect the issue doesn't reflect a neural network's
ability to learn systematic compositionality, but rather the lower coverage of
SCAN commands which appear at the end of an input sequence.

We extend our analysis by encoding SCAN as a SFST and find that it far exceeds
our maximum SFST size of $100$ states—we estimate the full SFST for SCAN has
$7,728$ states. With a finite dataset size of $20,000$ input–output pairs,
SCAN should not be possible to learn with a high accuracy. However, unlike
datasets of a similar size, it turns out that SCAN is nearly perfectly
learnable in our (random-split) experiments. This naturally raises the
question of what makes SCAN special. We find the answer within the structure
of the SFST that encodes it. SCAN's input alphabet includes primitive commands
such as walk, turn, run, turn left and turn right among others. Each primitive
command has a single accept state. Figure 4, shows a simplified transducer
encoding of the primitive jump. Although the SCAN SFST has over $7,000$
states, the majority of these states and the transitions between them are
repeated. Qualitatively, SCAN has $23$ states for encoding each primitive
command connected by $50$ transitions. While SCAN as a SFST would be very
large, its core logic is not so complex. Two of SCAN's design features in
particular explain this: reversal and repetition.

First, the transduction in SCAN reverses the order of commands so that jump
right turns into RTURN JUMP. Because of this reversal, the SFST encoding of
SCAN cannot have a single encoding for all SCAN primitives. Each must be
encoded separately, since the only transition in SCAN that can return anything
is the one leading to the final state. Recall Figure 4, which encodes the
basic logic of a single primitive such as jump, ignoring twice and thrice
commands. We need four of these blocks, each with $40$ states, to encode the
four SCAN primitives or just above what we find comfortably learnable in our
datasets.

Second, SCAN also has commands that allow for repetitions of sequences: twice,
thrice, and, and after, which additionally reverse the input strings for
output. Adding those primitives in will triple the number of states
transitioning to the end state. For example, from state $2$ in Figure 4, which
encodes jump right we will need two more states transitioning from it,
encoding jump right twice and jump right thrice. Repetition therefore expands
the number of states per SCAN primitive to $24$, putting it far above our
empirically learnable threshold.

This effect is exacerbated with the commands and and after, where for every
primitive in SCAN, each transition connecting to the end state—$10$ in Figure
4—will have the whole SFST itself appended to its end, thereby multiplying the
number of states each time. Since SCAN is learnable to a near $100$% accuracy,
we hypothesize that the LSTM must learn the repeated patterns between the
different primitives, allowing models to learn SCAN despite it exceeding the
ceiling of $40$ states.

We lend some support to this hypothesis by including these special properties
for random SFSTs in our datasets. We generate datasets with loops over the
$20$-state SFSTs, i.e., the largest feasibly learnable SFSTs. The loops add
twice and thrice commands to each dataset, increasing the number of states in
the SFST that generated them. We find no difference in learning these
augmented SFSTs, demonstrating that sub-graph structures in the SFSTs indeed
affect their learnability.

Therefore, we conclude that while SCAN's design contains many repeated blocks
(i.e., sub-graphs), the ability of our neural models to learn a SFST depends
on the complexity of the unique structure encoded by it, rather than these
repetitions.222Note that “blocks” are not visible to SFSTs. SFSTs have no way
to encode similarity between their parts. Viewing SCAN through the lens of
automata thus reveals the properties that enable the learnability of the
finite language it encodes.

However, the case of SCAN also reveals that to predict ease of learning
requires a concept of complexity that is not reflected in SFST state counts:
the duplicated sub-graph structures. We hope future work will use our method
and explore different types of automata to identify classes that robustly
encode learnability, bringing us closer to understanding how neural models
learn such functions.

##  6 Related Work

We build on two strains of research: the work on systematic compositionality
and on grammatical inference. This section contextualizes our work.

###  6.1 Compositionality Datasets

There is a growing number of artificial language datasets focused on
systematic compositionality. Lake and Baroni (2018) introduced a SCAN dataset,
made up of simple navigational text commands. The task is to translate the
command in the simple natural language into sequences of actions. One
successor to SCAN is the NACS dataset (Bastings et al., 2018), which is
comparable to SCAN, but instead of mapping multiple input signals to a single
duplicated output symbol (e.g., walk twice $\rightarrow$ WALK WALK), NACS does
the opposite (WALK WALK $\rightarrow$ walk twice). NACS is both theoretically
more difficult, and has been shown to be empirically more challenging (for
both GRU and RNNs with attention, Bastings et al. 2018). The same inversion
method applied to the transducers in our datasets may result in non-
deterministic (and thus non-subsequential) transductions, since deterministic
SFSTs are not closed under inverse (Mohri, 1997). Another successor is gSCAN
Ruis et al. (2020) which grounds SCAN-like commands in states of a grid world,
closer to Mikolov et al. (2016)'s grid world grounding for their agents. gSCAN
requires the agent to learn differences between sizes and colours of different
geometric shapes and interact with them, by moving them around the grid world.
Executing a gSCAN command is therefore much more difficult than to execute its
SCAN counterpart. As Ruis et al. (2020) asserts, the gSCAN dataset removes
artefacts in SCAN which are not central to the compositional generalization.
They find that models perform worse on gSCAN than on SCAN.

Like the work above, we investigate systematic compositionality, but unlike
the work above we generalise over the existing datasets by sampling from a
well-defined class of patterns.

###  6.2 Grammatical Inference.

Grammatical inference studies the general behavior of learning algorithms over
classes of functions generated by varieties of automata, rather than the
specific behavior of specific algorithms on specific functions. Our work
focuses on the learning of a restricted class of functions generated by a
correspondingly restricted class of finite-state transducers. This allows us
to synthesize our study of compositionality in neural models as rule-based
inference by neural models, which we can restrict in principled ways. Finite-
state machines generalise many techniques in NLP: probabilistic finite-state
automata, hidden Markov models, Markov chains, $n$-grams, probabilistic suffix
trees, deterministic stochastic probabilistic automata, weighted automata, and
other syntactic objects which generate distributions over sets of possible
infinite cardinality of strings, sequences, words, trees, and graphs (Vidal et
al., 2005). Many grammatical inference studies of neural networks test them on
samples drawn from some deterministic finite-state acceptor (DFA) (Cleeremans
et al., 1989), and some then use quantization, clustering, or spectral
algorithms to partition the RNN state space into an extractable DFA (Giles et
al., 1992; Eyraud and Ayache, 2020) (see Jacobsson 2005 for a review).

Other inference experiments using neural nets attempt to see if the regular
boundary can be empirically crossed. LSTMs (Hochreiter and Schmidhuber, 1997)
can perform dynamic counting and variably learn simple counter languages such
as some $k$-Dyck languages and $a^{n}b^{n}$ patterns (Weiss et al., 2018;
Suzgun et al., 2019; Bhattamishra et al., 2020; Hewitt et al., 2020), which
are generated by a finite-state machine with a counter on top (Schützenberger,
1962). In contrast, Avcu et al. (2017) show that LSTM and other RNN
architectures often fail to learn long-distance dependencies drawn from
simpler sub-regular language classes, even on large benchmarks (Mahalunkar and
Kelleher, 2019). Nelson et al. (2020) study the inference of sequence-to-
sequence networks, showing that RNN, LSTM, and GRU (Cho et al., 2014)
systematically fail to learn a wide range of regular string copying functions
generated from a family of two-way transducers, which characterize regular
string-to-string functions. When augmented with attention, they reliably learn
every function, and the attention history mirrors the derivations of the
corresponding two-way transducers. These independently productive strands of
work in compositionality and inference suggest that our work is a reasonable
starting point for future interactions.

##  7 Conclusion

We study if sequence-to-sequence models are capable of learning systematic
compositionality. To do this we sample arbitrary SFSTs and generate thousands
of unique string-to-string datasets. By knowing the structure of the
underlying SFSTs we can observe which of the underlying factors affect the
learnability of this class of languages. This allows us to make precise
statements about the learnability of systematic phenomena by a neural network.

We showed that neural models are indeed capable of learning generalized
systematic behaviour. Moreover, we find that neural models tend to either
generalize completely or fail miserably, with little middle ground. We pin
down the underlying cause as transition coverage. With this in mind, we can
now reason about possible cause for a pronounced performance difference
between random and lengthwise split observed on the SCAN dataset. We observe
that the cause here is likely not due to the lack of systematicity, but rather
the lack of transition coverage. We find that formalising a human designed
finite language as a SFST can reveal factors affecting its learnability which
would otherwise remain hidden. Specifically, repeating sub-graphs in its SFST
encoding can lead to learnability of SFSTs with a much higher number of
states, and a corresponding lower transition coverage, than one would expect
from an average dataset of similar number of states. Our formal language
approach therefore not only allows for novel insights into the constrains for
the neural models, but also a novel way of analysing the existing work.

## References

  * Andreas (2019) Jacob Andreas. 2019.  Measuring compositionality in representation learning.  In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_. OpenReview.net.
  * Avcu et al. (2017) Enes Avcu, Chihiro Shibata, and Jeffrey Heinz. 2017.  Subregular complexity and deep learning.  In _CLASP Papers in Computational Linguistics: Proceedings of the Conference on Logic and Machine Learning in Natural Language (LaML 2017), Gothenburg, 12 –13 June_ , pages 20–33.
  * Bahdanau et al. (2015) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2015.  Neural machine translation by jointly learning to align and translate.  In _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_.
  * Balle et al. (2017) Borja Balle, Rémi Eyraud, Franco M. Luque, Ariadna Quattoni, and Sicco Verwer. 2017.  Results of the sequence prediction challenge (SPiCe): a competition on learning the next symbol in a sequence.  In _Proceedings of The 13th International Conference on Grammatical Inference_ , volume 57 of _Proceedings of Machine Learning Research_ , pages 132–136, Delft, The Netherlands. PMLR.
  * Bastings et al. (2018) Jazmijn Bastings, Marco Baroni, Jason Weston, Kyunghyun Cho, and Douwe Kiela. 2018\.  Jump to better conclusions: SCAN both left and right.  _Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pages 47–55.
  * Bhaskar et al. (2020) Siddharth Bhaskar, Jane Chandlee, Adam Jardine, and Christopher Oakden. 2020.  Boolean monadic recursive schemes as a logical characterization of the subsequential functions.  In _Language and Automata Theory and Applications - LATA 2020_ , Lecture Notes in Computer Science, pages 157–169. Springer.
  * Bhattamishra et al. (2020) Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. 2020.  On the practical ability of recurrent neural networks to recognize hierarchical languages.  In _Proceedings of the 28th International Conference on Computational Linguistics_ , pages 1481–1494, Barcelona, Spain (Online). International Committee on Computational Linguistics.
  * Carroll (1871) Lewis Carroll. 1871.  _Through the Looking Glass_.  Macmillan London.
  * Cho et al. (2014) Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. 2014.  Learning phrase representations using RNN encoder–decoder for statistical machine translation.  In _Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 1724–1734, Doha, Qatar. Association for Computational Linguistics.
  * Choffrut (2003) Christian Choffrut. 2003.  Minimizing subsequential transducers: a survey.  _Theoretical Computer Science_ , 292(1):131 – 143.  Selected Papers in honor of Jean Berstel.
  * Chomsky (1957) Noam Chomsky. 1957.  _Syntactic Structures_.  Walter de Gruyter.
  * Cleeremans et al. (1989) Axel Cleeremans, David Servan-Schreiber, and James L. McClelland. 1989.  Finite state automata and simple recurrent networks.  _Neural Computation_ , 1(3):372–381.
  * Dankers et al. (2021) Verna Dankers, Elia Bruni, and Dieuwke Hupkes. 2021.  The paradox of the compositionality of natural language: a neural machine translation case study.  _arXiv preprint arXiv:2108.05885_.
  * Eyraud and Ayache (2020) Remi Eyraud and Stephane Ayache. 2020.  Distillation of weighted automata from recurrent neural networks using a spectral approach.  _arXiv preprint arXiv:2009.13101_.
  * Fodor and Pylyshyn (1988) Jerry A. Fodor and Zenon W. Pylyshyn. 1988.  Connectionism and cognitive architecture: A critical analysis.  _Cognition_ , 28(1-2):3–71.
  * Giles et al. (1992) C. Lee Giles, Clifford B. Miller, Dong Chen, Hsing-Hen Chen, Guo-Zheng Sun, and Yee-Chun Lee. 1992.  Learning and extracting finite state automata with second-order recurrent neural networks.  _Neural Computation_ , 4(3):393–405.
  * Hewitt et al. (2020) John Hewitt, Michael Hahn, Surya Ganguli, Percy Liang, and Christopher D. Manning. 2020.  RNNs can generate bounded hierarchical languages with optimal memory.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 1978–2010, Online. Association for Computational Linguistics.
  * de la Higuera (2010) Colin de la Higuera. 2010.  _Grammatical Inference: Learning Automata and Grammars_.  Cambridge University Press.
  * Hochreiter and Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber. 1997.  Long short-term memory.  _Neural Computation_ , 9(8):1735–1780.
  * Hockett (1959) Charles F. Hockett. 1959.  Animal ``languages'' and human language.  _Human Biology_ , 31(1):32–39.
  * Hupkes et al. (2020) Dieuwke Hupkes, Verna Dankers, Mathijs Mul, and Elia Bruni. 2020.  Compositionality decomposed: How do neural networks generalise? (extended abstract).  In _Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI-20_ , pages 5065–5069. International Joint Conferences on Artificial Intelligence Organization.  Journal track.
  * Jacobsson (2005) Henrik Jacobsson. 2005.  Rule extraction from recurrent neural networks: A taxonomy and review.  _Neural Computation_ , 17(6):1223–1263.
  * Kingma and Ba (2015) Diederik P. Kingma and Jimmy Ba. 2015.  Adam: A method for stochastic optimization.  In _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_.
  * Lake and Baroni (2018) Brenden Lake and Marco Baroni. 2018.  Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks.  In _International Conference on Machine Learning_ , pages 2873–2882. PMLR.
  * Mahalunkar and Kelleher (2019) Abhijit Mahalunkar and John Kelleher. 2019.  Multi-element long distance dependencies: Using SPk languages to explore the characteristics of long-distance dependencies.  In _Proceedings of the Workshop on Deep Learning and Formal Languages: Building Bridges_ , pages 34–43, Florence. Association for Computational Linguistics.
  * Marcus (1998) Gary F. Marcus. 1998.  Rethinking eliminative connectionism.  _Cognitive Psychology_ , 37(3):243–282.
  * Mikolov et al. (2016) Tomas Mikolov, Armand Joulin, and Marco Baroni. 2016.  A roadmap towards machine intelligence.  In _International Conference on Intelligent Text Processing and Computational Linguistics_ , pages 29–61. Springer.
  * Mohri (1997) Mehryar Mohri. 1997.  Finite-state transducers in language and speech processing.  _Computational Linguistics_ , 23(2):269–311.
  * Montague (1970) Richard Montague. 1970.  Universal grammar.  _Theoria_.
  * Nelson et al. (2020) Max Nelson, Hossep Dolatian, Jonathan Rawski, and Brandon Prickett. 2020.  Probing RNN encoder-decoder generalization of subregular functions using reduplication.  In _Proceedings of the Society for Computation in Linguistics 2020_ , pages 167–178, New York, New York. Association for Computational Linguistics.
  * Oncina et al. (1993) José Oncina, Pedro García, and Enrique Vidal. 1993.  Learning subsequential transducers for pattern recognition interpretation tasks.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 15(5):448–458.
  * Partee (1995) Barbara Partee. 1995.  Lexical semantics and compositionality.  _An Invitation to Cognitive Science: Language_ , 1:311–360.
  * Rawski and Heinz (2019) Jonathan Rawski and Jeffrey Heinz. 2019.  No free lunch in linguistics or machine learning: Response to pater.  _Language_ , 95(1):e125–e135.
  * Rogers and Pullum (2011) James Rogers and Geoffrey K. Pullum. 2011.  Aural pattern recognition experiments and the subregular hierarchy.  _Journal of Logic, Language and Information_ , 20(3):329–342.
  * Ruis et al. (2020) Laura Ruis, Jacob Andreas, Marco Baroni, Diane Bouchacourt, and Brenden M Lake. 2020\.  A benchmark for systematic generalization in grounded language understanding.  In _Advances in Neural Information Processing Systems 33 pre-proceedings (NeurIPS 2020)_.
  * Schützenberger (1962) Marcel P. Schützenberger. 1962.  Finite counting automata.  _Information and Control_ , 5:91–107.
  * Suzgun et al. (2019) Mirac Suzgun, Yonatan Belinkov, Stuart Shieber, and Sebastian Gehrmann. 2019.  LSTM networks can perform dynamic counting.  In _Proceedings of the Workshop on Deep Learning and Formal Languages: Building Bridges_ , pages 44–54, Florence. Association for Computational Linguistics.
  * Vidal et al. (2005) E. Vidal, F. Thollard, C. de la Higuera, F. Casacuberta, and R. C. Carrasco. 2005.  Probabilistic finite-state machines - part I.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 27(7):1013–1025.
  * Weiss et al. (2018) Gail Weiss, Yoav Goldberg, and Eran Yahav. 2018.  On the practical computational power of finite precision RNNs for language recognition.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)_ , pages 740–745, Melbourne, Australia. Association for Computational Linguistics.
  * Wu and Cotterell (2019) Shijie Wu and Ryan Cotterell. 2019.  Exact hard monotonic attention for character-level transduction.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 1530–1537, Florence, Italy. Association for Computational Linguistics.


# Importance of Aligning Training Strategy with Evaluation for Diffusion
Models
in 3D Multiclass Segmentation

Yunguan Fu
University College London
InstaDeep
& Yiwen Li
University of Oxford
& Shaheer U. Saeed
University College London
& Matthew J. Clarkson
University College London
& Yipeng Hu
University College London
University of Oxford

###### Abstract

Recently, denoising diffusion probabilistic models (DDPM) have been applied to
image segmentation by generating segmentation masks conditioned on images,
while the applications were mainly limited to 2D networks without exploiting
potential benefits from the 3D formulation. In this work, we studied the DDPM-
based segmentation model for 3D multiclass segmentation on two large
multiclass data sets (prostate MR and abdominal CT). We observed that the
difference between training and test methods led to inferior performance for
existing DDPM methods. To mitigate the inconsistency, we proposed a recycling
method which generated corrupted masks based on the model’s prediction at a
previous time step instead of using ground truth. The proposed method achieved
statistically significantly improved performance compared to existing DDPMs,
independent of a number of other techniques for reducing train-test
discrepancy, including performing mask prediction, using Dice loss, and
reducing the number of diffusion time steps during training. The performance
of diffusion models was also competitive and visually similar to non-
diffusion-based U-net, within the same compute budget. The JAX-based diffusion
framework has been released at https://github.com/mathpluscode/ImgX-DiffSeg.

_Keywords_ Image Segmentation $\cdot$ Diffusion Model $\cdot$ Prostate MR
$\cdot$ Abdominal CT

##  1 Introduction

Multiclass segmentation is one of the most basic tasks in medical imaging, one
that arguably benefited the most from deep learning. Although different model
architectures (Li et al., 2021; Strudel et al., 2021) and training strategies
(Fu et al., 2019; Li et al., 2022) have been proposed for specific clinical
applications, U-net (Ronneberger et al., 2015) trained through supervised
training remains the state-of-the-art and an important baseline for many (Ji
et al., 2022). Recently, denoising diffusion probabilistic models (DDPM) have
been demonstrated to be effective in a variety of image synthesis tasks (Ho et
al., 2020), which can be further guided by a scoring model to generate
conditioned images (Dhariwal and Nichol, 2021) or additional inputs (Ho and
Salimans, 2022). These generative modelling results are followed by image
segmentation, where the model generates segmentation masks by progressive
denoising from random noise. During training, DDPM is provided with an image
and a noise-corrupted segmentation mask, generated by a linear interpolation
between the ground-truth and a sampled noise. The model is then tasked to
predict the sampled noise (Amit et al., 2021; Kolbeinsson and Mikolajczyk,
2022; Wolleb et al., 2022; Wu et al., 2022) or the ground-truth mask (Chen et
al., 2022).

However, existing applications have been mainly based on 2D networks and, for
3D volumetric medical images, slices are segmented before obtaining the
assembled 3D segmentation. Challenges are often encountered for 3D images.
First, the diffusion model requires image and noise-corrupted masks as input,
leading to an increased memory footprint resulting in limited batch size and
potentially excessive training time. For instance, the transformer-based
architecture becomes infeasible without reducing model size or image, given
clinically or academically accessible hardware with limited memory. Second,
most diffusion models assume a denoising process of hundreds of time steps for
training and inference, the latter of which in particular leads to prohibitive
inference time (e.g., days on TPUs/GPUs).

This work addresses these issues by aligning training with evaluation
processes via recycling. As discussed in multiple studies (Chen et al., 2022;
Young et al., 2022; Kolbeinsson and Mikolajczyk, 2022; Lai et al., 2023),
noise does not necessarily disrupt the shape of ground truth masks and
morphological features may be preserved in noise-corrupted masks during
training. By training with recycling (Figure 1), the prediction from the
previous steps is used as input, i.e. rather than the ground truth used in
existing methods, for noisy mask sampling. This proposed training process
emulates the test process since the input is also from the previous
predictions at inference time for diffusion models, without access to ground
truth. Furthermore, this work directly predicts ground-truth masks instead of
sampled noise (Wu et al., 2022). This facilitates the direct use of Dice loss
in addition to cross-entropy during training, as opposed tp $L_{2}$ loss on
noise. Lastly, instead of denoising with at least hundreds of steps as in most
existing work, we propose a five-step denoising process for both training and
inference, resorting to resampling variance scheduling (Nichol and Dhariwal,
2021).

With extensive experiments in two of the largest public multiclass
segmentation applications, prostate MR (589 images) and abdominal CT images
(300 images) (Li et al., 2022; Ji et al., 2022), we demonstrated a
statistically significant improvement (between $0.015$ and $0.117$ in Dice
score) compared to existing DDPMs. Compared to non-diffusion supervised
learning, diffusion models reached a competitive performance (between $0.008$
and $0.015$ in Dice), with the same computational cost. With high transparency
and reproducibility, avoiding selective results under different conditions, we
conclude that the proposed recycling strategy using mask prediction setting
with Dice loss should be the default configuration for 3D segmentation
applications with diffusion models. We release the first unit-tested JAX-based
diffusion segmentation framework at https://github.com/mathpluscode/ImgX-
DiffSeg.

##  2 Related Work

The diffusion probabilistic model was first proposed by Sohl-Dickstein et al.
(2015) as a generative model for image sampling with a forward noising
process. Ho et al. (2020) proposed a reverse denoising process that estimates
the sampled error, achieving state-of-the-art performance in unconditioned
image synthesis at the time. Different conditioning methods were later
proposed to guide the sampling process toward a desired image class or prompt
text, using gradients from an external scoring model (Dhariwal and Nichol,
2021; Radford et al., 2021). Alternatively, Ho and Salimans (2022) showed that
guided sampling can be achieved by providing conditions during training.
Diffusion models have been successfully applied in medical imaging
applications to synthesise images of different modalities, such as
unconditioned lung X-Ray and CT (Ali et al., 2023), patient-conditioned brain
MR (Pinaya et al., 2022), temporal cardiac MR (Kim and Ye, 2022), and
pathology/sequence-conditioned prostate MR (Saeed et al., 2023). The
synthesised images have been shown to benefit pre-training self-supervised
models (Khader et al., 2022; Saeed et al., 2023) or support semi-supervised
learning (Young et al., 2022).

Besides image synthesis, Baranchuk et al. (2021) used pre-trained diffusion
models’ intermediate feature maps to train pixel classifiers for segmentation,
showing these unsupervised models capture semantics that can be extended for
image segmentation especially when training data is limited. Alternatively,
Amit et al. (2021) performed progressive denoising from random sampled noise
to generate segmentation masks instead of images for microscopic images (Amit
et al., 2021). At each step, the model takes a noise-corrupted mask and an
image as input and predicts the sampled noise. Similar approaches have been
also applied to thyroid lesion segmentation for ultrasound images (Wu et al.,
2022) and brain tumour segmentation for MR images with different network
architectures (Wolleb et al., 2022; Wu et al., 2022). Empirically, multiple
studies (Chen et al., 2022; Young et al., 2022; Kolbeinsson and Mikolajczyk,
2022) found the noise-corrupted mask generation, via linear interpolation
between ground-truth masks and noise, retained morphological features during
training, causing potential data leakage. Chen et al. (2022) therefore added
noises to mask analog bit and tuned its scaling. Young et al. (2022), on the
other hand, tuned the variance and scaling of added normal noise to reduce
information contained in noised masks. Furthermore, Kolbeinsson and
Mikolajczyk (2022) proposed recursive denoising instead of directly using
ground truth for noise-corrupted mask generation.

These works, although using different methods, are all addressing a similar
concern: the diffusion model training process is different from its evaluation
process, which potentially hinges the efficient learning. Moreover, most
published diffusion-model-based segmentation applications have been based on
2D networks. We believe such discrepancy would be more significant when
applying 3D networks to volumetric images due to the increased difficulty,
resulting in longer training and larger compute cost. In this work, building
on these recent developments, we focus on a consistent train-evaluate
algorithm for efficient training of diffusion models in 3D medical image
segmentation applications.

##  3 Method

Figure 1: Illustration of training with and without recycling and inference,
using mask prediction. For training without recycling (top left), the noisy
mask $\mathbf{x}_{t}$ is calculated using the ground truth mask. For training
with recycling (top right), $\mathbf{x}_{t}$ is calculated using prediction
from the previous step, which is similar to the inference (bottom right).

###  3.1 DDPM for Segmentation

The denoising diffusion probabilistic models (DDPM) (Sohl-Dickstein et al.,
2015; Ho et al., 2020; Nichol and Dhariwal, 2021; Kingma et al., 2021)
consider a forward process: given sample $\mathbf{x}_{0}\sim
q(\mathbf{x}_{0})$, noisy $\mathbf{x}_{t}$ for $t=1,\cdots,T$ follows a
multivariate normal distribution,
$q(\mathbf{x}_{t}\mid\mathbf{x}_{t-1})=\mathcal{N}(\mathbf{x}_{t};\sqrt{1-\beta_{t}}\mathbf{x}_{t-1},\beta_{t}\mathbf{I})$,
where $\beta_{t}\in[0,1]$. Given a sufficiently large $T$, $\mathbf{x}_{T}$
approximately follows an isotropic multivariate normal distribution
$\mathcal{N}(\mathbf{x}_{T};\mathbf{0},\mathbf{I})$. A reverse process is then
defined to denoise $\mathbf{x}_{t}$ at each step, for $t=T,\cdots,1$,
$p_{\theta}(\mathbf{x}_{t-1}\mid\mathbf{x}_{t})=\mathcal{N}(\mathbf{x}_{t-1};\bm{\mu}_{\theta}(\mathbf{x}_{t},t),\tilde{\beta}_{t}\mathbf{I})$,
with a predicted mean $\bm{\mu}_{\theta}(\mathbf{x}_{t},t)$ and a variance
schedule $\tilde{\beta}_{t}\mathbf{I}$.
$\tilde{\beta}_{t}=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}}\beta_{t}$,
where $\alpha_{t}=1-\beta_{t}$, $\bar{\alpha}_{t}=\prod_{s=0}^{t}\alpha_{s}$.
$\bm{\mu}_{\theta}(\mathbf{x}_{t},t)$ can be modelled in two different ways,

$\displaystyle\bm{\mu}_{\theta}(\mathbf{x}_{t},t)$
$\displaystyle=\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_{t}}{1-\bar{\alpha}_{t}}\mathbf{x}_{0,\theta}(\mathbf{x}_{t},t)+\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}}\sqrt{\alpha_{t}}\mathbf{x}_{t},\leavevmode\nobreak\
(\text{Predict}\leavevmode\nobreak\ \mathbf{x}_{0})$ (1)
$\displaystyle\bm{\mu}_{\theta}(\mathbf{x}_{t},t)$
$\displaystyle=\frac{1}{\sqrt{\alpha}_{t}}(\mathbf{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\bm{\epsilon}_{t,\theta}(\mathbf{x}_{t},t)),\leavevmode\nobreak\
(\text{Predict noise}\leavevmode\nobreak\ \bm{\epsilon}_{t})$ (2)

where $\mathbf{x}_{0,\theta}(\mathbf{x}_{t},t)$ or
$\bm{\epsilon}_{t,\theta}(\mathbf{x}_{t},t)$ are the learned neural network.

For segmentation, $\mathbf{x}$ represents a transformed probability with
values in $[-1,1]$. Particularly, $\mathbf{x}_{0}\in\\{1,-1\\}$ are binary-
valued, transformed from mask. $\mathbf{x}_{0,\theta}$ and $\mathbf{x}_{t}$
($t\geq 0$) have values in $[-1,1]$. Moreover, the networks
$\mathbf{x}_{0,\theta}(I,\mathbf{x}_{t},t)$ or
$\bm{\epsilon}_{t,\theta}(I,\mathbf{x}_{t},t)$ takes one more input $I$,
representing the image to segment.

###  3.2 Recycling

During training, existing methods samples $\mathbf{x}_{t}$ by interpolating
noise $\bm{\epsilon}_{t}$ and ground-truth $\mathbf{x}_{0}$, which results in
a certain level of data leak (Chen et al., 2022; Young et al., 2022).
Kolbeinsson and Mikolajczyk (2022) proposed recursive denoising, which
performed $T$ steps on each image progressively to use model’s predictions at
previous steps. However, this extends the training length $T$ times. Instead,
in this work, for each image, the time step $t$ is randomly sampled and the
model’s prediction $\mathbf{x}_{0,\theta}$ from the previous time step $t+1$
is recycled to replace ground-truth (Fig. 1).

Similar reuses of the model’s predictions have been previously applied in 2D
image segmentation (Chen et al., 2022), however $\mathbf{x}_{0,\theta}$ was
fed into the network along with $\mathbf{x}_{t}$ which requires additional
memories and still has data leak risks. A further difference to these previous
approaches is that, rather than stochastic recycling, usually with a
probability of $50\%$, it is always applied throughout the training (which was
empirically found to lead to more stable and performant model training).
Formally, the recycling technique at a sampled step $t$ is as follows,

$\displaystyle\mathbf{x}_{t+1}$
$\displaystyle=\sqrt{\bar{\alpha}_{t+1}}\mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t+1}}\bm{\epsilon}_{t+1},\leavevmode\nobreak\
\text{(Noise mask generation for}\leavevmode\nobreak\ t+1)$ (3)
$\displaystyle\mathbf{x}_{0,\theta}$
$\displaystyle=\text{StopGradient}(\mathbf{x}_{0,\theta}(I,t+1,\mathbf{x}_{t+1})),\leavevmode\nobreak\
\text{(Mask prediction)}$ (4) $\displaystyle\mathbf{x}_{t}$
$\displaystyle=\sqrt{\bar{\alpha}_{t}}\mathbf{x}_{0,\theta}+\sqrt{1-\bar{\alpha}_{t}}\bm{\epsilon}_{t},\leavevmode\nobreak\
\text{(Noise mask generation for}\leavevmode\nobreak\ t)$ (5)

where $\mathbf{x}_{0,\theta}$ is the predicted segmentation mask from $t+1$
using ground-truth, with gradient stopping. $\bm{\epsilon}_{t}$ and
$\bm{\epsilon}_{t+1}$ are two independently sampled noises. Recycling can be
applied to models predicting noise (see supplementary materials Section 1 for
derivation and illustration).

###  3.3 Loss

Given noised mask $\mathbf{x}_{t}$, time $t$, and image $I$, the loss can be,

$\displaystyle L_{\text{seg},\mathbf{x}_{0}}(\theta)$
$\displaystyle=\mathbb{E}_{t,\mathbf{x}_{0},\bm{\epsilon}_{t},I}\mathcal{L}_{\text{seg}}(\mathbf{x}_{0},\leavevmode\nobreak\
\mathbf{x}_{0,\theta}),$ (6) $\displaystyle
L_{\text{seg},\bm{\epsilon}_{t}}(\theta)$
$\displaystyle=\mathbb{E}_{t,\mathbf{x}_{0},\bm{\epsilon}_{t},I}\|\bm{\epsilon}_{t}-\bm{\epsilon}_{t,\theta}\|_{2}^{2},$
(7)

where model predict noise $\bm{\epsilon}_{t,\theta}$ or mask
$\mathbf{x}_{0,\theta}$ and $\mathcal{L}_{\text{seg}}$ represents a
segmentation-specific loss, such as Dice loss or cross-entropy loss. $t$ is
sampled from $1$ to $T$,
$\bm{\epsilon}_{t}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$, and
$\mathbf{x}_{t}(\mathbf{x}_{0},\bm{\epsilon}_{t})=\sqrt{\bar{\alpha}_{t}}\mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}}\bm{\epsilon}_{t}$.
When model predicts noise, segmentation loss can still be used as the mask can
be inferred via interpolation.

###  3.4 Variance resampling

During training or inference, given a variance schedule
$\\{\beta_{t}\\}_{t=1}^{T}$ for $T=1000$, a subsequence
$\\{\beta_{k}\\}_{k=1}^{K}$ for $K=5$ can be sampled with
$\\{t_{k}\\}_{k=1}^{K}$, where $t_{K}=T$,
$t_{1}=1$,$\beta_{k}=1-\frac{\bar{\alpha}_{t_{k}}}{\bar{\alpha}_{t_{k-1}}},\tilde{\beta}_{k}=\frac{1-\bar{\alpha}_{t_{k}-1}}{1-\bar{\alpha}_{t_{k}}}\beta_{t_{k}}$.
$\alpha_{k}$ and $\bar{\alpha}_{k}$ are recalculated correspondingly.

##  4 Experiment Setting

Prostate MR The data set111https://zenodo.org/record/7013610#.ZAkaXuzP2rM (Li
et al., 2022) contains $589$ T2-weighted image-mask pairs for $8$ anatomical
structures. The images were randomly split into non-overlapping training,
validation, and test sets, with $411$, $14$, $164$ images in each split,
respectively. All images were normalised, resampled, and centre-cropped to an
image size of $256\times 256\times 48$, with a voxel dimension of $0.75\times
0.75\times 2.5$ (mm).
Abdominal CT The data set222https://zenodo.org/record/7155725#.ZAkbe-zP2rO (Ji
et al., 2022) provides $200$ and $100$ CT image-mask pairs for $15$ abdominal
organs in training and validation sets. The validation set was randomly split
into non-overlapping validation and test sets, with $10$ and $90$ images,
respectively. HU values were clipped to $[-991,362]$ for all images. Images
were then normalised, resampled and centre-cropped to an image size of
$192\times 128\times 128$, with a voxel dimension of $1.5\times 1.5\times 5.0$
(mm).
Implementation 3D U-nets have four layers with $32$, $64$, $128$, and $256$
channels. For diffusion models, noise-corrupted masks and images were
concatenated along feature channels and time was encoded using sinusoidal
positional embedding (Rombach et al., 2022). Random rotation, translation and
scaling were adopted for data augmentation during training. The segmentation-
specific loss function is by default the sum of cross-entropy and foreground-
only Dice loss. When predicting noise, the $L_{2}$ loss has a weight of $0.1$
(Wu et al., 2022). All models were trained with AdamW for $12500$ steps and a
warmup cosine learning rate schedule. Hyper-parameter were configured
empirically without extensive tuning. Binary Dice score and $95\%$ Hausdorff
distance in mm (HD), averaged over foreground classes, were reported. Paired
Student’s t-tests were performed on Dice score to test statistical
significance between model performance with $\alpha=0.01$. Experiments were
performed using bfloat16 mixed precision on TPU v3-8. The code is available at
https://github.com/mathpluscode/ImgX-DiffSeg.

##  5 Results

Recycling The performance of the proposed diffusion model is summarised in
Table 1. With recycling, the diffusion-based 3D models reached a Dice score of
$0.830$ and $0.801$ which is statistically significantly ($p<0.001$) higher
than baseline diffusion with $0.815$ and $0.753$, for prostate MR and
abdominal CT respectively. Example predictions were provided in Fig. 2.
Diffusion vs Non-diffusion The diffusion model is also compared with U-net
trained via standard supervised learning in Table 2. Within the same computing
budget, the diffusion-based 3D model is competitive with its non-diffusion
counterpart. The results in Fig. 2 are also visually comparable. The
difference, however, remains significant with $p<0.001$.
Ablation studies Comparisons were performed on prostate MR data set for other
modifications, including: 1) predicting mask instead of noise 2) using Dice
loss in addition to cross-entropy, 3) using five steps denoising process
during training. Improvements in Dice score between $0.09$ and $0.117$ were
observed for all modifications (all p-values $<0.001$). The largest
improvement was observed when the model predicted segmentation masks instead
of noise. The results were found consistent with the consistency model (Song
et al., 2023), which requests diffusion models’ predictions of
$\mathbf{x}_{0}$ from different time steps to be similar. Such requirement is
implicitly met in our applications as the segmentation loss requires the
prediction to be consistent with the ground truth mask given an image. As a
result, the predictions from different time steps shall be consistent.
Limitation In general, all methods tend to perform better for large regions of
interest (ROI), and there is a significant correlation (Spearman $r>0.8$ and
$p<0.01$) between ROI (regions of interest) area and mean Dice score per
ROI/class, indicating room of future improvement by addressing small ROIs.

Figure 2: Example predictions from diffusion and non-diffusion U-net models. Table 1: DDPM with recycling. Recycling | Prostate MR | Abdominal CT
---|---|---
Dice Score | Hausdorff Dist. | Dice Score | Hausdorff Dist.
N | $0.815\pm 0.095$ | $5.485\pm 1.069$ | $0.753\pm 0.131$ | $9.526\pm 2.232$
Y | $0.830\pm 0.094$ | $5.424\pm 1.176$ | $0.801\pm 0.109$ | $9.125\pm 2.564$
Table 2: Comparison between diffusion models and non-diffusion models. Diffusion | Prostate MR | Abdominal CT
---|---|---
Dice score | Hausdorff Dist. | Dice score | Hausdorff Dist.
N | $0.838\pm 0.088$ | $5.197\pm 1.184$ | $0.816\pm 0.100$ | $9.091\pm 2.475$
Y | $0.830\pm 0.094$ | $5.424\pm 1.176$ | $0.801\pm 0.109$ | $9.125\pm 2.564$
Table 3: Ablation results (prostate MR). HD stands for Hausdorff distance.

Output | Dice Score | HD
---|---|---
Noise | 0.713 | 23.855
Logits | 0.830 | 5.424
(a)

Dice Loss | Dice Score | HD
---|---|---
N | 0.812 | 5.463
Y | 0.830 | 5.424
(b)

Steps | Dice Score | HD
---|---|---
1000 | 0.821 | 5.223
5 | 0.830 | 5.424
(c)

##  6 Discussion

In this work, we developed a novel denoising diffusion probabilistic model for
3D image multiclass segmentation. By recycling the model’s predictions at
previous time steps to replace ground truth during training, the method aligns
diffusion training and segmentation evaluation, resulting in significant
performance improvements compared to existing diffusion methods. Other
techniques mitigating training and test inconsistency further improved the
diffusion model’s performance. However, the diffusion model did not outperform
the non-diffusion-based segmentation models, which have long been well-
established. We believe it is important to report this lack of superior
performance in 3D medical image segmentation, especially when experiments are
limited to the same compute budget. Future work could consider other diffusion
models such as discrete diffusion and more memory-efficient implementation to
enable more complex architectures. Although the presented experimental results
primarily demonstrated methodological development, the fact that these were
obtained on two large clinical data sets represents a promising step towards
real-world applications. Localising multiple anatomical structures in prostate
MR images is key to MR-targeted biopsy, radiotherapy and tissue-preserving
focal treatment for patients with urological diseases, while abdominal organ
outlines can be directly used in planning gastroenterological procedures and
hepatic surgery.

##  7 Acknowledgement

This work was supported by the EPSRC grant (EP/T029404/1), the Wellcome/EPSRC
Centre for Interventional and Surgical Sciences (203145Z/16/Z), the
International Alliance for Cancer Early Detection, an alliance between Cancer
Research UK (C28070/A30912, C73666/A31378), Canary Center at Stanford
University, the University of Cambridge, OHSU Knight Cancer Institute,
University College London and the University of Manchester, and Cloud TPUs
from Google’s TPU Research Cloud (TRC).

## References

  * Li et al. [2021] Xiaoxiao Li, Yuan Zhou, Nicha Dvornek, Muhan Zhang, Siyuan Gao, Juntang Zhuang, Dustin Scheinost, Lawrence H Staib, Pamela Ventola, and James S Duncan.  Braingnn: Interpretable brain graph neural network for fmri analysis.  _Medical Image Analysis_ , 74:102233, 2021.
  * Strudel et al. [2021] Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid.  Segmenter: Transformer for semantic segmentation.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pages 7262–7272, 2021.
  * Fu et al. [2019] Yunguan Fu, Maria R Robu, Bongjin Koo, Crispin Schneider, Stijn van Laarhoven, Danail Stoyanov, Brian Davidson, Matthew J Clarkson, and Yipeng Hu.  More unlabelled data or label more data? a study on semi-supervised laparoscopic image segmentation.  _arXiv preprint arXiv:1908.08035_ , 2019.
  * Li et al. [2022] Yiwen Li, Yunguan Fu, Iani Gayo, Qianye Yang, Zhe Min, Shaheer Saeed, Wen Yan, Yipei Wang, J Alison Noble, Mark Emberton, et al.  Prototypical few-shot segmentation for cross-institution male pelvic structures with spatial registration.  _arXiv preprint arXiv:2209.05160_ , 2022.
  * Ronneberger et al. [2015] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In _Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18_ , pages 234–241. Springer, 2015.
  * Ji et al. [2022] Yuanfeng Ji, Haotian Bai, Jie Yang, Chongjian Ge, Ye Zhu, Ruimao Zhang, Zhen Li, Lingyan Zhang, Wanling Ma, Xiang Wan, et al.  Amos: A large-scale abdominal multi-organ benchmark for versatile medical image segmentation.  _arXiv preprint arXiv:2206.08023_ , 2022.
  * Ho et al. [2020] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Dhariwal and Nichol [2021] Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  _Advances in Neural Information Processing Systems_ , 34:8780–8794, 2021.
  * Ho and Salimans [2022] Jonathan Ho and Tim Salimans.  Classifier-free diffusion guidance.  _arXiv preprint arXiv:2207.12598_ , 2022.
  * Amit et al. [2021] Tomer Amit, Eliya Nachmani, Tal Shaharbany, and Lior Wolf.  Segdiff: Image segmentation with diffusion probabilistic models.  _arXiv preprint arXiv:2112.00390_ , 2021.
  * Kolbeinsson and Mikolajczyk [2022] Benedikt Kolbeinsson and Krystian Mikolajczyk.  Multi-class segmentation from aerial views using recursive noise diffusion.  _arXiv preprint arXiv:2212.00787_ , 2022.
  * Wolleb et al. [2022] Julia Wolleb, Robin Sandkühler, Florentin Bieder, Philippe Valmaggia, and Philippe C Cattin.  Diffusion models for implicit image segmentation ensembles.  In _International Conference on Medical Imaging with Deep Learning_ , pages 1336–1348. PMLR, 2022.
  * Wu et al. [2022] Junde Wu, Huihui Fang, Yu Zhang, Yehui Yang, and Yanwu Xu.  Medsegdiff: Medical image segmentation with diffusion probabilistic model.  _arXiv preprint arXiv:2211.00611_ , 2022.
  * Chen et al. [2022] Ting Chen, Lala Li, Saurabh Saxena, Geoffrey Hinton, and David J Fleet.  A generalist framework for panoptic segmentation of images and videos.  _arXiv preprint arXiv:2210.06366_ , 2022.
  * Young et al. [2022] Sean I Young, Adrian V Dalca, Enzo Ferrante, Polina Golland, Bruce Fischl, and Juan Eugenio Iglesias.  Sud: Supervision by denoising for medical image segmentation.  _arXiv preprint arXiv:2202.02952_ , 2022.
  * Lai et al. [2023] Zeqiang Lai, Yuchen Duan, Jifeng Dai, Ziheng Li, Ying Fu, Hongsheng Li, Yu Qiao, and Wenhai Wang.  Denoising diffusion semantic segmentation with mask prior modeling.  _arXiv preprint arXiv:2306.01721_ , 2023.
  * Nichol and Dhariwal [2021] Alexander Quinn Nichol and Prafulla Dhariwal.  Improved denoising diffusion probabilistic models.  In _International Conference on Machine Learning_ , pages 8162–8171. PMLR, 2021.
  * Sohl-Dickstein et al. [2015] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _International Conference on Machine Learning_ , pages 2256–2265. PMLR, 2015.
  * Radford et al. [2021] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In _International conference on machine learning_ , pages 8748–8763. PMLR, 2021.
  * Ali et al. [2023] Hazrat Ali, Shafaq Murad, and Zubair Shah.  Spot the fake lungs: Generating synthetic medical images using neural diffusion models.  In _Artificial Intelligence and Cognitive Science: 30th Irish Conference, AICS 2022, Munster, Ireland, December 8–9, 2022, Revised Selected Papers_ , pages 32–39. Springer, 2023.
  * Pinaya et al. [2022] Walter HL Pinaya, Petru-Daniel Tudosiu, Jessica Dafflon, Pedro F da Costa, Virginia Fernandez, Parashkev Nachev, Sebastien Ourselin, and M Jorge Cardoso.  Brain imaging generation with latent diffusion models.  _arXiv preprint arXiv:2209.07162_ , 2022.
  * Kim and Ye [2022] Boah Kim and Jong Chul Ye.  Diffusion deformable model for 4d temporal medical image generation.  In _Medical Image Computing and Computer Assisted Intervention–MICCAI 2022: 25th International Conference, Singapore, September 18–22, 2022, Proceedings, Part I_ , pages 539–548. Springer, 2022.
  * Saeed et al. [2023] Shaheer U. Saeed, Tom Syer, Wen Yan, Qianye Yang, Mark Emberton, Shonit Punwani, Matthew J. Clarkson, Dean C. Barratt, and Yipeng Hu.  Bi-parametric prostate mr image synthesis using pathology and sequence-conditioned stable diffusion.  _arXiv preprint arXiv:2303.02094_ , 2023.
  * Khader et al. [2022] Firas Khader, Gustav Mueller-Franzes, Soroosh Tayebi Arasteh, Tianyu Han, Christoph Haarburger, Maximilian Schulze-Hagen, Philipp Schad, Sandy Engelhardt, Bettina Baessler, Sebastian Foersch, et al.  Medical diffusion–denoising diffusion probabilistic models for 3d medical image generation.  _arXiv preprint arXiv:2211.03364_ , 2022.
  * Baranchuk et al. [2021] Dmitry Baranchuk, Ivan Rubachev, Andrey Voynov, Valentin Khrulkov, and Artem Babenko.  Label-efficient semantic segmentation with diffusion models.  _arXiv preprint arXiv:2112.03126_ , 2021.
  * Kingma et al. [2021] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho.  Variational diffusion models.  _Advances in neural information processing systems_ , 34:21696–21707, 2021.
  * Rombach et al. [2022] 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.
  * Song et al. [2023] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever.  Consistency models.  _arXiv preprint arXiv:2303.01469_ , 2023.

## Appendix

Recycling can be applied to models predicting noise as below (Figure 3),

$\displaystyle\mathbf{x}_{t+1}$
$\displaystyle=\sqrt{\bar{\alpha}_{t+1}}\mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t+1}}\bm{\epsilon}_{t+1},\leavevmode\nobreak\
\text{(Noise mask generation for}\leavevmode\nobreak\ t+1)$ (8)
$\displaystyle\bm{\epsilon}_{t+1,\theta}$
$\displaystyle=\text{StopGradient}(\bm{\epsilon}_{t+1,\theta}(I,t+1,\mathbf{x}_{t+1})),\leavevmode\nobreak\
\text{(Noise prediction)}$ (9) $\displaystyle\mathbf{x}_{0,\theta}$
$\displaystyle=\frac{1}{\sqrt{\bar{\alpha}_{t+1}}}(\mathbf{x}_{t+1}-\sqrt{1-\bar{\alpha}_{t+1}}\bm{\epsilon}_{t+1,\theta}),\leavevmode\nobreak\
\text{(Mask prediction)}$ (10) $\displaystyle\mathbf{x}_{t}$
$\displaystyle=\sqrt{\bar{\alpha}_{t}}\mathbf{x}_{0,\theta}+\sqrt{1-\bar{\alpha}_{t}}\bm{\epsilon}_{t}.\leavevmode\nobreak\
\text{(Noise mask generation for}\leavevmode\nobreak\ t)$ (11)

Figure 3: Illustration of training with and without recycling, and inference,
using noise prediction. For training without recycling (top left), the noisy
mask $\mathbf{x}_{t}$ is calculated using the ground truth mask. For training
with recycling (top right), $\mathbf{x}_{t}$ is calculated using prediction
from the previous step, which is similar to the inference (bottom right).