Buchholz

# Fourier Image Transformer

Tim-Oliver Buchholz  CSBD and MPI-CBG, Dresden, Germany  Florian Jug
Fondatione Human Technopole, Milano, Italy

###### Abstract

Transformer architectures show spectacular performance on NLP tasks and have
recently also been used for tasks such as image completion or image
classification. Here we propose to use a sequential image representation,
where each prefix of the complete sequence describes the whole image at
reduced resolution. Using such Fourier Domain Encodings (FDEs), an auto-
regressive image completion task is equivalent to predicting a higher
resolution output given a low-resolution input. Additionally, we show that an
encoder-decoder setup can be used to query arbitrary Fourier coefficients
given a set of Fourier domain observations. We demonstrate the practicality of
this approach in the context of computed tomography (CT) image reconstruction.
In summary, we show that Fourier Image Transformer (FIT) can be used to solve
relevant image analysis tasks in Fourier space, a domain inherently
inaccessible to convolutional architectures.

###### keywords:

Transformer | Fourier Space | Image Processing

<EMAIL_ADDRESS><EMAIL_ADDRESS>
##  1 Introduction

Transformer architectures are currently setting new standards on virtually all
natural language processing (NLP) tasks (1, 2). But also in other domains,
transformers can find application (3). Vision Transformers (ViT) (4), for
example, show SOTA results on image classification tasks. In the future we
expect an increasing number of problems to be solved with transformers and
here we also contribute such applications.

The key novelty of transformers is their self-attention mechanism (5),
allowing them to learn and utilize long ranging dependencies in data.
Recently, this mechanism is applied to longer and longer input sequences of
words or other elements, _e.g_. pixels (6).

\begin{overpic}[width=433.62pt,trim=0.0pt 50.18748pt 0.0pt
4.015pt]{figs/figure-fit_teaser.pdf} \end{overpic} Figure 1: We present
Fourier Image Transformers (FITs), realizations of Fast-Transformers (7), that
operate on images via a novel sequential image representation we call Fourier
Domain Encoding (FDE). We demonstrate the utility of FITs on two tasks. First,
an image super-resolution task (top row), where a given low-resolution image
is used to predict a higher resolution version. The second task is tomographic
reconstruction (bottom row), where a given set of projection images (a
sinogram) is transformed by an encoder-decoder FIT to improve the quality of
the reconstructed image w.r.t. filtered backprojection (FBP), a commonly used
method for tomographic reconstructions. Note that Fourier Image Transformers
solve both tasks in Fourier space, with the final results being obtained by an
inverse Fourier transformation (iFFT).

Initially, we wondered if pixel sequences are the only valid image
representation to train auto-regressive transformer models in the spirit of
(6). In particular, we want to use a representation where each prefix of such
a descriptive sequence encodes the full image at lower resolution. Hence, we
introduce Fourier Domain Encodings (FDEs), which do have this desired property
and, as we show, can successfully be used to train auto-regressive Fourier
Image Transformers (“FIT: SRes”) for super-resolution tasks.

Additionally, we show how an encoder-decoder based Fourier Image Transformer
(“FIT: TRec”) can be trained on a set of Fourier measurements and then used to
query arbitrary Fourier coefficients, which we use to improve sparse-view
computed tomography (CT) image restoration111Sparse-view CT image restoration
typically suffers from missing high frequency Fourier coefficients, leading to
unwanted image artefacts.. We demonstrate this by providing a given set of
projection Fourier coefficients to our encoder-decoder setup and use it to
predict Fourier coefficients at arbitrary query points. This allows us to
predict a dense, grid-sampled discrete Fourier spectrum of a high quality CT
reconstruction.

In Section 2 we review some of the related transformer and tomography
literature. In Section 3 we introduce our novel Fourier Domain Encoding (FDE)
and training strategies for auto-regressive and encoder-decoder transformer
models. In Section 4 we present our experiments and results on multiple
datasets.

##  2 Related Work

Transformer architectures are revolutionizing neural language processing
(NLP), replacing recurrent neural networks (RNNs) and long short-term memory
(LSTM) architectures in virtually all NLP tasks (1, 2). Our own work is based
on the seminal paper by Vaswani  _et al_. (5), which introduces the self-
attention mechanism that is key to the success of transformers. Since we work
on images, where input sequences tend to be long, we use fast-transformers, an
efficient approximation of softmax self-attention, as introduced by
Katharopoulos  _et al_. (7).

###  2.1 Attention is all you need

Vaswani  _et al_. (5) were the first to introduce transformers. More
specifically, they introduced an encoder-decoder structure, where the encoder
maps an input encoding $\mathbf{x}\in\mathbb{R}^{N\times F}$ into a continuous
latent space $\mathbf{z}\in\mathbb{R}^{N\times F}$, with $N$ corresponding to
the number of input tokens and $F$ representing the feature dimensionality per
token. This latent space embedding $\mathbf{z}$ is then given to the decoder,
which generates an $M$ long output sequence $\mathbf{y}\in\mathbb{R}^{M\times
F}$ iteratively, element by element. This auto-regressive decoding scheme
means that the decoder generated the $i$-th output token while not only
observing $\mathbf{z}$, but also all $i-1$ output tokens generated previously.

More formally, a transformer is a function $T:\mathbb{R}^{N\times
F}\xrightarrow{}\mathbb{R}^{N\times F}$, represented by $L$ transformer layers

$T_{l}(\mathbf{x})=f_{l}(A_{l}(\mathbf{x})+\mathbf{x}),$ (1)

with $A_{l}$ denoting a self-attention module and $f_{l}$ being a simple feed
forward network.

