# Parallelizing Legendre Memory Unit Training

Narsimha Chilkuri  Chris Eliasmith

###### Abstract

Recently, a new recurrent neural network (RNN) named the Legendre Memory Unit
(LMU) was proposed and shown to achieve state-of-the-art performance on
several benchmark datasets. Here we leverage the linear time-invariant (LTI)
memory component of the LMU to construct a simplified variant that can be
parallelized during training (and yet executed as an RNN during inference),
thus overcoming a well known limitation of training RNNs on GPUs. We show that
this reformulation that aids parallelizing, which can be applied generally to
any deep network whose recurrent components are linear, makes training up to
200 times faster. Second, to validate its utility, we compare its performance
against the original LMU and a variety of published LSTM and transformer
networks on seven benchmarks, ranging from psMNIST to sentiment analysis to
machine translation. We demonstrate that our models exhibit superior
performance on all datasets, often using fewer parameters. For instance, our
LMU sets a new state-of-the-art result on psMNIST, and uses half the
parameters while outperforming DistilBERT and LSTM models on IMDB sentiment
analysis.

Machine Learning, ICML

##  1 Introduction

LSTMs (Hochreiter & Schmidhuber, 1997), the most popular class of recurrent
neural networks (RNNs), process information in an inherently sequential
manner. This prevents parallelization within training examples, and thus, they
cannot fully leverage today’s commodity GPU hardware. This sequential nature
of RNNs is one of the critical reasons for the shift towards purely self-
attention based architectures such as the transformer and its derivatives
(Vaswani et al., 2017; Devlin et al., 2018; Brown et al., 2020), especially in
the domain of Natural Language Processing (NLP). Parallelization makes
training such networks far more efficient on GPUs, and thus allows them to be
applied to enormous datasets – for example, the Colossal Clean Crawled Corpus
(C4) which comprises 750GB of English text (Raffel et al., 2019). Hence, such
models, via unsupervised pre-training, make it possible for us to exploit
resources such as the internet,111https://commoncrawl.org which produces 20TB
of text data each month. A feat such as this, from the training perspective,
would be unimaginable using LSTM or other RNN based models.

In this paper, we exploit the idea that linear recurrence relations can be
solved. More specifically, a linear time-invariant (LTI) system’s state update
equation can be written in a non-sequential fashion (Åström & Murray, 2010).
This allows for the computation of the hidden state of an LTI system to be
done in parallel, thus overcoming the fundamental constraint that traditional
RNNs suffer from. Despite this desirable property, to our knowledge, such
systems have remained underexplored. This perhaps has to do with the general
belief that nonlinear dynamics are critical for solving complicated problems,
especially AI level tasks. Here, we challenge that notion by focussing our
attention on a specific LTI system, the memory component of the recently
proposed Legendre Memory Unit (Voelker et al., 2019), and demonstrating
superior performance of our model on a variety of tasks, including predicting
a nonlinear, chaotic dynamical system.

The Legendre Memory Unit (LMU) is an RNN cell that is constructed by coupling
an LTI system to a nonlinear one. The LTI system, known as the Delay Network
(Voelker & Eliasmith, 2018), is the core component of the LMU that projects a
sliding window of length $\theta$ of the input signal onto the Legendre basis
– hence the name Legendre Memory Unit. This construction, which employs the
Delay Network (DN) to act as the memory and a nonlinear recurrent layer to
compute arbitrary functions across time, has been shown to outperform LSTMs
and other RNNs on various tasks. Of special interest to us is the ability of
the LMU to handle temporal dependencies across 100,000 time-steps, which is
orders of magnitude better than the LSTM.

In the following pages, we start by simplifying the LMU such that recurrence
exists only in the linear system. After showing that the training of this
simplified variant can be parallelized, we turn to the question of its
effectiveness. We do so by first considering the two synthetic tasks that the
original LMU was validated on: psMNIST and Mackey-Glass. We show that our
variant exceeds the original LMU’s performance on both of these tasks, while
also establishing a new state-of-the-art result for RNNs on psMNIST. We also
test our models against LSTM models of comparable sizes on several NLP tasks.
We look at sentiment classification (IMDB), semantic similarity (QQP), and
natural language inference (SNLI), and show that our models achieve better
performance while using significantly fewer parameters (up to 650x). We then
briefly investigate the transfer learning abilities of our model by training a
language model on the (unlabelled) Amazon Reviews dataset and using that to
improve the IMDB score. Here, we show that it outperforms DistilBert, a
transfomer based model, while using 50% fewer parameters. We conclude our NLP
experiments by demonstrating superior performance on language modelling
(text8) and machine translation (IWSLT’15 En-Vi). Finally, we show that these
architectural modifications result in training times that are up to 200 times
faster, relative to the original LMU.

##  2 Related Work

There have been recent studies focused on the application of LMU to new
problems and on the modification of the original architecture. For example,
Blouw et al. (2020) used an LMU based architecture to achieve SotA results on
the problem of efficient key-word spotting, and Gu et al. (2020) recently
proposed a generalization of the original architecture, which they use to beat
the SotA result set by the LMU on psMNIST.

Our model, however, is more directly inspired by the success of self-
attention. Self-attention based architectures have come to replace RNN based
approaches for problems such as language modelling, machine translation, and a
slew of other NLP tasks (Radford et al., 2018; Raffel et al., 2019). Three
properties that make self-attention desirable over RNNs are: (1) it is better
at handling the challenging problem of long-range dependencies; (2) it is
purely feedforward; and (3) when the sequence length is smaller than the
dimension of representation, which is not uncommon in NLP applications, self-
attention is computationally cheaper than an RNN.

Of these three desirable properties, our model inherits the first one from the
original LMU, and satisfies properties (2) and (3) by construction.
Additionally, the capability to run our model in a recurrent manner during
inference can be advantageous in situations where there are memory constraints
or where low-latency is critical.

It should also be noted that prior works such as Balduzzi & Ghifary (2016) and
Martin & Cundy (2017) have explored the use of linear RNNs for sequence
modelling. Our work differs from theirs in a few ways. First, we do not learn
the recurrent connections, but instead make use of an LTI system, the DN, that
optimally implements the delay operation; the use of frozen connections also
leads to a more efficient parallelization scheme. Second, while previous works
restrict the recurrent relations to be elementwise (recurrent weight matrix is
diagonal), we use full matrices in the linear system. Finally, we test our
model on a wide range of tasks, well beyond the synthetic datasets (Martin &
Cundy, 2017) and language modelling (Balduzzi & Ghifary, 2016) considered
previously.

##  3 Background and LMU Variants

In this section we start by introducing the delay problem and show how a delay
is optimally realized by the DN, an LTI system. We then introduce the LMU
which employs a single-input DN coupled to a nonlinear dynamical system to
process sequential data. Finally, we introduce our model, which is obtained by
simplifying the original architecture. We show how this simplification allows
for training to be parallelized (often reducing the computation complexity),
while also retaining the ability to handle long range dependencies of the
original formulation.

###  3.1 Delay Network

##### Ideal Delay

A system is said be an ideal delay system if it takes in an input, $u(t)$, and
outputs a function, $y(t)$, which is the delayed version of the input.
Mathematically, this can be described in the following manner:

$y(t)=\mathcal{D}[u(t)]=\begin{cases}0\quad t<\theta\\\ u(t-\theta)\quad
t\geq\theta\end{cases},$ (1)

where $D$ is the ideal delay operator and $\theta\in\mathbb{R}$ is the length
of the delay. There are two things of note here: (1) the ideal delay system is
linear, i.e., for any two functions, $f(t)$ and $g(t)$, and any
$a,b\in\mathbb{R}$, it respects the following equation:

$\mathcal{D}[af(t)+bg(t)]=a\mathcal{D}[f(t)]+b\mathcal{D}[g(t)];$ (2)

and (2) although this problem looks deceptively simple, in fact it takes a
system with infinite memory to take in a continuous stream of input (with
unspecified frequency content), store it for $\theta$ seconds and then
reproduce the entire input without error.

These two considerations tell us that the optimal system that implements delay
must be linear and that even the most optimal physical implementation can only
approximately realize the delay.

##### Approximating Delay

If we are interested in constructing a dynamical system that implements delay,
thanks to the linearity constraint, we can narrow our search space from the
general system of ODEs of the form:

$\displaystyle\dot{{\bm{m}}}=\textbf{f}({\bm{m}},u),$ (3) $\displaystyle
y=\textbf{g}({\bm{m}},u),$ (4)

to just finding the four matrices $({\bm{A}},{\bm{B}},{\bm{C}},{\bm{D}})$ that
define an LTI system:

$\displaystyle\dot{{\bm{m}}}={\bm{A}}{\bm{m}}+{\bm{B}}u,$ (5) $\displaystyle
y={\bm{C}}{\bm{m}}+{\bm{D}}u.$ (6)

