# Convolve and Conquer: Data Comparison with Wiener Filters

Deborah Pelacani Cruz
Department of Earth Science and Engineering
Imperial College London
London, SW7 2BX
<EMAIL_ADDRESS>
&George Strong††footnotemark:
Department of Earth Science and Engineering
Imperial College London
London, SW7 2BX
<EMAIL_ADDRESS>
&Oscar Bates
Department of Bioengineering
Imperial College London
London, SW7 2BX
<EMAIL_ADDRESS>
&Carlos Cueto
Department of Bioengineering
Imperial College London
London, SW7 2BX
<EMAIL_ADDRESS>
&Jiashun Yao
Department of Earth Science and Engineering
Imperial College London
London, SW7 2BX
<EMAIL_ADDRESS>
&Lluis Guasch
Department of Earth Science and Engineering
Imperial College London
London, SW7 2BX
<EMAIL_ADDRESS>
These authors are equal contributors to this work

###### Abstract

Quantitative evaluations of differences and/or similarities between data
samples define and shape optimisation problems associated with learning data
distributions. Current methods to compare data often suffer from limitations
in capturing such distributions or lack desirable mathematical properties for
optimisation (e.g. smoothness, differentiability, or convexity). In this
paper, we introduce a new method to measure (dis)similarities between paired
samples inspired by Wiener-filter theory. The convolutional nature of Wiener
filters allows us to comprehensively compare data samples in a globally
correlated way. We validate our approach in four machine learning
applications: data compression, medical imaging imputation, translated
classification, and non-parametric generative modelling. Our results
demonstrate increased resolution in reconstructed images with better
perceptual quality and higher data fidelity, as well as robustness against
translations, compared to conventional mean-squared-error analogue
implementations.

##  1 Introduction

Data comparison is essential to optimisation problems that aim to learn data
distributions. Mean squared error (MSE) is the most commonly used measure
because: i) it is simple to implement and interpret, ii) it is computationally
cheap, and iii) it has convenient mathematical properties such as smoothness,
convexity, and differentiability [47]. Despite these advantages, MSE has some
well known limitations extensively discussed in the optimisation community
[46, 50, 9]. MSE-based reconstructions exhibit limited quantitative accuracy
and perceptual fidelity [35, 47], which arise from the assumption that all
data points are independent and of equal importance. This assumption results
in insensitivity to global spatial structure and texture, fixating only on the
properties of the individual elements independently; for example, brightness
and contrast of a single pixel in an image neglects information from
neighbouring pixels.

In the machine learning community, these issues have been partially addressed
by evolving towards complex network architectures, using techniques such as
convolutional layers, skip connections, or deeper models [13, 21, 39].
Networks are known to be biased towards low-frequency content (e.g spectral
bias and F-principle) [38, 51]. This is compounded by the fact that MSE
assumes that samples are normally distributed, which can result in high-
frequency degradation and limited structural fidelity [47]. A potential
solution is to modify the criteria used to quantify data differences, as will
be discussed in Section 2. We believe this research direction has been
comparatively overlooked, despite the fundamental role it plays across machine
learning problems.

We propose a new data comparison strategy using Wiener filters. Specifically,
we use information contained within Wiener filters that globally match sample
pairs through convolution to provide a quantitative measure of similarity. As
will be seen in Section 3, our notion of similarity is motivated by the
comparison between matching Wiener filters and the convolutional-identity
filter. In other words, the closer the Wiener filter coefficients are to a
Dirac-delta function, the more similar the sample pairs are. We showcase the
potential of our Wiener formulation as a learning metric in a machine learning
context with four experiments: compression of natural images, imputation of
medical data, classification with translated data, and a new class of non-
parametric generative models. In these experiments we show that networks can
better capture data distributions, boost perceptual quality using our Wiener-
based measurements compared to MSE, and that these measurements also permit
the construction of a generative model that can sample complex latent spaces.
Even though we apply this new formulation to tasks that are machine-learning
focused, we convey the message that this methodology can be used for any task
where data comparison is required and preservation of data correlation is
desired.

The structure of the paper is as follows. Section 2 reviews related metrics
and loss functions for optimisation problems. Section 3 starts with a
description of Wiener filters, followed by the formulation of a Wiener-based
loss function and a diffusion process. Section 4 showcases our four
experiments, presenting results that highlight the advantages of our method.
Then we discuss special hyper-parameters and computational costs. Finally we
present our conclusions and contributions.

##  2 Related Work

Various approaches have been proposed to mitigate the shortcomings of
conventional data comparison metrics such as MSE. The structural similarity
index measure (SSIM), first introduced by [48] and later proposed as a loss
function by [53], combines differences between the luminance, contrast, and
structure of two samples to provide a perceptual measure of similarity.
However, it has been known to produce unstable results, and further research
is required to understand its deficiencies [34]. More recently, [16] proposed
a promising approach termed the focal frequency loss, a Fourier domain
distance designed to address discrepancies between the Fourier domain
representation of real and generated images.

The use of deep neural networks for data comparison has become well
established due to their ability to learn hierarchical representations that
are sensitive to correlations within data. For instance, the perceptual loss,
proposed by [17], compares data representations within the latent space of
pre-trained convolutional neural networks. A drawback of this approach is that
the pre-trained networks used will incorporate potentially undesirable biases
from their architecture and the data that was originally used to pre-train
them. The discriminator from generative adversarial networks (GANs) [10] is
another neural-network-based loss function. However, adversarial training is
well known to suffer from instability problems.