In the self-attention module, the input $\mathbf{x}$ is mapped to queries
$Q=\mathbf{x}W_{Q}$, keys $K=\mathbf{x}W_{K}$ and values $V=\mathbf{x}W_{V}$
by matrix multiplication with learned matrices $W_{Q}\in\mathbb{R}^{F\times
D}$, $W_{K}\in\mathbb{R}^{F\times D}$ and $W_{V}\in\mathbb{R}^{F\times F}$.
The self-attention output is then computed by

$\begin{split}A_{l}(\mathbf{x})&=\text{softmax}\left(\frac{QK^{T}}{\sqrt{D}}\right)V,\end{split}$
(2)

with the softmax-function being applied per row. Intuitively, the softmax-
normalized similarity between computed keys and queries is used to obtain the
weighted sum over the values.

Typically, instead of a single self-attention module, multi-head attention is
being used. If that is the case, a transformer layer $T_{l}$ learns multiple
$W_{Q}$, $W_{K}$ and $W_{V}$, allowing the layer to simultaneously perform
multiple attention-based computations (5).

Since transformers do not explicitly encode the relative position between
input tokens, positional encodings are required whenever specific input
topologies need to be made accessible to the transformer. In (5), a useful 1D
positional encoding scheme was proposed. Later, Wang  _et al_. (8) generalized
this scheme to 2D topologies. In our own work, we have adopted this encoding
scheme, but use it not only to encode integer pixel-grid locations, but
arbitrary coordinates.

###  2.2 Fast-Transformers

While the advantage of transcending beyond CNN’s localized receptive fields by
introducing global attention proves beneficial for many learning tasks, the
big downside is the computational cost associated with it. Due to the required
matrix multiplication $QK^{T}$, the self-attention on an input sequence of
length $N$ requires $\mathcal{O}(N^{2})$ memory and time. Katharopoulos  _et
al_. (7) propose to generalize the self-attention for the $i$-th row to

$A_{l}^{i}(\mathbf{x})=\frac{\sum_{j=1}^{N}\text{sim}(Q_{i},K_{j})V_{j}}{\sum_{j=1}^{N}\text{sim}(Q_{i},K_{j})},$
(3)

with sim being a non-negative similarity function, which includes all kernels
$k(x,y)=\langle\Phi(x),\Phi(y)\rangle$ with $k(x,y):\mathbb{R}^{2\times
F}\xrightarrow{}\mathbb{R}^{+}$ where $\Phi$ is the feature mapping of the
kernel. Using the feature mapping $\Phi$ Equation 3 can be written as

$A_{l}^{i}(\mathbf{x})=\frac{\Phi(Q_{i})^{T}\sum_{j=1}^{N}\Phi(K_{j})V_{j}}{\Phi(Q_{i})^{T}\sum_{j=1}^{N}\Phi(K_{j})}.$
(4)

This reduces the time and memory complexity to $\mathcal{O}(N)$, because
$\sum_{j=1}^{N}\Phi(K_{j})V_{j}^{T}$ and $\sum_{j=1}^{N}\Phi(K_{j})$ can be
computed once and reused for each subsequent query computation (7).

Since the linearization of the softmax is infeasible Fast-Transformers use

$\Phi(\mathbf{x})=\text{elu}(\mathbf{x}+1)$ (5)

as feature mapping and show that their formulation performs on par with the
original softmax-attention, while significantly reducing time and memory
consumption.

###  2.3 Transformers in Computer Vision

The success of transformers in NLP has naturally raised the question if
computer vision tasks might as well benefit from transformer-like attention.
Recent work on image classification (9, 4) and pixel-by-pixel image
generation/completion (10, 6, 7) were among the first to successfully
demonstrate the applicability of transformers in the image-domain. In (10),
for example, the first $n$ pixels of the flattened input image are used to
condition a generative transformer setup that then predicts the remaining
image in an auto-regressive manner. Since image-based applications have to
deal with very long input sequences, the use of efficient transformer
implementations (7, 11, 12, 13) is essential (see discussion above).

###  2.4 Tomographic Image Reconstruction

In computed tomography (CT), the radon transform (14, 15) of a 2D sample
section is acquired by rotating a 1D detector array around the sample,
acquiring a series of density measurements at defined projection angles
$\alpha_{i}$. The most common way to visualize all acquired projections
(density measurements) is to collect them as pixel columns in a 2D image
called a sinogram. The most commonly used method to reconstruct the original
2D image from a sinogram is called filtered backprojection (FBP) (15, 16) (see
Figure 3).

The number of projections acquired at different angles directly influences the
quality of the reconstructed image, with more projections leading to improved
reconstruction quality. In practice, it is desirable to limit the number of
projections/views in order to reduce overall acquisition times and total
sample exposure. Such sparse-view acquisitions lead to artefacts in the final
tomographic reconstruction.

\begin{overpic}[trim=1.00374pt 130.48749pt 0.0pt
1.00374pt,clip,width=433.62pt]{figs/fig_fit-sres.pdf}
\put(0.5,2.0){\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\scriptsize{(a)}}
\put(17.2,2.0){\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\scriptsize{(b)}}
\put(34.1,2.0){\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\scriptsize{(c)}}
\put(51.0,2.0){\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\scriptsize{(d)}}
\put(67.6,2.0){\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\scriptsize{(e)}}
\put(84.5,2.0){\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}\scriptsize{(f)}}
\end{overpic} Figure 2: FIT for super-resolution. Low-resolution input images
are first transformed into Fourier space and then unrolled into an FDE
sequence, as described in Section 3.1. This FDE sequence can now be fed to a
FIT, that, conditioned on this input, extends the FDE sequence to represent a
higher resolution image. This setup is trained using an
$\mathcal{FC}\text{-}\mathcal{L}oss$ that enforces consistency between
predicted and ground truth Fourier coefficients. During inference, the FIT is
conditioned on the first $39$ entries of the FDE, corresponding to (a,d)
$3\times$ Fourier binned input images. Panels (b,e) show the inverse Fourier
transform of the predicted output, and panels (c,f) depict the corresponding
ground truth.