Following the derivation of the Delay Network in Voelker & Eliasmith (2018),
we start by considering the transfer function of the delay system, which for a
SISO system is defined as:

$G(s)=\frac{y(s)}{u(s)}=e^{-\theta s},$ (7)

where $y(s)$ and $u(s)$ are found by taking the Laplace transform of the input
and output functions in time. As expected, this defines an infinite
dimensional transfer function, capturing the intuitive difficulty of
constructing a continuous delay.

The transfer function can be converted to a finite, causal state space
realization if and only if it can be written as a proper222A ratio
$\frac{a(s)}{b(s)}$ is said be proper if the order of the numerator does not
exceed the order of the denominator. ratio of finite dimensional polynomials
in $s$ (Brogan, 1991). $G(s)$, however, is irrational, i.e., it cannot be
written as a proper, finite dimensional ratio. Therefore, making an
approximation is necessary.

We can achieve an optimal convergence rate (in the least-square error sense)
for rational approximants by means of Padé approximants (Partington, 2004).
Choosing the order of the numerator to be one less than the order of the
denominator and accounting for numerical issues in the state-space realization
(see Voelker & Eliasmith (2018) for details), gives the following canonical
realization:

$\displaystyle A_{i,j}=\frac{(2i+1)}{\theta}\begin{cases}-1\qquad i<j\\\
(-1)^{i-j+1}\quad i\geq j\end{cases},$ (8) $\displaystyle
B_{i}=\frac{(2i+1)(-1)^{i}}{\theta},$ (9) $\displaystyle
C_{i}=(-1)^{i}\sum_{l=0}^{i}{i\choose l}{i+l\choose j}(-1)^{l},$ (10)
$\displaystyle D=0,\quad i,j\in[0,d-1],$ (11)

where we refer to $d$ as the order of the system. The LTI system
$({\bm{A}},{\bm{B}},{\bm{C}},{\bm{D}})$ is what is known as a Delay Network
(DN). The order of the system, $d$, and the delay length, $\theta$ are the
main hyperparameters to choose when using a DN. Higher order systems require
more resources, but provide a more accurate emulation of the ideal delay.
Because we have used Padé approximants, each order is optimal for that
dimension of state vector ${\bm{m}}$.

##### Legendre Polynomials

Suppose we construct a system using the
$({\bm{A}},{\bm{B}},{\bm{C}},{\bm{D}})$ matrices define above, and provide it
with an input signal, $u(t)$. Given the state ${\bm{m}}_{t}$, we can use
${\bm{C}}$ to decode $u(t-\theta)$ to a degree of accuracy determined by the
order of our system, i.e.,

$u(t-\theta)\approx{\bm{C}}^{T}{\bm{m}}_{t}.$ (12)

Intuitively, given ${\bm{m}}_{t}$, it seems possible to decode not only
$u(t-\theta)$ but also
$u(t-\theta^{\prime})\,\,\forall\,\,0\leq\theta^{\prime}\leq\theta$. This can
be done using a slightly modified ${\bm{C}}$ for a given $\theta^{\prime}$:

$u(t-\theta^{\prime})\approx{\bm{C}}(\theta^{\prime})^{T}{\bm{m}}_{t},$ (13)

where

$C_{i}(\theta^{\prime})=(-1)^{i}\sum_{l=0}^{i}{i\choose l}{i+l\choose
j}\left(-\frac{\theta^{\prime}}{\theta}\right)^{l},0\leq\theta^{\prime}\leq\theta,$
(14)

and ${\bm{C}}(\theta^{\prime}=\theta)$ corresponds to the ${\bm{C}}$ defined
in equation (10). Interestingly, the functions in (14) turn out to be the
shifted Legendre polynomials.

###  3.2 Legendre Memory Unit

The LMU is obtained by coupling a single-input delay network to a nonlinear
dynamical system. The DN orthogonalizes the input signal across a sliding
window of length $\theta$, whereas the nonlinear system relies on this memory
to compute arbitrary functions across time. The state update equations that
define the LMU are given below:

$\displaystyle
u_{t}={\bm{e}}_{x}^{T}{\bm{x}}_{t}+{\bm{e}}_{h}^{T}{\bm{h}}_{t-1}+{\bm{e}}_{m}^{T}{\bm{m}}_{t-1},$
(15)
$\displaystyle{\bm{m}}_{t}=\bar{{\bm{A}}}{\bm{m}}_{t-1}+\bar{{\bm{B}}}u_{t},$
(16)
$\displaystyle{\bm{h}}_{t}=f({\bm{W}}_{x}{\bm{x}}_{t}+{\bm{W}}_{h}{\bm{h}}_{t-1}+{\bm{W}}_{m}{\bm{m}}_{t}),$
(17)

where $\bar{{\bm{A}}}$ and $\bar{{\bm{B}}}$ are the discretized
versions333Using zero-order hold and $dt=1$, exact discretization gives
$\bar{{\bm{A}}}=e^{{\bm{A}}}$ and
$\bar{{\bm{B}}}={\bm{A}}^{-1}(e^{{\bm{A}}}-{\bm{I}}){\bm{B}}$. of ${\bm{A}}$
and ${\bm{B}}$. These matrices are usually frozen during training, although
they need not be. The input to the LTI system, $u(t)$, is computed by
projecting the input to the RNN cell, ${\bm{x}}_{t}\in R^{d_{x}}$, the hidden
state, ${\bm{h}}_{t}\in\mathbb{R}^{d_{h}}$, and the memory state,
${\bm{m}}_{t}\in\mathbb{R}^{d}$, onto their respective encoding weights
$({\bm{e}}_{x},{\bm{e}}_{h},\text{and }{\bm{e}}_{m})$. The memory state,
hidden state and the input to the cell are then combined using the weight
matrices $({\bm{W}}_{x},{\bm{W}}_{h},\text{and }{\bm{W}}_{m})$ and passed
through the nonlinearity, $f$. The encoding vectors and the kernel matrices
are learned during training.

###  3.3 Our Model

In this paper, we propose a modified LMU architecture by making two changes to
the equations defined above. First, we remove certain connections in (15) and
(17), namely $({\bm{e}}_{h}^{T},{\bm{e}}_{m}^{T})$ and $({\bm{W}}_{h})$; this
is done to aid parallelization. Second, instead of projecting the input to the
RNN cell, ${\bm{x}}_{t}$ down to a scalar as in (15), we implement a general
affine transformation followed by an element-wise nonlinearity; the original
architecture is better suited for dealing with 1D or low-dimensional inputs,
and this is a straightforward generalization of the encoding equation for
higher dimensional inputs. Additionally, adding a bias term to equation (17),
we end up with the following:

$\displaystyle{\bm{u}}_{t}=f_{1}({\bm{U}}_{x}{\bm{x}}_{t}+{\bm{b}}_{u}),$ (18)
$\displaystyle{\bm{m}}_{t}=\bar{{\bm{A}}}{\bm{m}}_{t-1}+\bar{{\bm{B}}}{\bm{u}}_{t},$
(19)
$\displaystyle{\bm{o}}_{t}=f_{2}({\bm{W}}_{m}{\bm{m}}_{t}+{\bm{W}}_{x}{\bm{x}}_{t}+{\bm{b}}_{o}).$
(20)

Note that equations (18) and (20) are equivalent to having time-distributed
dense layers before and after equation (19). In general, we expect our model
to be modified and used in combination with other feed-forward layers,
including self-attention. For example, we found a gated architecture
(Srivastava et al., 2015) where equation (18) is modified to

$\displaystyle{\bm{u}}_{t}=f_{1}({\bm{W}}_{u}{\bm{x}}_{t}+{\bm{b}}_{u})\cdot{\bm{g}}+{\bm{x}}_{t}\cdot\left(1-{\bm{g}}\right),$

to work well for the addition problem (results not shown). The gate is defined
as ${\bm{g}}=\sigma({\bm{W}}_{g}{\bm{x}}_{t}+{\bm{b}}_{g})$, i.e., a sigmoid-
activated affine transformation where the bias is initialized to -1.

Additionally, now that the input to the LTI system in this case is a vector,
${\bm{u}}_{t}\in\mathbb{R}^{1\times d_{u}}$, we have that
${\bm{m}}_{t}\in\mathbb{R}^{d\cdot d_{u}}$, and equation (16) can be thought
of as implementing the following:

$\displaystyle{\bm{m}}_{t}=\text{reshape}(\bar{{\bm{A}}}\text{
reshape}({\bm{m}}_{t-1},(d,\,d_{u}))+\bar{{\bm{B}}}{\bm{u}}_{t},\,\,d\cdot
d_{u}).$ (21)

##### Parallel Training