Another increasingly popular data comparison metric is the Wasserstein
distance, which derives from optimal transport theory [37]. It has been
directly leveraged as a loss function to improve multi-label classification
[8], and the sliced Wasserstein distance, a common approximation, has been
used to formulate a non-parametric generative model [27] that bears conceptual
similarities to the Wiener diffusion we present in section 3.3. A limitation
of Wasserstein distances is that they are defined only for non-negative and
normalised signals. [42] provides a solution to this constraint, but the
proposed method requires the challenging determination of a hyper-parameter
controlling the tradeoff between the importance of amplitude and location
within a signal. Wasserstein distances are computationally expensive to
calculate, requiring the solution to an additional optimisation problem,
although they can be approximated using a faster entropy regularised
formulation [37].

The direct inspiration for this work comes from adaptive waveform inversion
(AWI) [49], a loss function that uses convolutional data-matching Wiener
filters for seismic imaging. Advanced loss functions, including optimal-
transport-based formulations [30], are a highly active research focus in
seismic imaging as they have a pivotal effect on the obtained outcome due to
the extreme non-linearity of the problem. AWI is mathematically and
conceptually the most similar technique to our proposed methods, as shall be
discussed in section 3.

##  3 Convolutional Data Comparison

###  3.1 Preliminary

The Wiener filter is central to all techniques presented in this work. It is a
well known tool that is frequently used within image processing, where it
typically arises as a deconvolutional filter that, when convolved with some
observed signal, produces the best least-squares estimate of an unknown target
signal of interest. Depending on the specific details of the application, this
unknown target signal may be a denoised or deblurred version of the observed
signal, hence the common use of Wiener filters for image denoising and
deblurring applications [15]. In this work, rather than using Wiener filters
for deconvolution, we exploit their properties to provide a global match
between signals.

A Wiener filter can be defined as the convolutional filter that provides the
best least-squares match between two signals. Mathematically, if we wish to
find the Wiener filter that matches two signals, $\mathbf{x}$ and
$\mathbf{y}$, expressed as column vectors, then we would have to minimise the
following expression with respect to $\mathbf{v}$:

$\underset{\mathbf{v}}{\textup{min}}\
\frac{1}{2}||\mathbf{Y}\mathbf{v}-\mathbf{x}||^{2}$ (1)

where $\mathbf{v}$ is the Wiener filter, and $\mathbf{Y}$ is a matrix that
represents convolution with signal $\mathbf{y}$. If signals $\mathbf{x}$ and
$\mathbf{y}$ are one-dimensional, then $\mathbf{Y}$ will take the form of a
Toeplitz matrix, with progressively lagged values of the signal $\mathbf{y}$
placed in each column [12]. If signals $\mathbf{x}$ and $\mathbf{y}$ are
vectors representing two-dimensional images, then $\mathbf{Y}$ is given by a
doubly block Toeplitz matrix [12]. The solution to this linear least-squares
problem is given by the well known equation:

$\mathbf{v}(\mathbf{x},\mathbf{y})=(\mathbf{Y}^{T}\mathbf{Y})^{-1}\mathbf{Y}^{T}\mathbf{x}$
(2)

where $\mathbf{Y}^{T}$ is the transpose of the Toeplitz matrix $\mathbf{Y}$,
that represents the cross-correlation with signal $\mathbf{y}$ [12]. In full,
this equation can be interpreted as the cross-correlation of $\mathbf{y}$ with
$\mathbf{x}$, deconvolved by the autocorrelation of $\mathbf{y}$.

In this work, we use full-lag Wiener filters that are of at least the same
size as signals $\mathbf{x}$ and $\mathbf{y}$. The receptive field of any
element within signal $\mathbf{x}$ is therefore equivalent to the entirety of
signal $\mathbf{y}$ and vice versa. This means that Wiener filter
$\mathbf{v}(\mathbf{x},\mathbf{y})$ globally matches signal $\mathbf{y}$ to
signal $\mathbf{x}$ through convolution, without assuming any local element-
wise relationship between any of the individual elements or structures
contained within signals $\mathbf{x}$ and $\mathbf{y}$. This sensitivity to
the global content within the samples is a core advantage of using Wiener
filters for data comparison over conventional methods, such as MSE, which is
sensitive only to local element-wise differences. It is worth noting that
truncated Wiener filters could of course be used in applications where it may
be advantageous to limit the receptive field in the data matching process,
though we leave further exploration of this topic for future research.

If signal $\mathbf{x}$ and $\mathbf{y}$ are identical, then the resulting
Wiener filter that best matches them is given by the Dirac delta function -
the identity function of the convolution operation. The Wiener filter that
would arise if $\mathbf{x}$ is a translated copy of $\mathbf{y}$ is simply a
Dirac delta function translated in the same manner. This is visually
demonstrated in Figure 1 by calculating the resulting Wiener filters that
match a sample from the MNIST data set to i) random noise, ii) an identical
copy of the sample, and iii) a translated copy of the sample. We leverage
these properties for a number of specific techniques and applications
throughout this work. However, we stress that our aim is to promote and
showcase the use of Wiener filters for data comparison in a general sense. As
we will demonstrate, Wiener-filter-based data comparison is highly versatile,
and the exact formulation that is adopted can be selected and customised to
the task and data at hand.

Figure 1: Top: visual illustration of two-dimensional Wiener filters that
match MNIST sample $\mathbf{y}$ to random noise $\mathbf{x}$. Middle: matching
$\mathbf{y}$ to $\mathbf{x}$ where $\mathbf{x}=\mathbf{y}$, demonstrating that
the Wiener filter when $\mathbf{x}=\mathbf{y}$ is given by the Dirac delta
function. Bottom: matching $\mathbf{y}$ to $\mathbf{x}$ where $\mathbf{x}$ is
a translation of $\mathbf{y}$.

An efficient algorithm for solving Toeplitz systems and calculating one-
dimensional Wiener filters is the Levinson-Durbin recursion [23, 7]. Although
extensions of this algorithm for block Toeplitz systems can be used to
calculate higher dimensional-Wiener filters [33], for computational efficiency
we opt to implement equation (2) in the Fourier-domain as follows:

$\mathbf{v}(\mathbf{x},\mathbf{y})=\mathcal{F}^{-1}\left\\{\frac{\mathcal{F}\left\\{\mathbf{x}\right\\}\otimes\mathcal{F}\left\\{\mathbf{y}\right\\}+\lambda}{\mathcal{F}\left\\{\mathbf{x}\right\\}\otimes\mathcal{F}\left\\{\mathbf{x}\right\\}+\lambda}\right\\}$
(3)

where $\mathcal{F}$ represents the fast Fourier transform (FFT), $\otimes$
represent the Hadamard product, the vinculum represents Hadamard division, and
$\lambda$ is a pre-whitening scalar used to stabilise the deconvolution. Note
that we add $\lambda$ to both numerator and denominator in (3) to ensure that
the Wiener filter $\mathbf{v}(\mathbf{x},\mathbf{y})$ when
$\mathbf{x}=\mathbf{y}$ is given by a Dirac delta function with an amplitude
of 1.

As mentioned in Section 2, adaptive waveform inversion (AWI) leverages a
Wiener filter based loss function to mitigate the non-linearity of seismic
inversion [49]. The AWI objective function is equivalent, up to a scaling
factor, to the Rayleigh quotient,
$\mathcal{R}(\mathbf{T},\mathbf{v}(\mathbf{x}_{\theta},\mathbf{y}))$, defined
as:

$\mathcal{R}(\mathbf{T},\mathbf{v}(\mathbf{x}_{\theta},\mathbf{y}))=\frac{||\mathbf{T}\mathbf{v}(\mathbf{x}_{\theta},\mathbf{y})||^{2}}{||\mathbf{v}(\mathbf{x}_{\theta},\mathbf{y})||^{2}}$
(4)

where $\mathbf{T}$ is termed the penalty matrix, $\mathbf{y}$ is a vector
containing the observed or target data, and $\mathbf{x_{\theta}}$ contains the
predicted data from a model with parameters $\theta$. A diagonal form for
$\mathbf{T}$ is used, representing the Hadamard product with a monotonically
decreasing function centered at zero-lag [49]. This penalises energy in the
Wiener filter that is not at zero-lag, thereby focusing energy towards a zero-
lag spike [49].

As discussed in [49], the normalisation term in the denominator of (4) is
crucial to prevent a trivial solution, where the numerator is minimised by
simply suppressing all of the values in the Wiener filter. The inclusion of
the normalisation term, however, incurs a significant negative consequence; it
renders (4) insensitive to the amplitudes within the Wiener filter and, by
definition, the predicted samples $\mathbf{x_{\theta}}$. Its only objective is
to focus energy to a spike at zero-lag. It does not care whether that spike is
equivalent to a Dirac delta function with an amplitude of 1. Within the
context of seismic imaging, for which AWI was originally developed, this issue
is of little consequence due to the heavy use of normalisation that is applied
to the data before comparison. However, in a more general machine-learning
context, both amplitude and phase information are vital [16], and so we derive
a new loss function that addresses this limitation. A final and important
consideration is that, under the current definition of the normal equation
(2), the Wiener filter in (4) matches the target $\mathbf{y}$ to the
prediction $\mathbf{x_{\theta}}$ through convolution. This is termed reverse
AWI, but it should be noted that an equally permissible formulation, termed
forward AWI, is obtained by instead optimising the Wiener filter that matches
the prediction $\mathbf{x_{\theta}}$ to the target $\mathbf{y}$ [49]. The
results presented in this work use the reverse Wiener filter formulation, in
accordance with (2), though we stress that both formulations are suitably
applicable and demonstrate similar properties.

###  3.2 Wiener Loss

Here we present a new loss function, termed the Wiener loss, that provides a
natural solution to the amplitude insensitivity of the AWI loss (4). The
Wiener loss, $\mathcal{L}_{\textsc{w}}$, is directly motivated by the
convolutional identity; its minimisation is equivalent to maximising the
likelihood of the Dirac delta function under a multivariate Gaussian
distribution whose mean is defined as data-matching full-lag Wiener filters,
thereby forcing the predicted samples towards the target samples in a global
sense. This likelihood is defined as follows:

$p(\boldsymbol{\delta}|\mathbf{v})=\underset{i=1}{\prod}\mathcal{N}(\boldsymbol{\delta}|\mathbf{v}(\mathbf{x}_{\theta}^{(i)},\mathbf{y}^{(i)}),\boldsymbol{\Sigma})$
(5)

where $\mathbf{x}_{\theta}^{(i)}$ is the $i^{\textup{th}}$ predicted sample
from a model with parameters $\theta$, $\mathbf{y}^{(i)}$ is the
$i^{\textup{th}}$ target sample, and $\boldsymbol{\Sigma}$ is the covariance
matrix. Taking the negative log-likelihood and neglecting terms that are not
dependent on the Wiener filters yields:

$-\textup{log}\
p(\boldsymbol{\delta}|\mathbf{v})\propto\underset{i=1}{\sum}\mathcal{L}_{\textsc{w}}(\mathbf{x}_{\theta}^{(i)},\mathbf{y}^{(i)}),$
(6)

where the Wiener loss $\mathcal{L}_{\textsc{w}}$ is defined as:

$\mathcal{L}_{\textsc{w}}(\mathbf{x}_{\theta},\mathbf{y})=\frac{1}{2}||\mathbf{W}\left\\{\mathbf{v}(\mathbf{x}_{\theta},\mathbf{y})-\boldsymbol{\delta}\right\\}||^{2},$
(7)