Reconstruction artefacts typically manifest as radial stripes radiating out of
image locations subject to large intensity gradients (see _e.g_. Figure 6).
The reason for these artefacts to appear is well understood. The Fourier slice
theorem states that the Fourier coefficients of each 1D projection at a given
angle $\alpha_{i}$ coincide with the 2D Fourier coefficients that lie on the
line that crosses the DC component at angle $\alpha_{i}$ (17). Hence, a sparse
sampling of projection angles leaves many 2D Fourier coefficients unobserved,
leading to said missing wedge artefacts. A plethora of methods to reduce these
artefacts was proposed (18, 19, 20, 21), but none of them is operating in
Fourier space. Below, we will be introducing a novel transformer-based
approach that aims at filling in all unobserved Fourier coefficients such that
missing-wedge artefacts will be reduced or, ideally, avoided.

##  3 Methods

###  3.1 Fourier Domain Encodings (FDEs)

To compute the Fourier Domain Encoding (FDE) for an image $\bm{x}$, we first
take the discrete Fourier transform (DFT), $\bm{X}=\mathcal{F}(\bm{x})$,
giving us the complex valued Fourier spectrum $\bm{X}$. The DC component of
$\bm{X}$ is in its center-most location. Concentric rings of Fourier
coefficients around the DC component are called Fourier rings. More central
Fourier rings contain lower frequencies, coefficients further away from the DC
component higher ones. Since we started with a real-valued image $\bm{x}$, we
know that the coefficients to the left of the DC component are redundant to
the ones on the right. Hence, we drop the left half of $\bm{X}$ and call the
remaining half $\bm{X}_{h}$. If one masks all Fourier ring coefficients up to
a radius $r$ and then back-transforms the result we receive
$\bm{x_{0-r}}=\mathcal{F}^{-1}(\bm{X\odot M_{0-r}})$, with $\bf{M_{0-r}}$
being a circular mask containing $1$ at every location from the DC component
up to the Fourier coefficients at distance $r$ and $0$’s beyond that. The
image $\bm{x_{0-r}}$ is a lower-resolution version of the original image
$\bm{x}$, its effective resolution depending on the value of $r$, see Figure 2
for an example. Hence, if we create a sequence of Fourier coefficients
starting from the DC component and followed by an unrolling of (half) Fourier
rings from $\bm{X}_{h}$, we end up with a sequential image representation

$\mathbf{S}=\text{unroll}(\bm{X}_{h})=\left[c_{1},c_{2},\dots,c_{N}\right]^{T},$
(6)

where $c_{i}$ – the words of our input sequence – are complex Fourier
coefficients. The sequence $\mathbf{S}$ has the desired prefix-properties we
need for our super-resolution task, which will be introduced in Section 3.3.

Figure 3: FIT for computed tomography. We propose an encoder-decoder based
Fourier Image Transformer setup for tomographic reconstruction. In 2D computed
tomography, 1D projections of an imaged sample (_i.e_. the columns of a
sinogram) are back-transformed into a 2D image. A common method for this
transformation is the filtered backprojection (FBP) (15, 16). Since each
projection maps to a line of coefficients in 2D Fourier space, a limited
number of projections in a sinogram leads to visible streaking artefacts due
to missing/unobserved Fourier coefficients. The idea of our FIT setup is to
encode all information of a given sinogram and use the decoder to predict
missing Fourier coefficients. The reconstructed image is then computed via an
inverse Fourier transform (iFFT) of these predictions. In order to reduce high
frequency fluctuations in this result, we introduce a shallow conv-block after
the iFFT (shown in black). We train this setup combining the
$\mathcal{FC}\text{-}\mathcal{L}oss$, see Section 3.2, and a conventional MSE-
loss between prediction and ground truth.

In order to proceed, we convert the complex Fourier coefficients $c_{i}$ into
normalized amplitudes

$a_{i}=\frac{2(|c_{i}|-a_{max})}{a_{max}-a_{min}}-1$ (7)

and phases

$\phi_{i}=\frac{\angle(c_{i})}{\pi},$ (8)

where $a_{min}$ and $a_{max}$ are minimum and maximum amplitudes computed over
all training images and the function $\angle$ returns the phase of a given
Fourier coefficient. Hence, the complex sequence $\mathbf{S}$ can now be
described by the normalized real-valued matrix

$\displaystyle\mathbf{C}=\begin{bmatrix}a_{1}&\cdots&a_{N}\\\
\phi_{1}&\cdots&\phi_{N}\end{bmatrix}^{T}$ (9)

with $\mathbf{C}\in\mathbb{R}^{N\times 2}$.

Our final goal is to transform each word $(a_{i},\phi_{i})$ into an
$F$-dimensional vector. To this end we feed $\mathbf{C}$ through a single
trainable linear layer that increases the feature dimensionality from $2$ to
$\frac{F}{2}$, to which we concatenate an $\frac{F}{2}$-dimensional 2D
positional encoding222 We slightly adapted the 2D positional encoding of (8)
to accept arbitrary (non integer) coordinates. that represents the original
polar coordinates of the Fourier coefficient in the original 2D Fourier
spectrum $\bm{X_{h}}$. The final FDE image sequence is therefore
$\bm{E}\in\mathbb{R}^{N\times F}$.