One of the motivations for the above mentioned architectural changes is that
the model now has only one recurrent connection: ${\bm{m}}_{t}$’s dependence
on itself from the past in equation (19). But because this is an LTI system,
standard control theory (Åström & Murray, 2010) gives us a non-iterative way
of evaluating this equation as shown below444Restricting ourselves to 1D input
for the purposes of illustration.

$\displaystyle{\bm{m}}_{t}=\sum_{j=1}^{t}\bar{{\bm{A}}}^{t-j}\bar{{\bm{B}}}u_{j}.$
(22)

Defining
${\bm{H}}=\begin{bmatrix}\bar{{\bm{A}}}^{0}\bar{{\bm{B}}}&\bar{{\bm{A}}}\bar{{\bm{B}}}&\ldots\end{bmatrix}\in\mathbb{R}^{d\times
n}$ and

$\displaystyle{\bm{U}}=\begin{bmatrix}u_{1}&u_{2}&u_{3}&\ldots&u_{n}\\\
&u_{1}&u_{2}&\ldots&u_{n-1}\\\ &&u_{1}&\ldots&u_{n-2}\\\ &&&\ddots&\vdots\\\
&&&&u_{1}\end{bmatrix}\in\mathbb{R}^{n\times n},$ (23)

the above convolution equation can alternatively be written as a matrix
multiplication:

$\displaystyle{\bm{m}}_{1:n}={\bm{H}}{\bm{U}},$ (24)

where $n$ is the sequence length. Given that ${\bm{H}}$ is the impulse
response of the LTI system, in practice we compute ${\bm{H}}$ by feeding in an
impulse to the RNN version of the DN (equation (19)). Note that in our method
the $\bar{{\bm{A}}}$ and $\bar{{\bm{B}}}$ matrices are frozen during training,
so the impulse response need only be computed once. In case of multi-
dimensional inputs, we would repeat the above computation several times, once
for each dimension of the input. It is also evident from the structure of the
${\bm{U}}$ matrix that although this reformulation turns the DN into a
feedforward layer, it still respects causality. In other words, the state
${\bm{m}}_{t}$ depends only on the inputs seen until that point of time, i.e.,
$u_{i}:i\leq t$.

Table 1: Complexity per layer and minimum number of sequential operations of various architectures. $n$ is the sequence length, $d_{x}$ is the input dimension, $d$ is the order, and $k$ is the size of the kernel. First three rows are as reported in Vaswani et al. (2017). Layer Type | Complexity | Sequential Ops
---|---|---
RNN | $O(n\cdot d^{2}_{x})$ | $\surd$
Convolution | $O(k\cdot n\cdot d_{x}^{2})$ | $\times$
Attention | $O(n^{2}\cdot d_{x})$ | $\times$
DN (19) | $O(n\cdot d^{2}\cdot d_{x})$ | $\surd$
DN (24) | $O(n^{2}\cdot d\cdot d_{x})$ | $\times$
DN (25) | $O(n\cdot d\cdot d_{x})$ | $\times$
DN (26) | $O(n\cdot\log{n}\cdot d\cdot d_{x})$ | $\times$

##### Complexity

With the new formulation, we immediately see that it is computationally (i.e.,
in terms of the number of operations) advantageous in situations where we only
need the final state (`return_sequences=False` in Keras terminology). Instead
of using (19) to explicitly simulate the first $n-1$ steps in-order to compute
the state at the time-step $n$, we can instead just compute

$\displaystyle{\bm{m}}_{n}={\bm{H}}{\bm{U}}_{:n},$ (25)

thus reducing the complexity of computing the final state, ${\bm{m}}_{n}$,
from $O(n\cdot d^{2}\cdot d_{x})$ to $O(n\cdot d\cdot d_{x})$, where $n$ is
the sequence length, $d$ is the order, and $d_{x}$ is the dimension of the
input. We show in Section (4.6) that using this implementation results in up
to 200x speedup.

The more general computation (24), although parallelizable, results in a
complexity of $O(n^{2}\cdot d\cdot d_{x})$. This can be made more efficient by
employing the convolution theorem which gives us an equivalent way of
evaluating the convolution in the Fourier space as555Assuming a padded Fourier
transform across the appropriate axis and automatic broadcasting when
computing element-wise multiplication. See LMUFFT code for details.

$\displaystyle{\bm{m}}_{1:n}=\mathcal{F}^{-1}\\{\mathcal{F}\\{{\bm{H}}\\}\cdot\mathcal{F}\\{{\bm{U}}_{:n}\\}\\}.$
(26)

Thanks to the fast Fourier transform, the above operation has a complexity of
$O(n\cdot\log_{2}{n}\cdot d\cdot d_{x})$. These, other algorithms, and their
complexities are reported in Table (1).

It was argued in Vaswani et al. (2017) that a self-attention layer is cheaper
than an RNN when the representation dimension of the input, $d_{x}$, is much
greater than the length of the sequence, $n$, which is seen in NLP
applications. For example, standard word or sub-word based machine translation
systems work with sequence lengths of about 100 and representation dimension
ranging from 300 to 512. The same argument holds for the DN layer too, since
we have found the inequality $\log_{2}n\cdot d\ll d_{x}$ to hold in all our
word-based NLP experiments – this excludes text8, which works at the level of
characters. More specifically, we use $d\leq 4$ for our word-based models.

##### Recurrent Inference

Machine learning algorithms are usually optimized for training rather than
deployment (Crankshaw, 2019), and because of that models need to be modified,
sometimes non-trivially, to be more suitable for inference. For example, one
of the metrics that is crucial for applications such as online Automatic
Speech Recognition (ASR) is low latency. Transformers have shown faster
training times and better performance than RNNs on ASR, but since they employ
(global) self-attention, which requires the entire input to begin processing,
they are not natural fits for such a task (Wang et al., 2020). Look-ahead
(Zhang et al., 2020) and chuck-based (Miao et al., 2020) approaches are
usually used to alter the architecture of self-attention for such purposes.
While our model can be trained in parallel, because of the equivalence of
equations (19) and (26), it can also be run in an iterative manner during
inference, and hence can process data in an online or streaming fashion during
inference.

##  4 Experiments

In the following experiments we compare our model against the LMU, LSTMs and
transformers. With these experiments, we focus on benchmarking rather than
establishing new state-of-the-art results. Hence, we stick to simple
architectures, constrain ourselves to train all the models, with the exception
of text8, using the Adam optimizer (Kingma & Ba, 2014) with all the default
parameter settings. For text8, we found it helpful to reduce the learning rate
by a factor of 10 halfway into training. Although unusual, we use the default
optimization settings even when transfer learning.

With models we compare against, we use results found in published work, even
when they make use of more sophisticated architectures, learning schedules or
regularization schemes.

We do not consider the capacity task from the original LMU paper here as they
employ just the delay layer without nonlinear dynamics in order to deal with
this problem, which makes their architecture essentially the same as ours.

###  4.1 psMNIST

The permuted sequential MNIST (psMNIST) dataset was introduced by Le et al.
(2015) to test the ability of RNNs to model complex long term dependencies.
psMNIST, as the name suggests, is constructed by permuting and then flattening
the $(28\times 28)$ MNIST images. The permutation is chosen randomly and is
fixed for the duration of the task, and in order to stay consistent, we use
the same permutation as Chandar et al. (2019) and Voelker et al. (2019). The
resulting dataset is of the form (samples, 784, 1), which is fed into a
recurrent network sequentially, one pixel at a time. We use the standard
50k/10k/10k split.

##### Architecture

As was pointed out in Voelker et al. (2019), in order to facilitate fair
comparison, RNN models being tested on this task should not have access to
more the $28^{2}=784$ internal variables. Keeping that in mind, and in order
to make direct comparisons to the original LMU model, we consider a model with
$d=468$ dimensions for memory. We set the dimension of the output state to
$346$ and use $\theta=784$. Our model uses 165k parameters, the same as all
the models reported in Table (3), except for the original LMU model, which
uses 102k parameters, and the HiPPO-LegS model, which is reported to use 512
hidden dimensions (number of parameters is unknown).

##### Results & Discussion

Test scores of various models on this dataset are reported in Table (3). Our
model not only surpasses the LSTM model, but also beats the current state-of-
the result of 98.3% set by HiPPO-LegS (Gu et al., 2020) recently. Thus, our
model sets a new state-of-the art result for RNNs of 98.49% on psMNIST. It is
interesting that our model, despite being simpler than the original LMU,
outperforms it on this dataset. This suggests that the main advantage of the
LMU over past models is the quality of its temporal memory, implemented by the
DN.

Table 2:  psMNIST results. The first three rows are from Voelker et al. (2019), and the fourth row is from Gu et al. (2020). Model | Accuracy
---|---
LSTM | 89.86
NRU | 95.38
LMU | 97.15
HiPPO-LegS | 98.3
Our Model | 98.49