with the inverse covariance matrix decomposing as
$\boldsymbol{\Sigma}^{-1}=\mathbf{W}^{T}\mathbf{W}$. As can be seen from (7),
the Wiener loss is equivalent to the square of the Mahalanobis distance [29]
between the Dirac delta function and a multivariate Gaussian with mean defined
by the data-matching Wiener filter
$\mathbf{v}(\mathbf{x}_{\theta},\mathbf{y})$ and covariance
$\boldsymbol{\Sigma}$, expressed as a function of samples
$\mathbf{x}_{\theta}$ and $\mathbf{y}$. Due to its least-squares nature, the
Wiener loss retains desirable optimisation properties inherent to MSE, such as
smoothness, convexity and differentiability. The form of matrix $\mathbf{W}$,
which can be interpreted as a whitening matrix, is considered a hyper-
parameter that can be used to tailor the behaviour of the Wiener loss to the
required task. This formulation is comparable to the one presented in (4), but
circumvents the normalisation of the filter to preserve signal amplitudes.

###  3.3 Wiener Diffusion

Here we introduce the Wiener diffusion, a new generative model that uses
Wiener filters to define a non-parametric energy based model, $E(\mathbf{x})$,
which represents the non-normalised negative log-likelihood of the data
distribution. The probability density of $E(\mathbf{x})$ is defined by the
Boltzmann distribution as:

$p(\mathbf{x})=\frac{\textup{exp}(-E(\mathbf{x}))}{\int\textup{exp}(-E(\mathbf{x}))\
d\mathbf{x}}.$ (8)

Direct calculation of the probability density $p(\mathbf{x})$ is
computationally intractable due to the normalisation term in the denominator,
however, it can be efficiently sampled using Langevin Markov Chain Monte Carlo
(MCMC) methods, which leverage the fact that the gradient of the log-energy
density is given by the negative gradient of the energy-based model [41]. The
Langevin MCMC diffusion process used to generate samples for $T$ steps is
defined by:

${\mathbf{x}_{t+1}\leftarrow\mathbf{x}_{t}-\frac{\alpha_{t}}{2}\nabla
E(\mathbf{x}_{t})+\mathbf{z}_{t},\ \textup{for}\ t=0,1,\cdots,T-1}$ (9)

where $\mathbf{x}_{t}$ is the sample undergoing diffusion at step $t$,
$\alpha_{t}$ represents the step size schedule and $\mathbf{z}_{t}$ represents
random noise sampled from a Gaussian distribution
$\mathbf{z}_{t}\sim\mathcal{N}(\mathbf{0},\beta_{t})$ defined by variance
schedule $\beta_{t}$ [36, 11]. The energy function that defines the Wiener
diffusion is given by the following:

$E(\mathbf{x})=\sum_{i=1}\frac{1}{2}\mathcal{R}(\mathbf{T},\mathbf{v}(\mathbf{x},\mathbf{y}_{i}))+\frac{\gamma}{2}||\boldsymbol{\delta}\otimes(\mathbf{v}(\mathbf{x},\mathbf{y}_{i})-\boldsymbol{\delta})||^{2}$
(10)

where the summation is performed over the samples of the dataset, $\gamma$ is
a scalar that controls the strength of the second term, and $\otimes$
represents the Hadamard product. The first term of equation (10), much like in
AWI, drives the data-matching Wiener filters towards zero-lag spikes by
minimising the Rayleigh quotient of a penalty matrix $\mathbf{T}$ and Wiener
filters $\mathbf{v}(\mathbf{x},\mathbf{y}_{i})$. The second term is a
corrective term that ensures the energy based model is sensitive to amplitude
information by honouring the convolutional identity, and encouraging the
zeroth-lag of the filters to have an amplitude of 1. Much like the
regularisation term used by [27], the noise term within the Langevin diffusion
process (9) helps to ensure that the generated samples do not “over-fit” to
the data distribution by collapsing onto the data samples.

The penalty matrix $\mathbf{T}$ defines the morphology of the energy
landscape. The precise form of $\mathbf{T}$ must therefore be carefully chosen
to produce a desirable energy landscape, for which the mechanics of the
gradient-based Langevin diffusion process (9) is able to generate diverse yet
accurate samples from the data distribution. Similarly to AWI, we adopt a
diagonal form for $\mathbf{T}$ that represents the Hadamard product with a
monotonically decreasing penalty function centered at zero-lag. An important
additional requirement is that the gradient of the penalty function must be
monotonically increasing with increasing proximity to the zeroth-lag. This
ensures that the gradient of the energy, and thus the relative size of the
update, is larger in the direction of samples $\mathbf{y}_{i}$ that are
closer, in a convolutional sense, to $\mathbf{x}_{t}$. If this requirement is
not met, the summation over data samples $\mathbf{y}_{i}$ in (10) will act to
produce an average of the data rather than to generate diverse samples.

##  4 Experiments

###  4.1 Autoencoder Image Compression

We initially showcase our Wiener loss (7) as an optimisation metric in a
CelebA [26] autoencoder reconstruction problem. The autoencoder is a shallow
network composed of 3 encoding and 3 decoding convolutional layers (see
suplemmentary material), with a latent size of 128 and Mish [32] activation
function. We construct a random train/validation split of 80/20 on the full
dataset of 202,599 64x64 images, and set up training with a batch size of 512
and a learning rate of $1e^{-3}$ using the Adam optimiser [20]. We formulate
our Wiener loss such that one two-dimensional filter is optimised for each RGB
channel, with each filter approximately twice the size of the image.

Figure 2: Autoencoder reconstruction outputs from the validation dataset.
Comparison of Ground Truth (top row), mean squared-error (middle row), and
Wiener loss (bottom row) as the optimisation functional.