Predicted output words $\mathbf{Z}=[z_{1},\dots,z_{k}]$, with
$z_{i}\in\mathbb{R}^{F}$, are fed through two linear layers that back-
transform the $F$-dimensional encoding of $z_{i}$ into predicted amplitudes
and phases $\hat{c}_{i}=(\hat{a}_{i},\hat{\phi}_{i})$, respectively. The
output of the phase-predicting layer is additionally passed through the
$tanh$-activation function to ensure that all phases are in $[-1,1]$.

\begin{overpic}[width=433.62pt]{figs/sres_quantitative.pdf}
\put(4.0,22.4){\hbox{\pagecolor{black}\makebox(1.0,1.0)[]{\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}(a)}}}
\put(20.0,22.4){\hbox{\pagecolor{black}\makebox(1.0,1.0)[]{\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}(b)}}}
\put(36.0,22.4){\hbox{\pagecolor{black}\makebox(1.0,1.0)[]{\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}(c)}}}
\put(51.9,22.4){\hbox{\pagecolor{black}\makebox(1.0,1.0)[]{\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}(d)}}}
\put(67.85,22.4){\hbox{\pagecolor{black}\makebox(1.0,1.0)[]{\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}(e)}}}
\put(83.8,22.4){\hbox{\pagecolor{black}\makebox(1.0,1.0)[]{\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}(f)}}}
\end{overpic} Figure 4: Super-resolution results on MNIST. The top left
triangles of (a,c,e) show $3\times$ binned (low-res) MNIST inputs. Conditioned
by these inputs, our trained FIT auto-regressively generates results as shown
in (b,d,f). Inputs and predictions are labeled with the peak signal-to-noise
ratio (PSNR) values computed w.r.t. ground truth images, shown in the lower-
right half of (a,c,e). The examples in (a-f) correspond to the 98th, 50th and
2nd percentile in terms of obtained PSNR over all MNIST prediction results.
Box-plots show the distribution of PSNR values of Fourier binned inputs and
predicted outputs, respectively (mean in dashed gold and median in solid
blue). The Fourier ring correlation plot shows how predicted Fourier
coefficients are improving w.r.t. true ground truth coefficients. Shaded areas
correspond to $+/-1$ standard deviation. The correlation for the first $5$
Fourier shells is $1$ because these shells have been used as inputs to the
FIT.  \begin{overpic}[width=433.62pt]{figs/sres_celeb_conditioning.pdf}
\end{overpic} Figure 5: Super-resolution results on CelebA. On $3$ input
images (rows $1$+$2$, $3$+$4$, and $5$+$6$) we show predictions of a trained
super-resolution FIT conditioned on $2$, $4$, $8$ and $16$ Fourier rings
(columns). Upper rows show the iFFT of the input FDEs used to condition the
FIT, lower rows depict the iFFT of the predicted results.

###  3.2 Fourier Coefficient Loss

We train our Fourier Image Transformers (FITs) with a loss function consisting
of two terms, $(i)$ the amplitude loss

$\mathcal{L}_{amp}(\hat{a}_{i},a_{i})=1+(\hat{a}_{i}-a_{i})^{2},$ (10)

computed between the predicted amplitudes $\hat{a}_{i}$ and the target
amplitudes $a_{i}$, and $(ii)$ the phase loss

$\mathcal{L}_{\angle}(\hat{\phi}_{i},\phi_{i})=2-cos(\hat{\phi}_{i}-\phi_{i}),$
(11)

with $\hat{\phi}_{i}$ the predicted phase and $\phi_{i}$ the corresponding
target phase.

The final Fourier coefficient loss $\mathcal{L_{FC}}$ is the multiplicative
combination of both individual losses, given by

$\mathcal{L_{FC}}(\hat{\mathbf{C}},\mathbf{C})=\frac{1}{N}\sum_{i=0}^{N}\mathcal{L}_{amp}(\hat{a}_{i},a_{i})\cdot\mathcal{L}_{\angle}(\hat{\phi}_{i},\phi_{i}).$
(12)

###  3.3 FIT for Super-Resolution

The Fourier domain encoding (FDE) we describe above is a sequential
representation of an input image $\bm{x}$, for which each prefix $pre$ encodes
a reduced resolution image $\bm{x}_{pre}$.

Hence, using FDE image sequences, we can train an auto-regressive Fourier
Image Transformer that takes an encoded sequence
$\bm{E}=[e_{1},\dots,e_{N-1}]$ as input and predicts the FDE sequence
$\bm{Z}=[z_{2},\dots,z_{N}]$. This FIT for super-resolution (“FIT: SRes”) is
trained w.r.t. the correct target sequence $\bm{C}=[c_{2},\dots,c_{N}]$, using
the previously introduced Fourier coefficient loss $\mathcal{L_{FC}}$, which
is computed using the back-transformed predicted amplitude and phase values
$\bm{\hat{C}}=[\hat{c}_{2},\dots,\hat{c}_{N}]$, where
$\hat{c}_{i}=(\hat{a}_{i},\hat{\phi}_{i})$, as explained in Section 3.1.

Once the transformer is trained, a prefix of a complete FDE sequence
$\bm{E}_{pre}=[e_{1},\dots,e_{|pre|}]$ is used to condition the transformer,
which is then used in iterations to auto-regressively predict the missing part
of the complete sequence $\bm{E}=[e_{1},\dots,e_{N}]$, _i.e_. filling in
predicted high-frequency information not contained in $\bm{E}_{pre}$ (see
Figure 2).