Table 3: Mackey-Glass results. Model | NRMSE
---|---
LSTM | 0.059
LMU | 0.049
Hybrid | 0.045
Our Model | 0.044

###  4.2 Mackey-Glass

Mackey-Glass equations are a set of two nonlinear differential equations that
were originally developed to model the quantity of mature cells in the blood.
The second of these equations is interesting because it can result in chaotic
attractors and is used to construct the dataset at hand. This is a time-series
prediction task where we need to predict 15 time-steps into the future.

##### Architecture

Voelker & Eliasmith (2018) compare three architectures on this dataset. The
first one uses 4 layers of LSTMs $(h=28)$, the second one replaces LSTMs with
LMUs $(d=4,\theta=4)$, and the final one replaces the first and third layers
in the first model with LMUs. We use a relatively simple architecture where we
combine our model (single layer) with an additional dense layer. We use
$d=40$, $\theta=50$, and 1 and 140 units in the input and output layers, with
the additional dense layer containing 80 units. We did not try other
variations. All the models contain about 18k parameters and are run for 500
epochs.

##### Results & Discussion

NRMSE scores on the test set are presented in Table (3). We see that our model
outperforms the other three models in accuracy and training time. In the
original paper Voelker & Eliasmith (2018) hypothesize that the superior
performance of the hybrid model is due to the presence of gates, but given
that our model lacks gating mechanisms, we think that it might have to do with
the LMU model being misparametrized.

###  4.3 Sentiment, Semantics and Inference

In this section, we explore the supervised and semi-supervised capabilities of
our model. More specifically, we first look at the tasks of sentiment analysis
(IMDB), semantic similarity (QQP) and natural language inference (SNLI), and
then improve upon the IMDB score using a language model pre-trained on the
(unlabelled) Amazon Reviews dataset (Ni et al., 2019).

#### Supervised

##### IMDB

The IMDB dataset (Maas et al., 2011) is a standard sentiment classification
task containing a collection of 50k highly polar reviews, with the training
and testing sets containing 25k reviews each. We use the standard pre-
processed dataset available from the Keras
website,666https://keras.io/api/datasets/imdb/. consider a vocabulary of 20k
words, and set the maximum sequence length to 500 words.

##### QQP

For Quora Question Pairs (QQP), given a pair of sentences, the task is to
identify whether the two sentences are semantically similar. In this case, we
experiment on two train/dev/test splits: 390k/8k/8k like in Shen et al.
(2018), and 280k/80k/40k like in Sharma et al. (2019). We use a vocabulary of
20k words and truncate the sentences to be less than 25 words.

##### SNLI

The Stanford Natural Language Inference777 Dataset and published results are
available at https://nlp.stanford.edu/projects/snli/. was released to serve as
a benchmark for evaluating machine learning systems on the task of natural
language inference. Given a premise, the task is to determine whether a
hypothesis is true (entailment), false (contradiction), or independent
(neutral). We use the standard 550k/10k/10k split, consider a vocabulary of
20k words, and set the maximum sequence length to 25 words.

##### Architecture

In our experiments, confusingly, we found the use of the DN, without any
nonlinearities, to work well. Therefore, we construct parameter efficient
models that employ just the DN layer, with $d=1$ and $\theta=\texttt{maxlen}$.
We use 300D Glove embeddings (840B Common Crawl; Pennington et al. (2014)) for
all our models. For the IMDB task, which is a single sentence task, we encode
the sentence and pass the encoded vector to the final classification layer.
For two-sentence tasks, QQP and SNLI, we encode the two sentences to produce
two vectors, and then pass the vector obtained by concatenating the two
vectors, their absolute difference, and their element-wise product to the
final classification layer. We compare our models against the LSTM models
described in (Gu et al., 2020) for IMDB, both (Shen et al., 2018) and (Sharma
et al., 2019) for QQP, and (Bowman et al., 2015) for SNLI. They all use at
least an order of magnitude more parameters than our models.

##### Results & Discussion

We present the results from the three experiments in Table (4). As we can see,
our simple models based on the DN alone do indeed outperform the LSTM models.
It is also noteworthy that our models use significantly fewer parameters than
the LSTM models: 160x on IMDB, 650x on QQP and 60x on SNLI.

Table 4: IMDB, QQP and SNLI results. IMDB result is from (Gu et al., 2020), QQP results are from (Shen et al., 2018) and (Sharma et al., 2019) respectively, and SNLI is from (Bowman et al., 2015). | LSTM | Our Model
---|---|---
| Param. | Acc. | Param. | Acc.
IMDB | 50k | 87.29 | 301 | 89.10
QQP | -/800k | 82.58/81.4 | 1201 | 86.95/85.36
SNLI | 220k | 77.6 | 3.6k | 78.85

#### Semi-Supervised

##### Amazon Reviews

NLP over the past couple of years has been dominated by transfer learning
approaches, where language models pre-trained on large corpora are used to
initialize general purpose fine-tuning architectures to help with various
downstream tasks. We therefore consider using a pre-trained language model on
the Amazon Reviews dataset to improve the previously achieved score on IMDB.
Here, we make direct comparisons to the LSTM model described in Radford et al.
(2017). While they use the entire dataset (82 million reviews) for pre-
training and train their model for a month on 4 GPUs, due to resource
constraints, we use less than 5% (3.7 million reviews) and train our model for
about 12 hours on a single GPU. We use a vocabulary of 30k words.

##### Architecture

We have found the repeating block architecture, where each block is composed
of our model, a few highway layers, and a dense layer, when used with skip
connections to work well (see Figure in supplementary materials). Since the
inputs, $d_{x}$, are high dimensional, using a large $\theta$, which would
have to be accompanied by a large order $d$ to maintain resolution, would
result in vectors that are $d_{x}\cdot d$ dimensional, which is not desirable.
Therefore, we instead work with smaller $\theta$ (and hence smaller $d$) and
use several repeating blocks to take long-term dependencies into account. This
is similar to how convolutional networks are used: small kernel sizes but many
convolutional layers. If $\theta_{i}$ is the setting we use with the DN in the
ith block, then the effective delay, $\theta_{e}=\sum_{i}\theta_{i}$. In this
specific case, we have $\theta_{i}=6\,\,\forall\,\,i$, and $\theta_{e}=30$.
Thus, the model has access to the past 30 tokens when making a prediction. In
terms of the fine-tuning performance, like Peters et al. (2018), we found that
using deep representations, i.e., a linear combination of representations from
all the blocks, to be better than using just the output of the top block. For
fine-tuning, we compute a weighted sum of the outputs from the language model,
and use that to classify a given review as expressing a positive or negative
sentiment. Even during fine-tuning, we use the the Adam optimizer with all the
default settings. We did not experiment with smaller learning rates or other
learning schedules.

##### Results & Discussion

Results for this experiment are presented in Table (5). Despite using a much
smaller pre-training dataset, training for just 12 hours, and using less than
50% of the parameters, our model outperforms the LSTM model described in
Radford et al. (2017). We also include a self-attention based model,
DistilBert (Sanh et al., 2019), for comparison; it must be noted however that
DistilBert was trained on a more general yet much larger dataset
(concatenation of English WikiPedia and Toronto Book Corpus).

###  4.4 Language Modelling

Table 5: IMDB results with pre-training. First row is from Radford et al. (2017), and the second row is from Sanh et al. (2019). Model | # parameters (Millions) | Accuracy
---|---|---
LSTM | 75 | 92.88
DistilBERT | 66 | 92.82
Our Model | 34 | 93.20

##### text 8

We evaluate our model on character level language modelling using the text8
dataset888http://mattmahoney.net/dc/textdata.. It contains 100MB of clean text
from Wikipedia and has an alphabet size of 27. As is standard, we use the
first 90MB as the training set, the next 5MB as the validation set and the
final 5MB as the test set. Following Mikolov et al. (2012) and Zhang et al.
(2016), we set the sequence length to 180.

##### Architecture

This architecture is similar in essence to the language model used in the
above section, except for the use of deep representations; we simply work with
the output from the top block.999We did not test the use of deep
representations in this context. For text8, minor changes were made to adapt
the model to deal with longer sequences, which is a consequence of modelling
at the level of characters, and to parameter match it to the LSTM model in
Zhang et al. (2016), which uses around 3.2 million weights. We employ three
blocks in this case and use all three DNs with the setting $\theta=15$.

##### Results & Discussion

We report the scores in bits per character in Table (6). We see that our model
performs better than the LSTM model of similar size. A couple of observations
regarding model optimization: 1) we noticed that the training loss plateaus
around the twelfth epoch, so we found it helpful it to decrease the learning
rate by a factor of 10 around then (this is the only dataset for which we
alter the optimization parameters); 2) despite that, the training slows down
after a few more epochs, and we stop the optimization after 20 epochs; we
believe that a more carefully designed learning schedule might be able to help
with faster convergence. We also observed that using a self-attention layer
after the final block helps with generalization; this perhaps points to the
fact that the model needs a context that is longer 45 tokens in order to make
predictions.