The outputs of the network are compared under two losses, the commonly used
MSE loss, and our Wiener loss with the aforementioned hyper-parameters. We
achieved best results with the Wiener loss implementation using a fixed pre-
whitening scalar, referred to in (3), of $\lambda=250$, and a diagonal form
for the whitening matrix $\mathbf{W}$ defined as a modified two-dimensional
Laplace density function (see supplementary materials). For both losses, 100
epochs were sufficient to achieve convergence. We observe improvements in
high-frequency image recovery with our Wiener loss (Figure 2), evidenced by
superior resolution of facial features such as eyes, mouth and hair. This
demonstrate the high-frequency feature recovery of our new loss formulation
and its ability to increase perceptual quality. The qualitative evaluation is
also corroborated by the metrics presented in Table 1, where an improvement on
the Learned Perceptual Image Patch Similarity (LPIPS) [52] and the Frechet
Inception Distance (FID) [14] are compelling. In contrast, our loss performs
loss on metrics that rely on pixel-wise comparison or take luminance into
account, like mean absolute error (MAE), MSE and SSIM. This can be attributed
to a limitation in amplitude recovery of signals under our proposed loss,
mostly visible by restricted saturation in the recovered images. This
limitation is still not yet well understood. We suspect that different
parametrization of the data and Wiener calculations could potentially overcome
this limitation, which will be explored in further work.

| MAE | MSE | Wiener Loss | SSIM | LPIPS | FID
---|---|---|---|---|---|---
CelebA Autoencoder MSE | $8.2\cdot 10^{-3}$ | $1.07\cdot 10^{-4}$ | $0.32$ | $0.83$ | $0.14$ | $0.63$
CelebA Autoencoder Wiener Loss | $1.12\cdot 10^{-2}$ | $1.38\cdot 10^{-4}$ | $8.54\cdot 10^{-2}$ | $0.81$ | $0.11$ | $0.25$
MRI Imputation Bicubic | $4.09\cdot 10^{-4}$ | $2.6\cdot 10^{-4}$ | $0.21$ | $0.8$ | $0.32$ | $4.12$
MRI Imputation MSE | $1.93\cdot 10^{-4}$ | $1.35\cdot 10^{-4}$ | $0.14$ | $0.91$ | $0.11$ | $1.66$
MRI Imputation Wiener Loss | $2.22\cdot 10^{-4}$ | $1.71\cdot 10^{-4}$ | $9.16\cdot 10^{-2}$ | $0.9$ | $7.93\cdot 10^{-2}$ | $0.23$
Table 1: Comparison metrics for results obtained in the CelebA autoencoder
example and MRI imputation. Metrics presented are averaged per sample and
evaluated on the validation dataset. Our method is capable of better restoring
the distribution of the ground-truth data, evidenced by improvements in the
LPIPS [52] and FID [14] metrics (highlighted in bold), but perform poorer
under mean absolute error (MAE), MSE and SSIM due to limitations in brightness
and luminance recoveries.

###  4.2 Image Imputation

Next, we apply our Wiener loss (7) to the imputation of medical images using
neural networks. Often, engineering datasets have sparse or missing
information, which are difficult to avoid due to poor design, unexpected
failures or restrictions in acquisition times (e.g. [25, 45, 54, 24, 2, 6]).
Interpolation is a popular strategy to imitate or simulate missing data, but
it is prone to failure when the amount of missing information is large [1].
Imputation [5] is an alternative strategy to estimate missing values using
statistics from across the dataset. In this experiment we focus on sparse
medical images, in particular Magnetic Resonance Imaging (MRI), which due to
time constraints are often sparse [5, 4], have varying resolutions [40], or
limited fields of view [2]. We propose i) that a supervised U-net [39],
trained using the Wiener loss, is capable of imputing missing data, and ii)
that the Wiener loss improves resolution of the imputed data when compared to
MSE and standard interpolation methods. Improving the quality of MRI
imputation will have knock-on effects for post-processing and image analysis
of clinical datasets [2]. We also expect imputation using a U-net [39] and
Wiener loss will be applicable to fields where sensors have variable
resolution in different dimensions [54, 24] and to super-resolution problems
in general [19].

Figure 3: Ground truth (a-b) and zoomed outputs from imputing under-sampled
scans (c) using bicubic interpolation (d), U-net [39] with MSE (e) and U-net
[39] with Wiener loss (f). The Wiener loss is able to recover finer structures
of the scan that are aliased or blurred in other methods. Results of (e-f) are
obtained from the validation set.

We use the MGH-USC Human Connectome Project (HCP) Young Adult database [44] to
test our claims. The dataset is comprised of 553 3D single-channel T1-weighted
MRI scans of healthy adult brains, which we reduce to 2D by retrieving the
mid-sagittal plane. The 2D scans are single-channel and of size 256 x 256. The
dataset has also been de-identified with normalised anterior filtering [31] in
order to remove identifying and sensitive anatomical features of patients. Our
network is a U-net [39] composed of 3 residual units, sequenced channels of
sizes 16, 32 and 64 (see supplementary materials), and Mish [32] as activation
function. Acquisition sparsity is simulated by applying a sampling mask to the
full-resolution scans, serving as input to our network (Figure 3). The masked
input samples have their outliers removed by capping values above the data’s
upper quartile range, and are scaled between 0 and 1 through a minimum-maximum
value-based normalisation method. We do so to prevent statistical bias while
preserving the distribution of the original data. Similarly to our previous
example, we used a train/validation split of 80/20, a modified two-dimensional
Laplace density function as whitening matrix $\mathbf{W}$ (see supplementary
materials), and a fixed pre-whitening scalar $\lambda=250$ for the 2D filter
computation. We use a batch size of 32, and employ the Adam optimiser [20]
with a learning rate of $1e^{-2}$. We then compare the output of the model to
the target scan and use this measure to train our network in a supervised
setting. Results from using a standard bicubic interpolation, the MSE loss,
and our Wiener loss for the U-net [39] imputation are presented in Figure 3.
For both MSE and Wiener loss experiments, 300 epochs achieved model
convergence.