Note that the proposed super-resolution setup operates exclusively on Fourier
domain encoded data. All final prediction images $\bm{\hat{x}}$ are generated
by computing the inverse Fourier transform on predictions $\bm{\hat{C}}$,
which are rearranged (rolled) into $\bm{\hat{X}_{h}}$ and completed to a full
predicted Fourier spectrum $\bm{\hat{X}}$, _i.e_.
$\bm{\hat{x}}=\mathcal{F}^{-1}(\text{roll}(\bm{\hat{C}}))$.

###  3.4 FIT for Tomography

Our Fourier Image Transformer setup for tomographic reconstruction (“FIT:
TRec”) is based on an encoder-decoder transformer architecture as shown in
Figure 3.

As input to the encoder we use the Fourier Domain Encoding (FDE) of a raw
sinogram $\bm{s}$. As described above, $\bm{s}$ consists of $P$ pixel columns
$[s_{1},\dots,s_{P}]$ of 1D projections of $\bm{x}$ at angles
$[\alpha_{1},\dots,\alpha_{P}]$. The Fourier slice theorem states, see also
Section 2.4, that the discrete 1D Fourier coefficients
$\bm{C}_{i}=\mathcal{F}(s_{i})$ coincide with the values of the 1D slice at
angle $\alpha_{i}$ through the 2D Fourier spectrum $\mathcal{F}(\bm{x})$. To
assemble the full FDE of a sinogram we need to combine all $\bm{C}_{i}$ with
the adequate positional encoding (using polar coordinates) of all Fourier
coefficients, as dictated by the Fourier slice theorem and sketched in Figure
3.

Hence, the encoder creates a latent space representation $\bm{Z}$ that encodes
the full input sinogram $\bm{s}$. This latent space encoding is then given as
input to the decoder. The decoder is then used to predict all Fourier
coefficients $\bm{\hat{C}}$, such that the predicted reconstruction
$\bm{\hat{x}}$ of $\bm{x}$ can be computed by
$\bm{\hat{x}}=\mathcal{F}^{-1}(\text{roll}(\bm{\hat{C}}))$, where roll
arranges the 1D sequence back into a discrete 2D Fourier spectrum. We call
this setup “FIT: TRec”.

Additionally, we propose a variation of this procedure, called “FIT: TRec +
FBP”, where the decoder not only receives the latent space encoding $\bm{Z}$,
but also FDEs of the Fourier coefficients
$\bm{C}_{\text{FBP}}=\mathcal{F}(\text{FBP}(\bm{s}))$, where FBP denotes the
function computing the filtered backprojection of a sinogram (see Figure 3).

Note that the implementation of “FIT: TRec” coincides with “FIT: TRec + FBP”,
with FBP being replaced by a function ZERO which returns $0$ for all inputs.

We train “FIT: TRec” and “FIT: TRec + FBP” using the $\mathcal{L_{FC}}$-loss
of Eq.12. Additionally, we introduced a residual convolution block consisting
of two convolutional layers ($3\times 3$ followed by $1\times 1$) with
$d_{\text{conv}}=8$ intermediate feature channels. This conv-block (conv)
receives the inverse Fourier transform of the predicted Fourier coefficients
$\bm{\hat{x}}=\mathcal{F}^{-1}(\text{roll}(\bm{\hat{C}}))$ as input and is
trained using the MSE-loss between the predicted real-space image
$\text{conv}(\bm{\hat{x}})$ and the known ground truth image $\bm{x}$. Hence,
the full loss is the sum over $\mathcal{L_{FC}}$ and the MSE-loss.

In order to speed up training, we start by feeding only a low-resolution
subset of all Fourier coefficients $\bm{C}_{i}=\mathcal{F}(s_{i})$ (and
$\bm{C}_{\text{FBP}}$), and successively increase this subset over training
until the full sets are used. This forces the FIT to first learn good low
resolution features and later learn to add suitable high resolution
predictions.

##  4 Experiments and Results

Here we describe the experiments we conduct with the previously described
super-resolution (“FIT: SRes”) and tomographic reconstruction (“FIT: TRec” and
“FIT: TRec + FBP”) setups. Note, all code used to create results reported in
this manuscript is available on
GitHub333https://github.com/juglab/FourierImageTransformer.

###  4.1 Data and Metrics for Super-Resolution

MNIST (22): Cropped to $27\times 27$ pixels with the default train-test split
(in $60^{\prime}000$ and $10^{\prime}000$ images, respectively). The train
images are further split into $55^{\prime}000$ samples for training and
$5^{\prime}000$ validation images.

CelebA $\mathbf{128\times 128}$ (23): Converted to gray scale and downscaled
to $63\times 63$ pixels. The images are randomly split into $20^{\prime}000$,
$5^{\prime}000$ and $5^{\prime}000$ training, validation and test samples,
respectively.

We evaluate all results using $(i)$ Fourier Ring Correlation (FRC) (24) in
Fourier space and $(ii)$ the peak signal-to-noise ratio (PSNR) in image space.

###  4.2 Data and Metrics for Computed Tomography

Data Preparation for Tomographic Reconstruction: As described in Section 2.4,
tomographic image reconstruction in 2D operates on a number of 1D projections
of a given true object $\bm{x}$. Our tomographic simulation process is based
on the work by Leuschner  _et al_. (25). Furthermore, we chose the detector
length to be equal to the width of our chosen object (_i.e_. ground truth
image) we apply our synthetic tomography pipeline on. To avoid spurious
contributions to individual projections, we need to set all pixel intensities
outside the largest image-centered circle to $0$ (hence, we see only circular
images in Figures 1, 3, and 6).