###  4.5 Translation

##### IWSLT’15 En-Vi

IWSLT’15 English to Vietnamese is a medium-resource language translation task
containing 133k sentence pairs. Following Luong & Manning (2015), we do not
perform any pre-processing other than replacing words that occur less
frequently than 5 by the <unk> token, which leaves us with vocabulary sizes of
17k and 7.7k for English and Vietnamese, respectively. We use the TED tst2012
as the validation set and TED tst2013 as the test set. We use 300D
representation for source and target embeddings.

##### Architecture

For translation, we use a standard encoder-decoder architecture inspired by
the Amazon reviews language model, and we also employ an attention layer to
help with translation. Our model’s about the same size as the the 24 million
LSTM model described in Luong & Manning (2015). Due to time constraints, the
architecture and hyperparamters for this problem were relatively
underexplored.

##### Results & Discussion

Case sensitive BLEU scores are reported in Table (6). Our model brings in an
improvement of 2.3 BLEU over the LSTM model. When tested on lower-cased text
(without altering the training procedure), we obtained a higher score of 26.2
BLEU. One major limiting factor of this analysis is that we use a small, fixed
vocabulary (17k and 7.7k words), with no way of dealing with out-of-vocabulary
words. For future work, we intend to experiment with an open-vocabulary
encoding scheme such as byte pair encoding (Sennrich et al., 2015).

Table 6: Language modelling and translation results. text8 score is from Zhang et al. (2016), and IWSLT score is from Luong & Manning (2015). a(case sensitive), b(case insensitive). Model | text8 | IWSLT’15
---|---|---
LSTM | 1.65 | 23.3
Our Model | 1.61 | 25.5a/26.2b

###  4.6 Training Time

Here we explore the effect of parallelization on training time. We refer to
architectures that implement the DN using equation (19) as the LTI version and
the ones that use either (25) or (26), depending on whether `return_sequences`
is false or true, as the parallel version.

Results are presented in Figure (1). On the left, we compare the speedup we
obtain by going from the original LMU to our model, in both LTI and parallel
forms, on psMNIST and Mackey-Glass. We first notice that switching from the
LMU to the LTI version results in non-trivial gains. This is solely due to the
reduction in the number of recurrent connections. On top of this, switching to
the parallel version, owing to long sequences (784 and 5000) and, in the case
of psMNIST, a reduction in the computational burden, results in substantial
improvements in training times of 220x and 64x respectively.101010We note that
for Mackey-Glass we use a (parameter matched) one layer LMU model instead of
the 4 layer model used in Voelker et al. (2019); the difference with respect
to the original architecture is more drastic, with our model (parallel) being
about 200x faster.

On the right of Figure (1), we look at how varying the sequence length of the
inputs effects the time it takes to complete one epoch. We see that the LTI
version exhibits a linear growth whereas the parallel one essentially stays
constant.

Figure 1: (left) The plot shows the speedup we obtain from switching from the
LMU to the LTI (in orange) and parallel implementation (in blue) of our model.
(right) This show the effect of increasing the sequence length on the LTI and
parallel versions of our model. All results were measured using a single GTX
1080.

##  5 Conclusion

With the intention of alleviating the long standing problem of speeding up RNN
training, we introduce a variant of the LMU architecture that can process data
in a feedforward or sequential fashion. When implemented as a feedforward
layer, it can utilize parallel processing hardware such as GPUs and thus is
suitable for large scale applications, and when run as an RNN, it is useful in
applications where the amount of available memory is a limitation. We
demonstrate the effectiveness of this architecture by experimenting on a range
of tasks, of varying scale and difficulty. We also briefly consider the
question of computational complexity of our model, and argue for its
suitability to large scale applications in the domain of NLP, a direction we
will pursue in the future.

We note that our method of parallelization applies to all deep architectures
with linear recurrent dependencies, and although we focus on a specific LTI
system throughout, we hope that our analysis highlights the utility of linear
systems for the purposes of machine learning. In sum, we believe that linear
systems offer a scalable solution to many time series problems without
sacrificing accuracy.

## Acknowledgements

We would like to thank Aaron Voelker, Vincenzo Heska, Alison Shi, Andreas
Stockel, and Terry Stewart for discussions that helped improve this paper.
This work was supported by CFI and OIT infrastructure funding as well as the
Canada Research Chairs program, NSERC Discovery grant 261453, AFOSR grant
FA9550-17-1-0026 and OGS graduate funding.

## References

  * Åström & Murray (2010) Åström, K. J. and Murray, R. M.  _Feedback systems: an introduction for scientists and engineers_.  Princeton university press, 2010.
  * Balduzzi & Ghifary (2016) Balduzzi, D. and Ghifary, M.  Strongly-typed recurrent neural networks.  In _International Conference on Machine Learning_ , pp. 1292–1300. PMLR, 2016.
  * Blouw et al. (2020) Blouw, P., Malik, G., Morcos, B., Voelker, A. R., and Eliasmith, C.  Hardware aware training for efficient keyword spotting on general purpose and specialized hardware.  _arXiv preprint arXiv:2009.04465_ , 2020.
  * Bowman et al. (2015) Bowman, S. R., Angeli, G., Potts, C., and Manning, C. D.  A large annotated corpus for learning natural language inference.  _arXiv preprint arXiv:1508.05326_ , 2015.
  * Brogan (1991) Brogan, W. L.  _Modern control theory_.  Pearson education india, 1991.
  * Brown et al. (2020) Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al.  Language models are few-shot learners.  _arXiv preprint arXiv:2005.14165_ , 2020.
  * Chandar et al. (2019) Chandar, S., Sankar, C., Vorontsov, E., Kahou, S. E., and Bengio, Y.  Towards non-saturating recurrent units for modelling long-term dependencies.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pp. 3280–3287, 2019.
  * Crankshaw (2019) Crankshaw, D.  _The Design and Implementation of Low-Latency Prediction Serving Systems_.  PhD thesis, UC Berkeley, 2019.
  * Devlin et al. (2018) Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Gu et al. (2020) Gu, A., Dao, T., Ermon, S., Rudra, A., and Re, C.  Hippo: Recurrent memory with optimal polynomial projections.  _arXiv preprint arXiv:2008.07669_ , 2020.
  * Hochreiter & Schmidhuber (1997) Hochreiter, S. and Schmidhuber, J.  Long short-term memory.  _Neural computation_ , 9(8):1735–1780, 1997.
  * Kingma & Ba (2014) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Le et al. (2015) Le, Q. V., Jaitly, N., and Hinton, G. E.  A simple way to initialize recurrent networks of rectified linear units.  _arXiv preprint arXiv:1504.00941_ , 2015.
  * Luong & Manning (2015) Luong, M.-T. and Manning, C. D.  Stanford neural machine translation systems for spoken language domains.  In _Proceedings of the International Workshop on Spoken Language Translation_ , pp. 76–79, 2015.
  * Maas et al. (2011) Maas, A., Daly, R. E., Pham, P. T., Huang, D., Ng, A. Y., and Potts, C.  Learning word vectors for sentiment analysis.  In _Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies_ , pp. 142–150, 2011\.
  * Martin & Cundy (2017) Martin, E. and Cundy, C.  Parallelizing linear recurrent neural nets over sequence length.  _arXiv preprint arXiv:1709.04057_ , 2017.
  * Miao et al. (2020) Miao, H., Cheng, G., Gao, C., Zhang, P., and Yan, Y.  Transformer-based online ctc/attention end-to-end speech recognition architecture.  In _ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pp. 6084–6088. IEEE, 2020.
  * Mikolov et al. (2012) Mikolov, T., Sutskever, I., Deoras, A., Le, H.-S., Kombrink, S., and Cernocky, J.  Subword language modeling with neural networks.  _preprint (http://www. fit. vutbr. cz/imikolov/rnnlm/char. pdf)_ , 8:67, 2012.
  * Ni et al. (2019) Ni, J., Li, J., and McAuley, J.  Justifying recommendations using distantly-labeled reviews and fine-grained aspects.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_ , pp. 188–197, 2019.
  * Partington (2004) Partington, J. R.  Some frequency-domain approaches to the model reduction of delay systems.  _Annual Reviews in Control_ , 28(1):65–73, 2004\.
  * Pennington et al. (2014) Pennington, J., Socher, R., and Manning, C. D.  Glove: Global vectors for word representation.  In _Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP)_ , pp. 1532–1543, 2014.
  * Peters et al. (2018) Peters, M. E., Neumann, M., Iyyer, M., Gardner, M., Clark, C., Lee, K., and Zettlemoyer, L.  Deep contextualized word representations.  In _Proc. of NAACL_ , 2018.
  * Radford et al. (2017) Radford, A., Jozefowicz, R., and Sutskever, I.  Learning to generate reviews and discovering sentiment.  _arXiv preprint arXiv:1704.01444_ , 2017.
  * Radford et al. (2018) Radford, A., Narasimhan, K., Salimans, T., and Sutskever, I.  Improving language understanding by generative pre-training.  2018\.
  * Raffel et al. (2019) Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _arXiv preprint arXiv:1910.10683_ , 2019.
  * Sanh et al. (2019) Sanh, V., Debut, L., Chaumond, J., and Wolf, T.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter.  _arXiv preprint arXiv:1910.01108_ , 2019.
  * Sennrich et al. (2015) Sennrich, R., Haddow, B., and Birch, A.  Neural machine translation of rare words with subword units.  _arXiv preprint arXiv:1508.07909_ , 2015.
  * Sharma et al. (2019) Sharma, L., Graesser, L., Nangia, N., and Evci, U.  Natural language understanding with the quora question pairs dataset.  _arXiv preprint arXiv:1907.01041_ , 2019.
  * Shen et al. (2018) Shen, D., Wang, G., Wang, W., Min, M. R., Su, Q., Zhang, Y., Li, C., Henao, R., and Carin, L.  Baseline needs more love: On simple word-embedding-based models and associated pooling mechanisms.  _arXiv preprint arXiv:1805.09843_ , 2018.
  * Srivastava et al. (2015) Srivastava, R. K., Greff, K., and Schmidhuber, J.  Training very deep networks.  In _Advances in neural information processing systems_ , pp. 2377–2385, 2015.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  In _Advances in neural information processing systems_ , pp. 5998–6008, 2017.
  * Voelker et al. (2019) Voelker, A., Kajić, I., and Eliasmith, C.  Legendre memory units: Continuous-time representation in recurrent neural networks.  In _Advances in Neural Information Processing Systems_ , pp. 15544–15553, 2019.
  * Voelker & Eliasmith (2018) Voelker, A. R. and Eliasmith, C.  Improving spiking dynamical networks: Accurate delays, higher-order synapses, and time cells.  _Neural computation_ , 30(3):569–609, 2018.
  * Wang et al. (2020) Wang, C., Wu, Y., Liu, S., Li, J., Lu, L., Ye, G., and Zhou, M.  Reducing the latency of end-to-end streaming speech recognition models with a scout network.  _arXiv preprint arXiv:2003.10369_ , 2020.
  * Zhang et al. (2020) Zhang, Q., Lu, H., Sak, H., Tripathi, A., McDermott, E., Koo, S., and Kumar, S.  Transformer transducer: A streamable speech recognition model with transformer encoders and rnn-t loss.  In _ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pp. 7829–7833. IEEE, 2020.
  * Zhang et al. (2016) Zhang, S., Wu, Y., Che, T., Lin, Z., Memisevic, R., Salakhutdinov, R. R., and Bengio, Y.  Architectural complexity measures of recurrent neural networks.  In _Advances in neural information processing systems_ , pp. 1822–1830, 2016.