In this heavily sparse problem, the results presented in Figure 3 show
substantial resolution improvements from using a supervised neural network to
perform the imputation task. Generalisation errors have been suppressed
despite the small dataset size. The use of our Wiener loss promotes imputation
outputs with higher resolution of smaller features, e.g. micro structures of
the cerebellum (see zoom), as well as continuity and sharpness of the skull,
when compared to outputs from the MSE optimisation and bicubic interpolation.
As previously mentioned, the reconstructions from the Wiener loss suffer from
luminance and brightness recovery constraints, performing with poorer scores
on the MAE, MSE, and SSIM metrics (Table 1). However, our Wiener loss excels
at recovering perceptual quality and is better suited to restore the ground-
truth distribution of the data. This is evidenced by lower scores on LPIPS
[52] and FID [14] metrics. The results of our loss have increased scan
resolution and contrast, which can have profound impacts during clinical
analysis.

###  4.3 Generative Modelling

The third application we present is the task of image generation using the
proposed Wiener diffusion. We validate this unique approach and demonstrate
its behaviour by applying it to the latent space of convolutional autoencoders
trained on two well known datasets, MNIST [22] and CelebA [26]. Although we
showcase latent space generation here, we emphasise that the Wiener diffusion
can be applied directly to natural images; this will be the focus of future
research. The autoencoder architectures used comprise 6 convolutional layers,
2 dense layers, and employ the Leaky ReLU activation function [28], with a
latent space bottleneck that is a vector of length 31 for MNIST, and 51 for
CelebA. For both experiments, we use 20,000 samples to define the energy
function (10) for the Wiener diffusion, with the initial sample
$\mathbf{x}_{t=0}$ drawn from a normal distribution with a variance equivalent
to the variance of the latent bottleneck. For the Wiener filter computation,
we employ a pre-whitening scalar of $\lambda=1e^{-1}$, and use a modified one-
dimensional Laplace density function to define the diagonal values of
$\mathbf{T}$ (see supplementary materials). The MNIST diffusion is performed
over $T=200$ steps, with a cosine step schedule decreasing from
$\alpha_{t}=500$ to $\alpha_{t}=1$, and a cosine variance schedule increasing
from $\beta_{t}=0.01$ to $\beta_{t}=0.8$. For CelebA, the diffusion is run for
$T=400$ steps, with a cosine step schedule decreasing from $\alpha_{t}=3000$
to $\alpha_{t}=300$ and a cosine variance schedule increasing from
$\beta_{t}=0.1$ to $\beta_{t}=4$. The generated latent samples are visualised
in Figure 4 using the corresponding decoders, alongside a visualisation
demonstrating how the data-matching Wiener filters are progressively focused
towards zero-lag delta functions throughout the diffusion process.

We stress that Wiener diffusion is a non-parametric method and, although we
use it on the latent space of trained autoencoders, the generative process
does not require the training of a model. The benefits of this are that it is
well suited to scenarios with limited data, and is also a good candidate for
controllable generation as the diffusion process can be conditioned, at
inference time, by specifying the data that defines the energy function (10);
future work will explore these hypothesized advantages.

Figure 4: Latent samples generated using the Wiener diffusion, visualised
using the decoder, and an illustration of the diffusion process alongside the
30 closest Wiener filters to the sample $\mathbf{x}_{t}$ (as defined by (10))
for MNIST (top) and CelebA (bottom).

###  4.4 Image Classification

In this final experiment, we further evidence our claim that Wiener filters
provide a versatile and customisable tool for data comparison. The task we
explore is the classification of digits from the affNIST dataset [43], which
consist of padded MNIST images that have undergone affine transformations,
using only the original (but equivalently padded) MNIST data as the training
set (see Figure 5). We first evaluate a baseline model; the k-nearest
neighbours (KNN) algorithm [3]. Using a 5-fold cross-validated grid-search on
the training set, the optimal configuration of the baseline KNN [3] was
obtained with the Manhattan distance metric and a k of 3. This model entirely
failed to classify samples from the affNIST test set, achieving an accuracy of
$\mathbf{14}\boldsymbol{\%}$. The reason for such poor performance is due to
translation, an affine transformation that, from Figure 5, appears to be
responsible for the predominant differences between the training and testing
data distributions.

Figure 5: Image classification task to use the MNIST training data (left) to
predict the affNIST testing data (right). A KNN classifier with the Manhattan
distance metric achieves and accuracy of $\mathbf{14}\boldsymbol{\%}$, whereas
a KNN classifier with our proposed translation invariant metric
$\mathcal{L}_{\textsc{ti}}$ achieves an accuracy of
$\mathbf{61}\boldsymbol{\%}$

.

We therefore formulate a translation invariant comparison function
$\mathcal{L}_{\textsc{ti}}$, by considering the maximum value of a
standardised data-matching Wiener filter as:

$\mathcal{L}_{\textsc{ti}}(\mathbf{x},\mathbf{y})=-\textup{max}\left(\frac{\mathbf{v}(\mathbf{x},\mathbf{y})-\mu_{\mathbf{v}}}{\sigma_{\mathbf{v}}}\right)$
(11)

where $\mu_{\mathbf{v}}$ is the mean value of filter
$\mathbf{v}(\mathbf{x},\mathbf{y})$, and $\sigma_{\mathbf{v}}$ is its standard
deviation. If the energy of $\mathbf{v}(\mathbf{x},\mathbf{y})$ is focused in
a narrow spike, then the process of standardisation amplifies the magnitude of
the spike to ensure unit standard deviation. On the other hand, if the energy
of $\mathbf{v}(\mathbf{x},\mathbf{y})$ is unfocused and distributed across all
lags, then unit standard deviation is achieved without requiring a large
maximum value in the standardised filter. $\mathcal{L}_{\textsc{ti}}$ is
therefore only sensitive to whether the energy in the filter is focused or
not; it does not care where the energy is focused in the filter and, as is
illustrated in Figure 1, is thus invariant to translation. Using
$\mathcal{L}_{\textsc{ti}}$ as a distance metric in KNN [3] , with a k of 10
on the same task achieves a dramatically improved test accuracy of
$\mathbf{61}\boldsymbol{\%}$. A promising direction for future research is to
use the ideas discussed here to formulate optimisation metrics for deep
learning where translational invariance is required.