MNIST (22): Data is split and preprocessed as described in Section 4.1.
Additionally, for visualization purposes, we min-clip all pixel intensities
within the before-mentioned largest image-centered circle to $50$. Finally, we
use this data to compute $P=7$ equally spaced projections which we assemble in
sinograms
$\bm{s}^{j}_{\text{MNIST}}=[s^{j,1}_{\text{MNIST}},\dots,s^{j,P}_{\text{MNIST}}]$.

Kanji (26): Data is randomly split into $50^{\prime}000$ train,
$5^{\prime}000$ validation and $5^{\prime}000$ test samples and all images are
cropped to $63\times 63$ pixels, which are otherwise processed as described
for MNIST. Finally, we use this data to compute $P=33$ equally spaced
projections which we assemble in sinograms
$\bm{s}^{j}_{\text{Kanji}}=[s^{j,1}_{\text{Kanji}},\dots,s^{j,P}_{\text{Kanji}}]$.

LoDoPaB (25): The original train- and validation-data is first reduced to
$4^{\prime}000$ and $400$ randomly chosen images respectively and for testing
all $3^{\prime}553$ images are used. All selected images are downscaled to
$111\times 111$ pixels and we compute $P=33$ equally spaced projections like
for the Kanji data.

All tomographic reconstruction experiments with “FIT: TRec” and “FIT: TRec +
FBP” and the FBP baseline are evaluated using peak signal-to-noise ratio
(PSNR) w.r.t. available ground truth.

###  4.3 Training Setup for Super-Resolution

We use a $F=256$ dimensional FDE, with the positional encoding being based on
polar coordinates, _i.e_. Fourier coefficients of same frequency have the same
radius. The FDE is passed to a causal-linear transformer (7) with $8$ layers,
$8$ self-attention heads, a query and value dimensionality of $32$, dropout of
$0.1$, attention dropout of $0.1$, and a dimensionality of the feed-forward
network of $1024$.

This setup is trained auto-regressively, _i.e_. with a triangular attention
mask. We use the rectified Adam optimizer (RAdam) (27) with an initial
learning rate of $0.0001$ and weight decay of $0.01$ for $100$ epochs. The
batch size is $32$. We half the learning rate on plateauing validation loss.

###  4.4 Results for Super-Resolution Tasks

Quantitative results for all conducted super-resolution experiments on MNIST
data are shown in Figure 4, where we show $(i)$ $3\times$ binned low-res MNIST
input images corresponding to $\bm{E}_{pre}=[c_{1},\dots,c_{39}]$, also
sketched in Figure 2, $(ii)$ corresponding ground truth MNIST images, $(iii)$
the predictions of the “FIT: SRes” network trained on the MNIST data as
described in Section 3.3, $(iv)$ two box-plots showing the distribution of
PSNR values computed between the ground truth images and the downscaled inputs
and predicted outputs, respectively, and $(v)$ Fourier ring correlation plots
showing the correlation between our predicted Fourier coefficients and the
corresponding ground truth.

In Figure 5, we show $3$ sequences of super-resolution results obtained with a
FIT trained on the CelebA data. For each image, we conditioned the trained
transformer on $2$, $4$, $8$, and $16$ Fourier rings, respectively. This
corresponds to low-resolution images subject to $16\times$, $8\times$,
$4\times$, and $2\times$ binning in Fourier space, respectively.

###  4.5 Training Setup for Computed Tomography

Dataset | Method | PSNR
---|---|---
MNIST (22) | Baseline: FBP | $17.87$
FIT: TRec + FBP (Ours) | $27.85$
FIT: TRec (Ours) | $\mathbf{27.90}$
Ablation: Only FBP | $26.89$
Ablation: Only Conv-Block | $22.53$
Kanji (26) | Baseline: FBP | $22.06$
FIT: TRec + FBP (Ours) | $\mathbf{30.72}$
FIT: TRec (Ours) | $25.99$
Ablation: Only FBP | $30.49$
Ablation: Only Conv-Block | $26.92$
LoDoPaB (25) (downscaled) | Baseline: FBP | $26.89$
FIT: TRec + FBP (Ours) | $\mathbf{30.98}$
FIT: TRec (Ours) | $21.90$
Ablation: Only FBP | $30.74$
Ablation: Only Conv-Block | $30.70$
Table 1: Quantitative tomographic reconstruction results. We report the
average peak signal-to-noise ratio (PSNR) with respect to ground truth, for
each of the three used datasets. For each dataset, we compare the results of
our “FIT: TRec + FBP” and “FIT: TRec” setups to results obtained with the
filtered backprojection (FBP) (15, 16) baseline, and the two ablation studies
described in Section 6.

Like before, we consistently use $F=256$ dimensional FDEs, and employ the
linear encoder and decoder method by Katharopoulos  _et al_. (7). More
specifically, we use $4$ transformer layers, $8$ self-attention heads per
layer, a query and value dimensionality of $32$, dropout of $0.1$, attention
dropout of $0.1$, and a dimensionality of the feed-forward network of $1024$.
The residual conv-block has $d_{\text{conv}}=8$ intermediate feature channels.

All networks are optimized using RAdam (27), with an initial learning rate of
$0.0001$ and weight decay of $0.01$ for $300$ (MNIST), $120$ (Kanji), and
$350$ (LoDoPaB) epochs. The batch size is $32$. We half the learning rate on
plateauing validation loss.