## Supplementary Materials

Figure 2: Illustration of the language model used for pre-training on the
Amazon Reviews dataset. Although the actual model uses five blocks
(combination of DN, Dense and Highway), we only show two blocks in the above
figure.


# Accelerated Motion Correction for MRI using Score-Based Generative Models

Brett Levac† Ajil Jalal⋆ Jonathan I. Tamir† † University of Texas, Austin, ECE
⋆University of California, Berkeley, EECS

###### Abstract

Magnetic Resonance Imaging (MRI) is a powerful medical imaging modality, but
unfortunately suffers from long scan times which, aside from increasing
operational costs, can lead to image artifacts due to patient motion. Motion
during the acquisition leads to inconsistencies in measured data that manifest
as blurring and ghosting if unaccounted for in the image reconstruction
process. Various deep learning based reconstruction techniques have been
proposed which decrease scan time by reducing the number of measurements
needed for a high fidelity reconstructed image. Additionally, deep learning
has been used to correct motion using end-to-end techniques. This, however,
increases susceptibility to distribution shifts at test time (sampling
pattern, motion level). In this work we propose a framework for jointly
reconstructing highly sub-sampled MRI data while estimating patient motion
using score-based generative models. Our method does not make specific
assumptions on the sampling trajectory or motion pattern at training time and
thus can be flexibly applied to various types of measurement models and
patient motion. We demonstrate our framework on retrospectively accelerated 2D
brain MRI corrupted by rigid motion.

Keywords: Magnetic Resonance Imaging (MRI), Motion Correction, Accelerated
Reconstruction, Joint Estimation, Score Models, Generative Models

##  1 Introduction

MRI is a powerful medical imaging modality that boasts superior soft tissue
contrast which can be leveraged to diagnose a variety of adverse health
conditions. Compared to other medical imaging methods, however, MRI is
relatively slow. Many previous works have developed methods to decrease the
data requirements for solving the ill-posed inverse problem of sub-sampled
image reconstruction [1, 2], including through deep learning [3, 4, 5, 6, 7].
Even with sub-sampling, MRI is highly susceptible to patient motion [8], which
leads to inconsistencies in measurement data at the time of reconstruction.
The image degradation due to motion-corrupt data can manifest as blurring or
ghosting which possibly render the image nondiagnostic, and often necessitates
reacquiring the scan. It has been reported that on average motion corrupted
scans in a pediatric setting have led some hospitals to incur nearly $\$365K$
in additional costs each year due to motion artifacts [9].

Many approaches to address motion have been proposed, including prospective
methods used at scan time like hard and soft gating. These methods measure
motion during the scan using motion navigators [10] or external measurement
devices such as respiratory bellows, nuclear magnetic resonance probes, and
optical tracking. While powerful, these approaches require modification to the
pulse sequence or additional hardware and can ultimately increase scan time.
Alternative approaches, which are considered retrospective methods,
incorporate the motion into the MRI forward model and solve a modified
reconstruction problem after the scan has been completed [11, 12, 13].

Recent developments in deep learning have allowed for end-to-end supervised
training methods in which a neural network (NN) is used to map motion corrupt
inputs to (ideally) motion free outputs [14]. These techniques have shown
great promise, however, recent work has displayed that the network’s
performance at test time is highly dependent on the motion simulation
parameters used at training time [15]. More recently, supervised learning has
been used as one step in a larger iterative algorithm that jointly solves for
the image and the motion parameters [16]. These methods have primarily been
applied to the low acceleration regime, and may still suffer from overfitting
to the specific motion simulation at training time.

In parallel, promising recent work in deep learning based image reconstruction
has used generative modeling to learn image priors separate from the
measurement forward model [17, 5]. By explicitly decoupling the forward model
and the image prior, generative model based techniques have an increased
robustness to test time shifts in the forward operator such as sampling
trajectory [5]. As a result, the reconstruction can flexibly handle other
modifications to the forward model without retraining the generative prior.

The aim of this work is to extend recent advances in generative models to a
joint estimation technique which can correct motion at high acceleration rates
by leveraging recent advances in score models for inverse problems and their
robustness to forward model shifts. We study the utility of our method in
reconstructing retrospectively sub-sampled multi-coil brain MRI corrupted by
simulated rigid motion.

##  2 Methods

###  2.1 Image Formation and Motion Simulation

We restrict our experiments to correcting in-plane, rigid body motion which is
a common assumption for motion correction techniques applied to brain imaging
[12, 13, 16]. It is known from basic Fourier theory that translations in image
space result in linear phase shifts in the Fourier domain (k-space), while
rotations in image space result in rotations in k-space. This allows us to
model patient motion as an alteration to the motion-free MRI forward operator.
The multi-coil MRI forward model with motion admits the following form:

$y=L_{\phi}N_{R_{\theta}K}Sx+w,w\sim{\cal N}(0,\sigma^{2}I).$ (1)

Here $x\in\mathbb{C}^{N}$ is the motion-free image, $S\in\mathbb{C}^{NC\times
N}$ contains the $C$ coil sensitivity maps, $y\in\mathbb{C}^{MC}$ is the
corrupted k-space measurements, $R_{\theta}$ denotes a set of rotation
matrices for each motion state angle ($\theta$) observed during the
acquisition, $L_{\phi}$ is a diagonal matrix which applies linear phase shifts
to acquired k-space depending upon the object’s translation ($\phi$) observed
during acquisition, $K$ defines a set of k-space coordinates which are the
intended trajectory of the acquisition, and $N_{R_{\theta}K}$ is a linear
function of the image $x$ which takes the Non-uniform Fast Fourier Transform
(NUFFT) of $Sx$ at the coordinates $R_{\theta}K$. Here we see that rotations
in image space are implemented by rotation of the corresponding coordinates in
the NUFFT, while translations are implemented with the corresponding phase
shift in k-space. For the remainder of this paper we will denote the forward
operation as:

$y=A_{\kappa}x+w,w\sim{\cal N}(0,\sigma^{2})$ (2)

where $A_{\kappa}$ includes the NUFFT and coil map weighting while it also
implicitly capturing all rotation and translation effects, at all times
throughout the scan, in $\kappa$. Our goal is to jointly estimate $x$ and
$\kappa$. We note that $y$ is a linear function of $x$ but is non-linear in
$\kappa$. This means that finding a global solution, even in the fully sampled
case, is not guaranteed due to the non-convexity introduced by estimating
$\kappa$.

Additionally, for the remainder of the paper we make the assumption, both when
simulating and estimating motion, that k-space points acquired in the same
repetition time (TR) are assumed to have the same motion state. This
assumption is well supported as in common sequences like fast spin echo (FSE),
timing between phase encode (PE) lines in a single echo train are on the order
of milliseconds while points in different TRs are separated on the order of
seconds [15]. In particular a fully sampled scan in our simulations has
$384\times 384$ k-space points and we choose our (fully sampled) sequence to
have an echo train length (ETL) of $8$ and $48$ TRs. We therefore simulate
$48$ different angles, x-translations, and y-translations in the fully sampled
case. In all experiments, scan acceleration ($R$) is applied along the TR
dimension. For example, this means that for $R=4$, $ETL=8$ and $TRs=12$
resulting in the simulation and estimation of $12$ different motion states,
corresponding to a 4x faster scan. We note that modeling motion across TRs is
not strictly necessary, and in addition the model can extended to non-rigid
motion by modifying the forward operator to include a parameterized
deformation field which can not be easily included in the description of the
NUFFT. In this case the motion state is expanded from three parameters to a
larger basis set.

###  2.2 Score Models

Recent work in deep learning has led to impressive progress in the area of
generative modeling [18, 19], and in particular, score-based generative models
[17]. Given i.i.d. training samples drawn from a high-dimensional probability
density $p$, score-based generative models can estimate the _score_ of $p$,
defined as $\nabla_{x}\log p(x)$, by training a NN via denoising score
matching [17]. These models can produce new samples from $p$ by initializing
$x_{0}\sim{\cal N}(0,I_{N})$, and running annealed Langevin dynamics for $T$
steps:

$\displaystyle x_{t+1}=x_{t}+\eta_{t}\nabla\log p(x_{t})+z_{t},$ (3)

where $z_{t}\sim{\cal N}(0,\,2\eta_{t}I_{N})$ is i.i.d. Gaussian. As
$\eta_{t}\to 0,T\to\infty$, the distribution of $x_{T}$ converges to $p$ [17].

If we know the distribution of $y|x$, Eqn (3) can be modified using Bayes’
rule to sample from the _posterior distribution_ $p(x|y)$, and this procedure
has been successful in inverse problems such as compressed sensing MRI. We
emphasize that during training, the score model is agnostic to the MRI forward
model, which enhances its robustness to forward model shifts when applied to
inverse problems [5].

###  2.3 Score-Based Accelerated Motion Correction

As stated above, prior works have shown positive results when using score
models to solve ill-posed inverse problems like sub-sampled image
reconstruction [5, 6, 7]. In prior work, however, the forward model is assumed
to be fixed and known throughout the reconstruction procedure. In our work,
however, we assume that our forward operator $A_{\kappa}$ belongs to a
restricted class of operators with unknown parameters $\kappa$ which must be
learned during the inference. Another way of viewing this problem is as
posterior sampling from the joint distribution $p(x,\kappa|y)$. Under this
joint posterior, we arrive at the following update equations:

$\displaystyle x_{t+1}$
$\displaystyle=x_{t}+\eta_{t}\nabla_{x_{t}}\log{p(x_{t},\kappa_{t}|y)}+z_{t},$
(4) $\displaystyle\kappa_{t+1}$
$\displaystyle=\kappa_{t}+\lambda_{t}\nabla_{\kappa_{t}}\log{p(x_{t},\kappa_{t}|y)}+\xi_{t},$
(5)

where $z_{t}\sim{\cal N}(0,2\eta_{t}I),\xi_{t}\sim{\cal N}(0,2\lambda_{t}I)$
are i.i.d. Gaussian.

We assume independence between $x$ and $\kappa$ and use Bayes’ rule to further
simplify the update equations to

$\displaystyle x_{t+1}$
$\displaystyle=x_{t}+\eta_{t}\nabla_{x_{t}}\left(\log{p(y|x_{t},\kappa_{t})}+\log{p(x_{t})}\right)+z_{t},$
(6) $\displaystyle\kappa_{t+1}$
$\displaystyle=\kappa_{t}+\lambda_{t}\nabla_{\kappa_{t}}\left(\log{p(y|x_{t},\kappa_{t})}+\log{q(\kappa_{t})}\right)+\xi_{t}.$
(7)

Inserting the likelihood function given by the forward model in Eqn (2) we get
the following final update steps:

$\displaystyle x_{t+1}$
$\displaystyle=x_{t}+\eta_{t}\left(\frac{A_{\kappa_{t}}^{H}(y-A_{\kappa_{t}}x_{t})}{\sigma^{2}}+f(x_{t})\right)+z_{t},$
(8) $\displaystyle\kappa_{t+1}$
$\displaystyle=\kappa_{t}+\lambda_{t}\nabla_{\kappa_{t}}\left(\|y-A_{\kappa_{t}}x\|^{2}_{2}+\log{q(\kappa_{t})}\right)+\xi_{t},$
(9)

where $f(x_{t})$ is the score-model trained to estimate $\nabla\log p(x)$. We
note that this formulation allows for the incorporation of different priors
over the motion parameters $\kappa$.

###  2.4 Implementation Details

The score model ($f$) in Eq. (8) was trained on 14,539 T2-weighted brain
images from the fastMRI dataset [20]. Sensitivity maps were estimated using
ESPIRiT [21]. Simulated ground truth motion states ($\kappa_{GT}$) were chosen
as $\kappa_{\mathrm{GT}}\sim U(-2,2)$ i.i.d. for each TR. This level was
chosen based off the severity of the resulting artifacts which visually
matched motion-corrupt data from prior work [14] and led to nondiagnostic
image quality even at low accelerations. The prior over motion parameters used
in our score update steps was $\kappa\sim U(-15,15)$ so as not to assume
perfect knowledge over the prior distribution which generated the motion
parameters.

Although the use of the NUFFT operator in our forward model allows for any
non-Cartesian sampling trajectory, the results we present here are for
Cartesian trajectories with equispaced sub-sampling (Fig. 1).

Figure 1: (a) $R=4$ Sampling trajectory where each color denotes a separate
TR. 12 TRs total each having ETL$=8$. (b) corresponding motion simulation
(rotation, translation) for each TR (12 motion states in total).

We compared our method to two different L1-wavelet [2] reconstructions,
implemented using BART [22], where (1) uses motion corrupt data and (2) uses
motion free data. For prior methods which rely on SENSE based schemes [1] for
the final reconstruction [12, 13, 16] (1) represents the worst case scenario
while (2) represents the best case scenario. Our method, however, jointly
estimated motion parameters and the image. We performed experiments on 11
validation slices and 39 test slices from the T2-weighted brain fastMRI
dataset at retrospective acceleration factors of ($R=4,4.8,6$). Each slice was
independently corrupted by motion and additive Gaussian noise.

All reconstruction hyperparameters for the three methods were tuned on the
validation set and held fixed on the test set. Normalized root mean squared
error (NRMSE) was used to evaluate the performance of the methods. Our
algorithm was implemented in PyTorch and gradients for $\kappa_{t}$ in Eq. (9)
were calculated through auto-differentiation. Inference was run on one Nvidia
RTX 3090. As all training, validation, and test data was taken from a publicly
available dataset [20] and motion was retrospectively simulated, ethical
approval was not required.

##  3 Results

Quantitative NRMSE metrics are shown for each method in Table 1 at
accelerations of ($R=4,4.8,6$). Our reconstruction often converges to motion
estimates with an offset from the ground-truth due to the ambiguity between
two data-consistent reconstructions with a fixed motion offset. We therefore
align our reconstructions to the motion-free ground-truth before evaluating
NRMSE. Example reconstructions for $R=4.8$ and $R=6$ are shown in Fig. 3. We
visualize in Fig. 2 the intermediate reconstruction images along with the
estimated motion trajectories at various points during inference.

Table 1: Reconstruction NRMSE averaged over the test set. | R$=4$ | R$=4.8$ | R$=6$
---|---|---|---
| L1 Wavelet
---
Motion Corrupt
$0.343$ | $0.345$ | $0.363$
| L1 Wavelet
---
Motion Free
$0.092$ | $0.115$ | $0.158$
Ours | $\mathbf{0.090}$ | $\mathbf{0.104}$ | $\mathbf{0.145}$
Figure 2: (Top) Image progression as inference proceeds. (Bottom) Motion
estimation (rotation, x-translation, y-translation) as inference proceeds. In
both figures each _Step_ indicates a different noise level in the Langevin
sampling inference.