##  5 Considerations

Throughout experimentation we observe training and performance sensitivity to
the choice of hyper-parameters that govern the application of our methods.
These include: i) the whitening matrix $\mathbf{W}$ that controls the
behaviour of the Wiener loss (7), ii) the penalty matrix $\mathbf{T}$ that
defines the shape of the energy landscape in (10), and iii) the pre-whitening
scalar $\lambda$ that stabilises the Wiener filter computation (3). Future
work will explore the systematic change in behaviour as a result of varying
these hyper-parameters. For instance, using optimal whitening matrices (as
discussed by [18]), investigating non-diagonal penalty functions, and choosing
different pre-whitening values for the numerator vs the denominator which can
amplify or suppress particular frequencies.

Finally, we found that the Fourier-domain implementation of the Wiener loss
remains computationally competitive to when compared to the MSE loss (see
Figure in supplementary materials). For MSE, the computational complexity is
$O(n)$ and the memory scales as $O(n^{2})$. For the Wiener loss
implementation, computational complexity is $O(log\,n)$ when implemented in
the Fourier domain using FFT and $O(n^{2})$ in the data domain using Levinson-
Durbin recursion; while memory scales as $O(n^{2})$ in the Fourier domain and
$O(n^{4})$ in the data domain. We note here that both Fourier and data domain
implementations are valid and produce equivalent results, but the later proved
to be computationally intractable for most practical deep-learning problems.
We also note that the current implementations are not fully optimised and that
performance can be further improved.

##  6 Conclusion

Despite the immense evolution of machine learning over the last decade in
terms of architecture complexity, limited progress has been made in developing
optimal strategies to quantify data comparisons effectively. We address this
issue by proposing a new method based on Wiener filters that promote
contextual awareness in data samples. At its core lays the principle of
convolutional identity as a measure of similarity between data sample pairs.

The results we present here indicate that this formulation is effective in a
number of unrelated tasks. We demonstrate enhanced perceptual quality in
autoencoder reconstructions, resolution improvements in imputation problems,
translation invariance in classification tasks, and define a novel non-
parametric generative model.

Our main contribution is a conceptual principle that holds immense potential
across computer vision problems, and more general machine learning
applications. So far, we have demonstrated its potential and problem-agnostic
nature in four different tasks, and we are currently working towards improving
our understanding of the fundamental assumptions that underlay its apparent
superior performance.