\begin{overpic}[width=420.61192pt,trim=0.0pt 1087.06125pt 0.0pt
0.0pt,clip,right]{figs/trec_mnist_ta7_trec_trecfbp_quali_v0.1.24_zero.pdf}
\put(0.0,8.0){\rotatebox{90.0}{\scriptsize{Input Sinograms ($7$ Projections,
MNIST)}}} \end{overpic}\begin{overpic}[width=420.61192pt,trim=0.0pt
1087.06125pt 0.0pt
21.58061pt,clip,right]{figs/trec_kanji_ta33_trec_trecfbp_quali_v0.1.24_zero.pdf}
\put(0.0,8.0){\rotatebox{90.0}{\scriptsize{Input Sinograms ($33$ Projections,
Kanji)}}} \end{overpic}\begin{overpic}[width=420.61192pt,trim=0.0pt
1087.06125pt 0.0pt
21.58061pt,clip,right]{figs/trec_lodopab_ta33_trec_trecfbp_quali_v0.1.24_zero.pdf}
\put(0.0,6.0){\rotatebox{90.0}{\scriptsize{Input Sinograms ($33$ Projections,
LoDoPaB)}}} \end{overpic}

Figure 6: Tomographic reconstruction results. We show three qualitative
results for each used datasets. From left to right, we show the input
sinogram, reconstruction results obtained with filtered backprojection (FBP)
(15, 16), our results obtained with the “FIT: TRec” setup, our results
obtained with the “FIT: TRec + FBP” setup, and the corresponding ground truth
images. In the top left corner of each reconstruction we show the peak signal-
to-noise ratio (PSNR) with respect to the ground truth image.

*: Ablation Studies We propose two ablation setups for all tomographic reconstruction experiments.

First, we ask what influence the encoded latent space information $\bm{Z}$,
_i.e_. the output of the encoded sinogram, has on the quality of the overall
reconstruction $\bf{\hat{x}}$. To that end, we perform ablation experiments
for all $3$ datasets, for which we do not feed $\bm{Z}$ to the decoder.
Technically this is implemented by replacing the decoder by an encoder network
(since only one input remains to be fed). We label these experiments “Only
FBP”.

The second ablation study asks, to what degree the conv-block contributes to
the overall reconstruction performance, _i.e_. we want to verify that the
convolution block alone is not sufficient to solve the task at hand. Hence, we
train the conv-block on pairs of images $(\text{FBP}(\bm{s}),\bm{x})$, _i.e_.
the filtered backprojection of sinograms $\bm{s}$ and their corresponding
ground truth images $\bm{x}$. We label these experiments “Only Conv-Block”.

For all ablation experiments all hyper-parameters not explicitly mentioned
above are kept unchanged.

###  4.6 Results for Tomographic Reconstruction Tasks

Qualitative tomographic reconstruction results for all three datasets we used
are shown in Figure 6. For each dataset, we show three input sinograms, the
reconstruction baseline obtained via filtered backprojection (FBP), our
results obtained via “FIT: TRec”and “FIT: TRec + FBP”, and the corresponding
ground truth images.

In Table 1, PSNR numbers for all three datasets using the FBP baseline, our
“FIT: TRec + FBP” and “FIT: TRec” training setups, and both ablation studies
are given.

##  5 Discussion

We proposed the idea of Fourier Domain Encodings (FDEs), a novel sequential
image encoding, for which each prefix represents the whole image at reduced
resolution, and demonstrated the utility of FDEs for solving two common image
processing tasks with transformer networks, _i.e_. super-resolution and
tomographic image reconstruction.

For the super-resolution task we showed that Fourier Image Transformer can be
trained such that, when conditioned on an FDE corresponding to a low-
resolution input image, can auto-regressively predict an extended FDE sequence
that can be back-transformed into a higher resolution output. It is obvious,
the information required to generate a higher resolution image must be stored
in the trained network, and we have shown in Figure 5, how this learned prior
completes very low to moderate resolution inputs in sensible ways. It is
curious to see that eyes are the first high-resolution structures filled in by
the trained FIT. We believe that this is a direct consequence of all training
images being registered such that the eyes are consistently at the same
location.

For the tomographic reconstruction task, we employ an encoder-decoder
transformer that encodes a given FDE sequence corresponding to a given
sinogram and can then be used to predict Fourier coefficients at arbitrary
query locations. We used the decoder to predict all Fourier coefficients of a
fully reconstructed image, which we could then visualize via inverse Fourier
transformation (iFFT). We noticed that introducing a shallow residual
convolution block after the iFFT reduces unwanted high frequency fluctuations
in predicted results. While we see that this procedure leads to very
convincing results on MNIST, for more complex datasets, results quickly
deteriorate. Hence, we proposed to additionally feed Fourier coefficients
obtained by filtered backprojection (FBP) into the decoder. This leads to much
improved results that outperform the FBP baseline, showing that the FIT does
contribute to solving the reconstruction task.

Our results show that transformers, currently the dominant approach for
virtually all NLP tasks, can successfully be applied to complex and relevant
tasks in computer vision. While this is encouraging, we see a plethora of
possibilities for future improvements. For example, the transformers we used
are rather small. We believe that an up-scaled version of our training setups
with more attention heads and more layers would already lead to much improved
results444https://tinyurl.com/needadviceforafriend. Still, we also believe
that there is plenty of room for methodological improvements that do not
require more computational resources, making this line of research also
accessible to many other research labs around the globe.

##  6 Acknowledgments