##  4 Discussion

From our results we see that, _even with no motion_ , L1-wavelet based
reconstruction methods hit a performance limit at high 2D acceleration
factors, in agreement with previously reported literature [3, 4, 5]. In
contrast, our method, _which was blind to the motion_ , is able to both
estimate the motion trajectories while also reconstructing a sub-sampled image
with high fidelity (Fig. 3). Although the performance of score models in the
presence of corrupted forward operators has not been displayed previously, we
showed that by viewing the inference procedure as posterior sampling jointly
over the image and motion parameters we can achieve high fidelity motion
correction at large acceleration rates for retrospectively corrupted data.

We note that the training and inference processes for solving inverse problems
with score models are unique when compared to end-to-end supervised methods.
Score models are trained only on ground-truth images which, in part, makes the
inference process agnostic to shifts in sampling pattern, and motion
distribution. On the other hand, end-to-end methods are known to degrade when
the test time distributions do not align with the training distributions,
which would potentially require re-training for different motion corruption
[15]. Although we only simulated and estimated one motion state per TR, this
can be relaxed to estimate motion states for every PE at the cost of larger
computational burden. However, assuming that motion occurs only between PEs
neglects higher order effects like magnetization history which is left for
future work.

Figure 3: Reconstruction results for L1 Wavelet (with and without motion) and
our method for (A) $R=4.8$, and (B) $R=6$ acceleration factors.

##  5 Conclusion

Our proposed method of jointly optimizing motion-corrupt MRI with score-based
generative models provides high fidelity images compared to even the best case
scenario (motion free measurements) for L1 Wavelet, which is similar to
previously proposed joint motion correction techniques [12, 13, 16] but at
higher acceleration levels. This retrospective study, however, did not compare
to end-to-end supervised methods which are becoming popular motion correction
techniques [14], but are also not typically not used at high accelerations,
nor was it applied to prospectively acquired motion-corrupt MRI.

##  6 Acknowledgments

This work was supported by NSF IFML 2019844, NIH U24EB029240, and Oracle for
Research Fellowship.

## References

  * [1] Klaas P. Pruessmann, Markus Weiger, Markus B. Scheidegger, and Peter Boesiger,  “Sense: Sensitivity encoding for fast mri,”  Magnetic Resonance in Medicine, vol. 42, no. 5, pp. 952–962, 1999\.
  * [2] Michael Lustig, David Donoho, and John M Pauly,  “Sparse mri: The application of compressed sensing for rapid mr imaging,”  Magnetic Resonance in Medicine: An Official Journal of the International Society for Magnetic Resonance in Medicine, vol. 58, no. 6, pp. 1182–1195, 2007.
  * [3] Kerstin Hammernik, Teresa Klatzer, Erich Kobler, Michael P. Recht, Daniel K. Sodickson, Thomas Pock, and Florian Knoll,  “Learning a variational network for reconstruction of accelerated mri data,”  Magnetic Resonance in Medicine, vol. 79, no. 6, pp. 3055–3071, 2018\.
  * [4] Hemant K. Aggarwal, Merry P. Mani, and Mathews Jacob,  “Modl: Model-based deep learning architecture for inverse problems,”  IEEE Transactions on Medical Imaging, vol. 38, no. 2, pp. 394–405, 2019.
  * [5] Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jonathan I Tamir,  “Robust compressed sensing mri with deep generative priors,”  Advances in Neural Information Processing Systems, 2021.
  * [6] Guanxiong Luo, Martin Heide, and Martin Uecker,  “MRI reconstruction via data driven markov chain with joint uncertainty estimation,”  arXiv preprint arXiv:2202.01479, 2022.
  * [7] Hyungjin Chung and Jong Chul Ye,  “Score-based diffusion models for accelerated mri,”  Medical Image Analysis, vol. 80, pp. 102479, 2022.
  * [8] Maxim Zaitsev, Julian Maclaren, and Michael Herbst,  “Motion artifacts in mri: A complex problem with many partial solutions,”  Journal of Magnetic Resonance Imaging, vol. 42, no. 4, pp. 887–901, 2015.
  * [9] Jakob M Slipsager, Stefan L Glimberg, Jes Søgaard, Rasmus R Paulsen, Helle H Johannesen, Pernille C Martens, Alka Seth, Lisbeth Marner, Otto M Henriksen, Oline V Olesen, et al.,  “Quantifying the financial savings of motion correction in brain mri: a model-based estimate of the costs arising from patient head motion and potential savings from implementation of motion correction,”  Journal of Magnetic Resonance Imaging, vol. 52, no. 3, pp. 731–738, 2020.
  * [10] Joseph Y. Cheng, Tao Zhang, Nichanan Ruangwattanapaisarn, Marcus T. Alley, Martin Uecker, John M. Pauly, Michael Lustig, and Shreyas S. Vasanawala,  “Free-breathing pediatric mri with nonrigid motion correction and acceleration,”  Journal of Magnetic Resonance Imaging, vol. 42, no. 2, pp. 407–420, 2015.
  * [11] Sajan Goud Lingala, Edward DiBella, and Mathews Jacob,  “Deformation corrected compressed sensing (dc-cs): A novel framework for accelerated dynamic mri,”  IEEE Transactions on Medical Imaging, vol. 34, no. 1, pp. 72–85, 2015.
  * [12] Lucilio Cordero-Grande, Rui Pedro A. G. Teixeira, Emer J. Hughes, Jana Hutter, Anthony N. Price, and Joseph V. Hajnal,  “Sensitivity encoding for aligned multishot magnetic resonance reconstruction,”  IEEE Transactions on Computational Imaging, vol. 2, no. 3, pp. 266–280, 2016.
  * [13] Melissa W Haskell, Stephen F Cauley, and Lawrence L Wald,  “Targeted motion estimation and reduction (TAMER): data consistency based motion mitigation for mri using a reduced model joint optimization,”  IEEE transactions on medical imaging, vol. 37, no. 5, pp. 1253–1265, 2018.
  * [14] Patricia M Johnson and Maria Drangova,  “Conditional generative adversarial network for 3d rigid-body motion correction in mri,”  Magnetic Resonance in Medicine, vol. 82, no. 3, pp. 901–910, 2019\.
  * [15] Brett Levac, Sidharth Kumar, Sofia Kardonik, and Jonathan I. Tamir,  “Fse compensated motion correction for mri using data driven methods,”  in Medical Image Computing and Computer Assisted Intervention – MICCAI 2022, Cham, 2022, pp. 707–716, Springer Nature Switzerland.
  * [16] Melissa W Haskell, Stephen F Cauley, Berkin Bilgic, Julian Hossbach, Daniel N Splitthoff, Josef Pfeuffer, Kawin Setsompop, and Lawrence L Wald,  “Network accelerated motion estimation and reduction (NAMER): convolutional neural network guided retrospective motion correction using a separable motion model,”  Magnetic resonance in medicine, vol. 82, no. 4, pp. 1452–1461, 2019\.
  * [17] Yang Song and Stefano Ermon,  “Generative modeling by estimating gradients of the data distribution,”  Advances in Neural Information Processing Systems, pp. 11918–11930, 2019.
  * [18] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio,  “Generative adversarial networks,”  Communications of the ACM, vol. 63, no. 11, pp. 139–144, 2020.
  * [19] Diederik P Kingma and Max Welling,  “Auto-encoding variational bayes,”  arXiv preprint arXiv:1312.6114, 2013.
  * [20] Florian Knoll, Jure Zbontar, Anuroop Sriram, Matthew J Muckley, Mary Bruno, Aaron Defazio, Marc Parente, Krzysztof J Geras, Joe Katsnelson, Hersh Chandarana, et al.,  “fastmri: A publicly available raw k-space and dicom dataset of knee images for accelerated mr image reconstruction using machine learning,”  Radiology. Artificial intelligence, vol. 2, no. 1, 2020.
  * [21] Martin Uecker, Peng Lai, Mark J. Murphy, Patrick Virtue, Michael Elad, John M. Pauly, Shreyas S. Vasanawala, and Michael Lustig,  “Espirit—an eigenvalue approach to autocalibrating parallel mri: Where sense meets grappa,”  Magnetic Resonance in Medicine, vol. 71, no. 3, pp. 990–1001, 2014\.
  * [22] Martin Uecker, Jonathan I Tamir, Frank Ong, and Michael Lustig,  “The bart toolbox for computational magnetic resonance imaging,”  in Proc Intl Soc Magn Reson Med, 2016, vol. 24.