## References

  * [1] Iman Aganj, Boon Thye Thomas Yeo, Mert R Sabuncu, and Bruce Fischl.  On removing interpolation and resampling artifacts in rigid image registration.  IEEE Transactions on Image Processing, 22(2):816–827, 2012.
  * [2] Mikael Brudfors, Yael Balbastre, Parashkev Nachev, and John Ashburner.  A tool for super-resolving multimodal clinical mri.  arXiv preprint arXiv:1909.01140, 2019.
  * [3] Thomas Cover and Peter Hart.  Nearest neighbor pattern classification.  IEEE transactions on information theory, 13(1):21–27, 1967.
  * [4] Adrian V Dalca, Katherine L Bouman, William T Freeman, Natalia S Rost, Mert R Sabuncu, and Polina Golland.  Population based image imputation.  In International Conference on Information Processing in Medical Imaging, pages 659–671. Springer, 2017.
  * [5] Adrian V Dalca, John Guttag, and Mert R Sabuncu.  Unsupervised data imputation via variational inference of deep subspaces.  arXiv preprint arXiv:1903.03503, 2019.
  * [6] Chao Dong, Chen Change Loy, Kaiming He, and Xiaoou Tang.  Image super-resolution using deep convolutional networks.  IEEE transactions on pattern analysis and machine intelligence, 38(2):295–307, 2015.
  * [7] James Durbin.  The fitting of time-series models.  Revue de l’Institut International de Statistique, pages 233–244, 1960.
  * [8] Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio.  Learning with a wasserstein loss.  Advances in neural information processing systems, 28, 2015.
  * [9] Bernd Girod.  What’s wrong with mean-squared error?  Digital images and human vision, pages 207–220, 1993.
  * [10] 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, 63(11):139–144, 2020.
  * [11] Ulf Grenander and Michael I Miller.  Representations of knowledge in complex systems.  Journal of the Royal Statistical Society: Series B (Methodological), 56(4):549–581, 1994.
  * [12] Per Christian Hansen.  Deconvolution and regularization with toeplitz matrices.  Numerical algorithms, 29(4):323–378, 2002.
  * [13] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
  * [14] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  Advances in neural information processing systems, 30, 2017.
  * [15] Anil K Jain.  Fundamentals of digital image processing.  Prentice-Hall, Inc., 1989.
  * [16] Liming Jiang, Bo Dai, Wayne Wu, and Chen Change Loy.  Focal frequency loss for image reconstruction and synthesis.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 13919–13929, 2021.
  * [17] Justin Johnson, Alexandre Alahi, and Li Fei-Fei.  Perceptual losses for real-time style transfer and super-resolution.  In European conference on computer vision, pages 694–711. Springer, 2016.
  * [18] Agnan Kessy, Alex Lewin, and Korbinian Strimmer.  Optimal whitening and decorrelation.  The American Statistician, 72(4):309–314, 2018.
  * [19] Jiwon Kim, Jung Kwon Lee, and Kyoung Mu Lee.  Accurate image super-resolution using very deep convolutional networks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1646–1654, 2016.
  * [20] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [21] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  Proceedings of the IEEE, 86(11):2278–2324, 1998.
  * [22] Yann LeCun and Corinna Cortes.  MNIST handwritten digit database.  2010\.
  * [23] Norman Levinson.  The wiener (root mean square) error criterion in filter design and prediction.  Journal of Mathematics and Physics, 25(1-4):261–278, 1946.
  * [24] Keqiang Li, Xiao Wang, Youchun Xu, and Jianqiang Wang.  Density enhancement-based long-range pedestrian detection using 3-d range data.  IEEE Transactions on Intelligent Transportation Systems, 17(5):1368–1380, 2015.
  * [25] Bin Liu and Mauricio D Sacchi.  Minimum weighted norm interpolation of seismic records.  Geophysics, 69(6):1560–1568, 2004.
  * [26] 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.
  * [27] Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian-Robert Stöter.  Sliced-wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions.  In International Conference on Machine Learning, pages 4104–4113. PMLR, 2019.
  * [28] Andrew L Maas, Awni Y Hannun, Andrew Y Ng, et al.  Rectifier nonlinearities improve neural network acoustic models.  In Proc. icml, volume 30, page 3. Atlanta, Georgia, USA, 2013.
  * [29] Goeffrey J McLachlan.  Mahalanobis distance.  Resonance, 4(6):20–26, 1999.
  * [30] Ludovic Métivier, Romain Brossier, Quentin Merigot, and Edouard Oudet.  A graph space optimal transport distance as a generalization of lp distances: application to a seismic imaging inverse problem.  Inverse Problems, 35(8):085001, 2019.
  * [31] Mikhail Milchenko and Daniel Marcus.  Obscuring surface anatomy in volumetric imaging data.  Neuroinformatics, 11(1):65–75, 2013.
  * [32] Diganta Misra.  Mish: A self regularized non-monotonic neural activation function.  arXiv preprint arXiv:1908.08681, 4(2):10–48550, 2019.
  * [33] Bruce Ronald Musicus.  Levinson and fast choleski algorithms for toeplitz and almost toeplitz matrices.  1988\.
  * [34] Jean-François Pambrun and Rita Noumeir.  Limitations of the ssim quality metric in the context of diagnostic imaging.  In 2015 IEEE International Conference on Image Processing (ICIP), pages 2960–2963. IEEE, 2015.
  * [35] Thrasyvoulos N Pappas, Robert J Safranek, and Junqing Chen.  Perceptual criteria for image quality evaluation.  Handbook of image and video processing, 110, 2000.
  * [36] Giorgio Parisi.  Correlation functions and computer simulations.  Nuclear Physics B, 180(3):378–384, 1981.
  * [37] Gabriel Peyré, Marco Cuturi, et al.  Computational optimal transport: With applications to data science.  Foundations and Trends® in Machine Learning, 11(5-6):355–607, 2019.
  * [38] Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville.  On the spectral bias of neural networks.  In International Conference on Machine Learning, pages 5301–5310. PMLR, 2019.
  * [39] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In International Conference on Medical image computing and computer-assisted intervention, pages 234–241. Springer, 2015.
  * [40] François Rousseau, Alzheimer’s Disease Neuroimaging Initiative, et al.  A non-local approach for image super-resolution using intermodality priors.  Medical image analysis, 14(4):594–605, 2010.
  * [41] Yang Song and Diederik P Kingma.  How to train your energy-based models.  arXiv preprint arXiv:2101.03288, 2021.
  * [42] Matthew Thorpe, Serim Park, Soheil Kolouri, Gustavo K Rohde, and Dejan Slepčev.  A transportation $\mathrm{L}^{p}$ distance for signal analysis.  Journal of mathematical imaging and vision, 59:187–210, 2017.
  * [43] Tijmen Tieleman.  affnist.  https://www.cs.toronto.edu/~tijmen/affNIST/, 2013.
  * [44] David C Van Essen, Kamil Ugurbil, Edward Auerbach, Deanna Barch, Timothy EJ Behrens, Richard Bucholz, Acer Chang, Liyong Chen, Maurizio Corbetta, Sandra W Curtiss, et al.  The human connectome project: a data acquisition perspective.  Neuroimage, 62(4):2222–2231, 2012.
  * [45] Benfeng Wang, Ning Zhang, Wenkai Lu, and Jialin Wang.  Deep-learning-based seismic data interpolation: A preliminary result.  Geophysics, 84(1):V11–V20, 2019.
  * [46] Zhou Wang and Alan C Bovik.  Modern image quality assessment.  Synthesis Lectures on Image, Video, and Multimedia Processing, 2(1):1–156, 2006.
  * [47] Zhou Wang and Alan C Bovik.  Mean squared error: Love it or leave it? a new look at signal fidelity measures.  IEEE signal processing magazine, 26(1):98–117, 2009.
  * [48] Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli.  Image quality assessment: from error visibility to structural similarity.  IEEE transactions on image processing, 13(4):600–612, 2004.
  * [49] Michael Warner and Lluís Guasch.  Adaptive waveform inversion: Theory.  Geophysics, 81(6):R429–R445, 2016.
  * [50] Hong Ren Wu and Kamisetty Ramamohan Rao.  Digital video image quality and perceptual coding.  CRC press, 2017.
  * [51] Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma.  Frequency principle: Fourier analysis sheds light on deep neural networks.  arXiv preprint arXiv:1901.06523, 2019.
  * [52] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang.  The unreasonable effectiveness of deep features as a perceptual metric.  In CVPR, 2018.
  * [53] Hang Zhao, Orazio Gallo, Iuri Frosio, and Jan Kautz.  Loss functions for image restoration with neural networks.  IEEE Transactions on computational imaging, 3(1):47–57, 2016.
  * [54] Junxuan Zhao, Hao Xu, Hongchao Liu, Jianqing Wu, Yichen Zheng, and Dayong Wu.  Detection and tracking of pedestrians and vehicles using roadside lidar sensors.  Transportation research part C: emerging technologies, 100:68–87, 2019.