The authors would like to thank Mangal Prakash for discussions and help to
improve the text. We also want to thank the Scientific Computing Facility at
MPI-CBG for giving us access to HPC infrastructure. Funding was provided from
the Max-Planck Society under project code M.IF.A.MOZG8106, the core budget of
the Max-Planck Institute of Molecular Cell Biology and Genetics (MPI-CBG), the
Human Technopole, and the BMBF under codes 031L0102 (de.NBI) and 01IS18026C
(ScaDS2), as well as by the DFG under code JU3110/1-1 (FiSS) Additionally, we
thank Yannic Kilcher and his YouTube
channel555https://www.youtube.com/c/YannicKilcher which greatly facilitated
our desire to think about applications of transformers on images.

##  7 Bibliography

## References

  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Radford et al. (2018) Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever.  Improving language understanding by generative pre-training, 2018.
  * Lee et al. (2019) Jinhyuk Lee, Wonjin Yoon, Sungdong Kim, Donghyeon Kim, Sunkyu Kim, Chan Ho So, and Jaewoo Kang.  BioBERT: a pre-trained biomedical language representation model for biomedical text mining.  _Bioinformatics_ , 09 2019.  ISSN 1367-4803.  10.1093/bioinformatics/btz682.
  * Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In _Advances in neural information processing systems_ , pages 5998–6008, 2017.
  * Chen et al. (2020) Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever.  Generative pretraining from pixels.  In _International Conference on Machine Learning_ , pages 1691–1703. PMLR, 2020.
  * Katharopoulos et al. (2020) A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret.  Transformers are rnns: Fast autoregressive transformers with linear attention.  In _Proceedings of the International Conference on Machine Learning (ICML)_ , 2020.
  * Wang and Liu (2020) Zelun Wang and Jyh-Charn Liu.  Translating math formula images to latex sequences using deep neural networks with sequence-level training.  _International Journal on Document Analysis and Recognition (IJDAR)_ , pages 1–13, 2020.
  * Ramachandran et al. (2019) Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jon Shlens.  Stand-alone self-attention in vision models.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.
  * Parmar et al. (2018) Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran.  Image transformer.  In _International Conference on Machine Learning_ , pages 4055–4064. PMLR, 2018.
  * Kitaev et al. (2020) Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya.  Reformer: The efficient transformer.  _arXiv preprint arXiv:2001.04451_ , 2020.
  * Wang et al. (2020) Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma.  Linformer: Self-attention with linear complexity.  _arXiv preprint arXiv:2006.04768_ , 2020.
  * Bello (2021) Irwan Bello.  Lambdanetworks: Modeling long-range interactions without attention.  In _International Conference on Learning Representations_ , 2021.
  * Radon (1917) J. Radon.  Über die Bestimmung von Funktionen durch ihre Integralwerte längs gewisser Mannigfaltigkeiten.  _Akad. Wiss._ , 69:262–277, 1917.
  * Kak et al. (2002) Avinash C Kak, Malcolm Slaney, and Ge Wang.  Principles of computerized tomographic imaging, 2002.
  * Ramesh et al. (1989) GR Ramesh, N Srinivasa, and K Rajgopal.  An algorithm for computing the discrete radon transform with some applications.  In _Fourth IEEE Region 10 International Conference TENCON_ , pages 78–81. IEEE, 1989.
  * Bracewell (1956) Ronald N Bracewell.  Strip integration in radio astronomy.  _Australian Journal of Physics_ , 9(2):198–217, 1956.
  * Chen et al. (2017) Hu Chen, Yi Zhang, Mannudeep K Kalra, Feng Lin, Yang Chen, Peixi Liao, Jiliu Zhou, and Ge Wang.  Low-dose ct with a residual encoder-decoder convolutional neural network.  _IEEE transactions on medical imaging_ , 36(12):2524–2535, 2017.
  * Jin et al. (2017) Kyong Hwan Jin, Michael T McCann, Emmanuel Froustey, and Michael Unser.  Deep convolutional neural network for inverse problems in imaging.  _IEEE Transactions on Image Processing_ , 26(9):4509–4522, 2017.
  * Adler and Öktem (2018) Jonas Adler and Ozan Öktem.  Learned primal-dual reconstruction.  _IEEE transactions on medical imaging_ , 37(6):1322–1332, 2018.
  * Hauptmann et al. (2018) Andreas Hauptmann, Felix Lucka, Marta Betcke, Nam Huynh, Jonas Adler, Ben Cox, Paul Beard, Sebastien Ourselin, and Simon Arridge.  Model-based learning for accelerated, limited-view 3-d photoacoustic tomography.  _IEEE transactions on medical imaging_ , 37(6):1382–1393, 2018.
  * LeCun and Cortes (2010) Yann LeCun and Corinna Cortes.  MNIST handwritten digit database.  2010\.
  * Liu et al. (2015) Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang.  Deep learning face attributes in the wild.  In _Proceedings of International Conference on Computer Vision (ICCV)_ , December 2015.
  * Van Heel et al. (1982) M Van Heel, W Keegstra, W Schutter, and EJF Van Bruggen.  Arthropod hemocyanin structures studied by image analysis.  _Life Chem. Rep. Suppl_ , 1:69–73, 1982.
  * Leuschner et al. (2019) Johannes Leuschner, Maximilian Schmidt, Daniel Otero Baguer, and Peter Maaß.  The lodopab-ct dataset: A benchmark dataset for low-dose ct reconstruction methods.  _arXiv preprint arXiv:1910.01113_ , 2019.
  * Clanuwat et al. (2018) Tarin Clanuwat, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto, and David Ha.  Deep learning for classical japanese literature, 2018.
  * Liu et al. (2019) Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han.  On the variance of the adaptive learning rate and beyond.  _arXiv preprint arXiv:1908.03265_ , 2019.