# Temporal Wasserstein non-negative matrix factorization for non-rigid motion
segmentation and spatiotemporal deconvolution

Erdem Varol1,3,4, Amin Nejatbakhsh2,3,4, Conor McGrory3,4
Department of Statistics1, Department of Neurobiology and Behavior2
Center for Theoretical Neuroscience3, Zuckerman Institute4
Columbia University

###### Abstract

Motion segmentation for natural images commonly relies on dense optic flow to
yield point trajectories which can be grouped into clusters through various
means including spectral clustering or minimum cost multicuts. However, in
biological imaging scenarios, such as fluorescence microscopy or calcium
imaging, where the signal to noise ratio is compromised and intensity
fluctuations occur, optical flow may be difficult to approximate. To this end,
we propose an alternative paradigm for motion segmentation based on optimal
transport which models the video frames as time-varying mass represented as
histograms. Thus, we cast motion segmentation as a temporal non-linear matrix
factorization problem with Wasserstein metric loss. The dictionary elements of
this factorization yield segmentation of motion into coherent objects while
the loading coefficients allow for time-varying intensity signal of the moving
objects to be captured. We demonstrate the use of the proposed paradigm on a
simulated multielectrode drift scenario, as well as calcium indicating
fluorescence microscopy videos of the nematode Caenorhabditis elegans (C.
elegans). The latter application has the added utility of extracting neural
activity of the animal in freely conducted behavior.

Figure 1: Schematic of the proposed spatio-temporal matrix factorization
model. Pixels of correlated movement are locally modelled as being c  Figure
2: Simulated results of two dimensionmal random walk of five particles Top:
Simulated particle data, Middle: Spatio-temporal footprints of NMF components
colored by component, Bottom: Spatio-temporal footprints of TWNMF components,
colored by component

##  1 Introduction

In the past five years, new imaging techniques[14] have made it possible to
record 3D volumetric data from whole animals with high speed and single-neuron
resolution. When used with calcium indicators, these techniques can be used to
observe large-scale patterns of neural activity with high spatial and temporal
precision. Extracting traces of neural activity from this data is a difficult
problem that has been addressed using a variety of approaches. Recently, non-
negative matrix factorization (NMF) methods[26] have been proposed for solving
this problem. These methods extract neural activity from imaging data by
building a low-rank representation of the data, expressed as a product of two
matrices that encode the spatial footprints of neurons and their temporal
firing activity respectively. This general approach has proven to be very
effective for denoising and demixing calcium signals in a variety of
applications. However, one shortcoming of traditional NMF methods is their
assumption that the spatial footprint of neurons is invariant across time. In
whole-animal recordings of behaving C. elegans, this assumption is unrealistic
due to the non-linear movement of the worm, causing NMF methods to perform
poorly at extracting calcium signals. Tracking methods that explicitly model
neuron motion also fail at this task, because they require high-fidelity
detection of neurons, which is often not realizable in biological signal-to-
noise regimes. Instead, to solve this problem, we propose a novel method
called Temporal Wasserstein Non-negative Matrix Factorization (TWNMF), a
nonlinear extension of traditional NMF methods that allows the spatial
footprint of neurons to evolve in time. We describe how this approach can be
used to track and extract calcium traces from moving neurons and demonstrate
its effectiveness on imaging data of C. elegans. We also show how this
technique can be used for the slightly related problem of extracting voltage
traces from multielectrode array data in the presence of motion drift.

###  1.1 Paper organization

In section 2, we briefly revised the prior work on motion segmentation and
optimal transport. In section 3, we pedagogically build up our formulation
using the classical non-negative matrix factorization and its variants using
the Wasserstein metric. In section 4, we demonstrate three important
applications of our method: multielectrode array signal demixing under motion
drift (4.1), calcium signal recovery in moving C.elegans worms (4.2), and
motion segmentation of freely moving whole worm body (4.3). In section 5, we
discuss extensions and computational bottlenecks of the proposed method.

##  2 Related work and contributions

###  2.1 Motion Segmentation

Motion segmentation is the problem of automatically identifying independently
moving objects in a video.[25] Early methods such as [21] used factorization
approaches to accomplish this, under the assumption that the objects in
question were undergoing rigid motion. More recent approaches solve this
problem by estimating optical flow – the inferred ”motion” of pixels between
successive images. [1] Optical flow can be estimated in a variety of ways,
including using deep neural networks [8]. Once the optical flow between each
successive pair of images has been estimated, objects can be identified by
either clustering the long-term point trajectories derived from the flows [13]
[23], modeling the pixels and trajectories as a graph and computing minimum-
cost multi-cuts [9], or using additional features to group the trajectories
[3] [20] [6]. The major problem with applying these approaches to calcium
imaging data is that they perform badly when the pixel intensity of the moving
objects in the image changes rapidly, which happens in scenarios where calcium
indicators are used to measure neural activity.

###  2.2 Optimal Transport Methods

Optimal transport theory [22] has been recently applied to image processing
problems, specifically through the use of the Wasserstein distance as a
similarity metric between images. Cuturi et al. [5] use this metric to compute
prototypical images from sets of similar images. Rolet et al. [16] and Quian
et al. [15] use it as the loss function for non-negative matrix factorization
methods, an approach that we call Wasserstein non-negative matrix
factorization (WNMF). Our proposed method here is an extension of the WNMF
algorithm that incorporates temporal connections and constraints.

###  2.3 Contributions

Our contributions are three-fold:

  * •

We cast non-rigid motion segmentation as a first-principles non-linear matrix
factorization problem utilizing the geometric structure that the Wasserstein
metric provides for data fidelity terms.

  * •

The basis elements resulting from the matrix factorization provide spatio-
temporal footprints of moving objects while the temporal coefficients capture
the temporal appearance intensities effectively modeling motion segmentation
and spatio-temporal deconvolution simultaneously.

  * •

We apply the proposed technique to demixing simulated voltage data in motion
drifting multielectrode scenarios.

  * •

We demonstrate the ability to capture neural activity in the moving worm C.
elegans in a completely unsupervised way for the first time.

##  3 Method

First we introduce notation. Let $\textbf{X}\in\mathbb{R}^{D\times T}$ be a
non-negative matrix containing the imaging data, with $T$ frames of $d$ pixels
or voxels each. The goal of non-negative matrix factorization (NMF) is to
represent X as a product of two non-negative matrices
$\textbf{D}\in\mathbb{R}^{D\times K}$ and $\textbf{C}\in\mathbb{R}^{K\times
T}$, where $K\ll d$, effectively compressing X to a lower-dimensional latent
space. Each column of $\mathbf{D}$ represents the spatial footprint of one of
the $K$ latent components, and each row of $\mathbf{C}$ represents its
activity across time. To compute $\mathbf{D}$ and $\mathbf{C}$ for a given
$\mathbf{X}$, we minimize the cost function
$E_{NMF}(\mathbf{D},\mathbf{C})=\|\mathbf{X}-\mathbf{DC}\|_{F}^{2}$, subject
to the constraint that $\mathbf{D}$ and $\mathbf{C}$ are non-negative. This
can be done efficiently using a simple iterative procedure [10]. If the
signal-to-noise ratio in our data is sufficiently large, the components we
recover correspond to spatial regions with correlated temporal activity, with
$\mathbf{D}$ encoding their spatial footprints and $\mathbf{C}$ encoding their
relative intensities over time. However, one of the assumptions in NMF is that
the spatial footprint of time-varying activity is fixed in time and does not
model the possibility of motion.

Algorithm 1 [$\bm{\alpha}$,$\bm{\beta}$,$\bm{\pi}$]=
SINKHORN(${\bm{a}},{\bm{b}},{\mathbf{M}},\gamma$)

Input:  ${\bm{a}},{\bm{b}}\in\mathbb{R}^{d}$ such that ${\bm{a}}\geq
0,{\bm{b}}\geq 0$, ${\bm{a}}^{T}\bm{1}=1,{\bm{b}}^{T}\bm{1}=1$ (input
histograms), ${\mathbf{M}}\in\mathbb{R}^{d\times d}$ (ground cost),
$\gamma\geq 0$ (entropic regularization)

Set: $\mathbf{K}=e^{-\gamma{\mathbf{M}}}$,
$\bm{u}=\bm{1}_{d}$,$\bm{v}=\bm{1}_{d}$

while not converged do

$\bm{u}\leftarrow{\bm{a}}/{\mathbf{K}}\bm{v}$ (elementwise)

$\bm{v}\leftarrow{\bm{b}}/{\mathbf{K}}\bm{u}$ (elementwise)

end whilereturn Dual variables $\bm{\alpha}=-\gamma\log{\bm{u}}$,
$\bm{\beta}=-\gamma\log{\bm{v}}$, transportation plan
$\bm{\pi}=diag({\bm{u}}){\mathbf{K}}\text{diag}({\bm{v}})$

Algorithm 2 Temporal Wasserstein Non-negative Matrix Factorization (TWNMF)

Input:  Signal: ${\mathbf{X}}\in{\mathbb{R}}^{d\times T}$, ground cost:
${\mathbf{M}}\in{\mathbb{R}}^{d\times d}$, entropic regularization:
$\gamma>0$, Temporal entropy penalty: $\lambda_{T}\geq 0$, Spatial entropy
penalty: $\lambda_{D}$,Step size: $\nu\geq 0$

Require: Data normalization: ${\mathbf{X}}_{t}^{T}\bm{1}=1$, ${\mathbf{X}}\geq
0$

Initialization:${\mathbf{D}}_{\ell}\leftarrow\text{Unif}[0,1]^{d\times K}$
(column sum normalized),
${\mathbf{C}}_{\ell}\leftarrow\text{Unif}[0,1]^{K\times|T_{\ell}|}$ (row sum
normalized such that ${\mathbf{D}}_{\ell}{\bm{c}}_{t}$ is a histogram)

while not converged do

for $\ell=1,\ldots,L$ do

for $t\in T_{\ell}$ do

$[\bm{\alpha}^{\ell,t},\bm{\beta}^{\ell,t},\bm{\pi}^{\ell,t}]\leftarrow$SINKHORN(${\bm{x}}_{t}$,${\mathbf{D}}_{\ell}{\bm{c}}_{t}$,${\mathbf{M}}$,$\gamma$)

$[\bm{\psi}^{\ell,k},\bm{\omega}^{\ell,k},\bm{\rho}^{\ell,t}]\leftarrow$SINKHORN(${\mathbf{D}}_{\ell}^{k}$,${\mathbf{D}}_{\ell+1}^{k}$,${\mathbf{M}}$,$\gamma$)

for k=1,…,K do

${\mathbf{D}}_{\ell}^{k}\leftarrow[{\mathbf{D}}_{\ell}^{k}-\nu\nabla_{{\mathbf{D}}_{\ell}^{k}}\mathcal{L}({\mathbf{D}},{\mathbf{C}},\bm{\pi},\bm{\rho},\bm{\alpha},\bm{\beta},\bm{\psi},\bm{\omega})]_{+}$

${\bm{c}}_{t}\leftarrow[{\bm{c}}_{t}-\nu\nabla_{{\bm{c}}_{t}}\mathcal{L}({\mathbf{D}},{\mathbf{C}},\bm{\pi},\bm{\rho},\bm{\alpha},\bm{\beta},\bm{\psi},\bm{\omega})]_{+}$

Column normalize ${\mathbf{D}}_{\ell}^{k}$

Row normalize ${\bm{c}}_{t}$ such that ${\mathbf{D}}_{\ell}^{k}{\bm{c}}_{t}$
is a histogram

end for

end for

end for

end while

return Spatiotemporal footprints: $\\{{\mathbf{D}}_{\ell}\\}_{\ell=1}^{L}$,
spatiotemporal loading coefficients: $\\{{\mathbf{C}}_{\ell}\\}_{\ell=1}^{L}$,
Dictionary-to-data segmentation push-forwards: $\\{\bm{\rho}^{\ell,t}\\}$

Figure 3: One dimensional data simulating multielectrode array motion drift,
x-axis in subplots denotes the electrode positions, while y-axis denotes the
voltage captured across time. Rows correspond to spatio-temporal footprints of
five different neurons measured by the electrode array. 1st column: The ground
truth spatio-temporal footprints of five neurons, 2nd column: Spatio-temporal
footprints of TWNMF components, 3rd column: Temporal components of WNMF (blue)
compared with actual voltage signal (red), 4th column Spatio-temporal
footprints of WNMF[16] components, 5th column: Temporal components of WNMF
(blue) compared with actual voltage signal (red), 6th column:Spatio-temporal
footprints of NMF[10] components, 7th column: Temporal components of NMF
(blue) compared with actual voltage signal (red) Figure 4: Simulated
multielectrode array voltage data.

###  3.1 Wasserstein distance

Wasserstein distances have gained popularity in recent years as a way to
measure the similarity between images. These distances work by solving
different forms of the optimal transport problem, a well-studied problem that,
in its most general form, asks how to efficiently transport mass from one
spatial distribution to another, given an underlying cost metric[22]. To
compute the OT distance between two images, we treat pixel intensity as mass
and use the euclidean distance between pixels as the transportation cost.
Computing this exactly is prohibitively expensive for most image-processing
applications, but fast and accurate approximations [4] exist.

For two input histograms ${\bm{a}},{\bm{b}}\in\mathbb{R}^{d}$, the Wasserstein
metric $W_{1}$ can be expressed as the solution to a constrained linear
program:

$\displaystyle
W_{1}({\bm{a}},{\bm{b}})=\underset{{\mathbf{P}}}{\min}~{}\sum_{i,j=1}^{d}m_{i,j}P_{i,j}$
$\displaystyle\text{subject
to}~{}\sum_{j=1}^{d}P_{i,j}={\bm{a}},~{}~{}\sum_{i=1}^{d}P_{i,j}={\bm{b}}$ (1)

where $m_{i,j}$ is the ground cost of moving a unit of mass from histogram bin
$i$ to histogram bin $j$. The ground cost is typically set to be the euclidean
distance of the pixel locations i.e.
$m_{i,j}=\|{\bm{p}}_{i}-{\bm{p}}_{j}\|_{2}$ where ${\bm{p}}_{i}$ is the
coordinates of the ith pixel. Solving this linear program in general requires
$O(d^{3})$ operations and furthermore may not yield a unique optima since it
is possible for entire facets of the feasible region polytope (i.e.
$\\{\mathcal{P}:{\mathbf{P}}\in\mathbb{R}^{d\times
d}~{}~{}|~{}~{}{\mathbf{P}}\bm{1}={\bm{a}},{\mathbf{P}}^{T}\bm{1}={\bm{b}}\\}$
to yield a minimum cost solution.

Cuturi [4] has proposed a relaxation of equation 3.1 by introducing an
entropic regularizer:

$\displaystyle
W_{1}^{\gamma}({\bm{a}},{\bm{b}})=\underset{{\mathbf{P}}}{\min}~{}\sum_{i,j=1}^{d}m_{i,j}P_{i,j}+\gamma(P_{i,j}\log
P_{i,j}-1)$ $\displaystyle\text{subject
to}~{}\sum_{j=1}^{d}P_{i,j}={\bm{a}},~{}~{}\sum_{i=1}^{d}P_{i,j}={\bm{b}}$ (2)

which is strongly convex and can be optimized using fixed point Sinkhorn
algorithm[18] iterations in $O(d\log d)$ operations. The outline of the
Sinkhorn algorithm is provided in algorithm 1.

Figure 5: Top: Moving C. elegans fluorescence microscopy video, 2nd row:
Zoomed region, 3rd row: Spatio-temporal footprints of NMF components colored
by component, 4th row: Spatio-temporal footprints of W-NMF components, colored
by component, Bottom:  Calcium traces extracted from NMF (left) and W-NMF
(right) components.

###  3.2 Wasserstein non-negative matrix factorization

While the euclidean loss takes into account the magnitude difference two
vectors, it does not take into account the underlying manifold that the bases
of the vectors reside in. On the other hand, the wasserstein metric not only
takes into account the differences in the loadings of the basis vectors but
also the distances of the basis vectors itself. To illustrate this, let
${\bm{a}}=[1,0,0]$,${\bm{b}}=[0,1,0]$ and ${\bm{c}}=[0,0,1]$ be three vectors
in $\mathbb{R}^{3}$. Under the euclidean metric $d_{E}$,
$d_{E}({\bm{a}},{\bm{b}})=d_{e}({\bm{a}},{\bm{c}})$. However, if we endow this
metric space with the ground cost matrix
${\mathbf{M}}=\bigg{[}\begin{matrix}0&1&2\\\ 1&0&1\\\
2&1&0\end{matrix}\bigg{]}$, then it is easy to see that
$W_{1}({\bm{a}},{\bm{b}})=1$ and $W_{1}({\bm{a}},{\bm{c}})=2$. Thereby, if
there were to be a flow of movement from ${\bm{a}}$ to ${\bm{b}}$ or
${\bm{c}}$, the motion ${\bm{a}}\rightarrow{\bm{b}}$ would be the most likely
one.

Thus, the transportation plan ${\mathbf{P}}$ that is computed to obtain the
Wasserstein metric inherently resembles a flow field of mass from source
histogram to the target histogram, making it an attractive choice for modeling
motion and mass transfer.

Thus using the Wasserstein metric instead of the euclidean distance in the
loss function non-negative matrix factorization allows for wiggle room for the
spatial footprints to capture motion varying signal captured in the data
matrix ${\mathbf{X}}$. Intuitively, by using the Euclidean loss, NMF methods
assume that object shapes in the data can be explained as a set of static
spatial footprints perturbed by independent pixel noise. For WNMF methods,
this perturbation is replaced by ’motion noise’ – a more reasonable assumption
for the case where the objects are moving. The first use of the entropically
regularized Wasserstein metric in non-negative matrix factorization was
introduced by Rolet et al.[16] and expanded by Schmitz et al. [17]. The new
cost function minimized by WNMF is
$E_{WNMF}(\mathbf{Y},\mathbf{C})=\sum_{t=1}^{T}W^{\gamma}({\mathbf{X}}_{t},\mathbf{D}{\bm{c}}_{t})$,
where $W^{\gamma}$ is the smoothed Wasserstein distance.

While the advantage of WNMF is that it can capture spatial footprints of
moving objects, the fundamental limitation of it is that, if the motion
variability is too high due to very long time courses, the uncertainty in the
flow fields from a static dictionary to arbitrarily distant (in a Wasserstein
sense) data entries, yields poor reconstruction and thus low data fidelity.

###  3.3 Temporal Wasserstein non-negative matrix factorization

To alleviate the problem of modeling large motion variability using a single
dictionary, we propose to use a temporally constrained piecewise Wasserstein
non-negative matrix factorization. In our formulation, we break the data
matrix ${\mathbf{X}}$ into blocks of length $L$ and model that each of these
blocks is factored as a product of a dictionary ${\mathbf{D}}_{\ell}$ and a
coefficient matrix ${\mathbf{C}}_{\ell}$. To model for motion within these
time blocks, we utilize the entropic Wasserstein metric as the data fidelity
loss. To yield temporally coherent motion tracks, we further constrain the
corresponding terms in the sequential dictionary blocks to be close in the
Wasserstein metric space by adding the terms
$W({\mathbf{D}}_{\ell}^{k},{\mathbf{D}}_{\ell+1}^{k})$. Furthermore, we
penalize the negative entropy, $H({\mathbf{D}})$, of the dictionary elements
to yield spatially compact regions. Lastly, we penalize the entropy of the
loading coefficients to yield spatially smooth signals $H({\mathbf{C}})$. This
formulation yields the objective function below:

$\displaystyle\underset{{\mathbf{D}},{\mathbf{C}}}{\min}\sum_{\ell=1}^{L}\sum_{t\in
T_{\ell}}W^{\gamma}({\bm{x}}_{t},{\mathbf{D}}_{\ell}{\bm{c}}_{t})+\lambda_{T}\sum_{\ell=1}^{L}\sum_{k=1}^{K}W^{\gamma}({\mathbf{D}}_{\ell}^{k},{\mathbf{D}}_{\ell+1}^{k})$
$\displaystyle-\lambda_{D}H({\mathbf{D}})-\lambda_{C}H({\mathbf{C}})$ (3)

We term this model, temporal Wasserstein non-negative matrix factorization
(TWNMF). As in NMF, this objective is non-convex but convex in ${\mathbf{D}}$
and ${\mathbf{C}}$, separately. To optimize the objective, we form the
Lagrangian system that incorporates the dual variables
$\bm{\alpha,\beta,\psi,\omega}$ to account for the Wasserstein metric
constraints:

$\displaystyle\mathcal{L}({\mathbf{D}},{\mathbf{C}},\bm{\pi},\bm{\rho},\bm{\alpha},\bm{\beta},\bm{\psi},\bm{\omega})$
$\displaystyle\sum_{\ell=1}^{L}\sum_{t\in
T_{\ell}}\sum_{i=1}^{d}\sum_{j=1}^{d}\pi_{i,j}^{\ell,t}m_{i,j}+\gamma(\pi_{i,j}^{\ell,t}\log\pi_{i,j}^{\ell,t}-1)$
$\displaystyle+\lambda_{T}\sum_{\ell=1}^{L}\sum_{k=1}^{K}\sum_{i=1}^{d}\sum_{j=1}^{d}\rho_{i,j}^{\ell,k}m_{i,j}+\gamma(\rho_{i,j}^{\ell,k}\log\rho_{i,j}^{\ell,k}-1)$
$\displaystyle+\lambda_{D}\sum_{\ell=1}^{L}\sum_{k=1}^{K}\sum_{i=1}^{d}(D_{\ell,i}^{k}\log
D_{\ell,i}^{k}-1)$
$\displaystyle+\lambda_{C}\sum_{\ell=1}^{L}\sum_{k=1}^{K}\sum_{t\in
T_{\ell}}(C_{\ell,t}^{k}\log C_{\ell,t}^{k}-1)$
$\displaystyle+{\bm{\alpha}^{\ell,t}}^{T}(\bm{\pi}^{\ell,t}\bm{1}-{\bm{x}}_{t})+{\bm{\beta}^{\ell,t}}^{T}({\bm{\pi}^{\ell,t}}^{T}\bm{1}-{\mathbf{D}}_{\ell}{\bm{c}}_{t})$
$\displaystyle+{\bm{\psi}^{\ell,k}}^{T}({\bm{\rho}^{\ell,k}}\bm{1}-{\mathbf{D}}_{\ell}^{k})+{\bm{\omega}^{\ell,k}}^{T}({\bm{\rho}^{\ell,k}}^{T}\bm{1}-{\mathbf{D}}_{\ell+1}^{k})$
(4)

The optimization routine for equation 3.3 is outlined in algorithm 2. For each
time block, we optimize for the dual variables $\bm{\alpha,\beta}$ using the
Sinkhorn iteration using the subroutine 1. Furthermore, the Sinkhorn iterates
are used to optimize the temporal glueing dual variables $\bm{\psi,\omega}$.
Once the dual variables are updated, we perform projectced gradient descent on
the dictionary elements ${\mathbf{D}}$ and ${\mathbf{C}}$, respecting non-
negativity. To obtain spatial tracks from ${\mathbf{D}}^{\ell,k}$, we use the
transportation plan $\bm{\rho}_{\ell}$ to obtain the push-forward flow field
onto the data matrix using
${\mathbf{T}}=\bm{\rho}^{\ell,t}{\mathbf{D}}^{\ell,k}$.

##  4 Results

To demonstrate the utility of our method in performing Spatio-temporal
deconvolution, we experimented with three different applications. First, we
attempted to tackle the motion drift problem in multielectrode array voltage
deconvolution. Then we showed that TWNMF can be used to extract calcium signal
in freely moving C. elegans videos in full spatial resolution, constraining
our analysis to a subregion of voxels such that the computational cost of
performing optimal transport is manageable. Lastly, we attempted to motion
segment a full-body video of C.elegans after it has gone dimensionality
reduction using the matching pursuit algorithm.

###  4.1 Multielectrode array motion drift

To test the effectiveness of TWNMF for recovering neural voltage traces from
multielectrode array recordings[12, 19] in the case of motion drift, we
created data simulating 5 neurons recorded by a drifting multielectrode array
with 100 electrodes. We model five neurons drifting in a one-dimensional
random walk while exhibiting a random phase changed sinusoidal voltage
pattern. We then ran the TWNMF, WNMF[16], and NMF[26] algorithms on the time
series of electrode values, and compared the traces extracted by each
algorithm to the true voltage traces (Figure 3). In this simulation, TWNMF
recovered the original traces with the highest accuracy. The traces extracted
by WNMF were reasonably accurate for the middle time points but differed
significantly from the true signals at time points close to the beginning and
end of the series. This is most likely because the WNMF method models all
object motion as noise, and therefore cannot extract traces from cells that
move significantly from their original positions. The traces extracted by NMF
sometimes captured part of the true signal but were inaccurate for most time
points. This is because NMF does not model cell motion at all, and attempts to
construct a static representation for all cells.

Method | Spatial accuracy | Temporal accuracy
---|---|---
TWNMF | 0.81 $\pm$ 0.07 | 0.62 $\pm$ 0.33
WNMF[16] | 0.72 $\pm$ 0.13 | 0.39 $\pm$ 0.21
NMF[26] | 0.27 $\pm$ 0.37 | 0.05 $\pm$ 0.30
Table 1: Quantitative evaluation of the proposed method, WNMF, and NMF on the
simultaneous motion segmentation and voltage demixing in the simulated
multielectrode drift data. Proposed TWNMF recovers the spatial footprints of
drifting neuron signal with high accuracy while at the same time, the temporal
voltage signal is accurately deconvolved. Figure 6: Motion segmentation of
moving c. elegans worm using proposed method and K-means with K-shortest paths
algorithm (KM+KSP). Top: The raw point cloud of matching pursuit hits of C.
elegans neurons. Middle: Motion segmentation components using proposed TWNMF
algorithm. Bottom: Motion segmentation components using K-means stiched
together with K-shortest paths[2].

###  4.2 Extracting calcium traces from C. elegans videos

Next, we used TWNMF to try to extract calcium traces from real C. elegans
imaging data (Figure 5). For computational efficiency purposes, the sample
video we used was a sum-projection of a zoomed-in section of the original 3D
video. For comparison, we also extracted traces using NMF. The spatial
components extracted by TWNMF correctly identified and tracked the cells, and
the calcium traces extracted by the algorithm also appear reasonable. The
components extracted by NMF, however, are much more distributed spatially, and
their corresponding calcium traces do not look realistic.

###  4.3 Motion segmentation in C. elegans videos

To see if TWNMF could track groups of cells in C. elegans imaging data, we ran
it on a 3D, full-worm fluorescent microscopy video and plotted the time series
of components it extracted (Figure 6). The video was taken from a freely
behaving, transgenic strain of C. elegans called NeuroPAL [24]. The video was
imaged at a temporal resolution of 5Hz and spatial resolution of 0.3 $\mu m$
$\times$ 0.3 $\mu m$ $\times$ 1 $\mu m$ (x,y,z) at 900 $\times$ 600 $\times$
30 voxels. To make TWNMF computationally tractable on the data, each frame was
represented as a weighted combination of 500 Gaussian functions with identical
covariances, whose means and weights were computed using a greedy matching-
pursuit (MP) algorithm [7]. For comparison, a more traditional motion
segmentation approach was also applied to the data. In this method, the
Gaussian components for each frame were first clustered using the k-means
algorithm [11] and then connected across frames using a k-shortest-paths
algorithm [2] to produce the tracked groups of cells.

The motion segmentation results of TWNMF demonstrate the ability to accurately
track consistent parts of the C.elegans worm under large motion between
frames. In comparison, performing K-means clustering which is followed by
K-shortest paths routing to stitch together segments yields a motion
segmentation that is highly inconsistent over time, where clusters highly
drift from one part of the worm to another.

##  5 Conclusion

In this paper, we advocated for a new paradigm for motion segmentation using
optimal transport theory. Using the Wasserstein metric to encourage data
fidelity and temporal continuity enables the formulation of motion
segmentation of non-rigid moving bodies as a non-linear matrix factorization
problem. This contrasts with the latest state of the art techniques for motion
segmentation which heavily rely on optical flow. Our formulation can be cast a
simple piecewise convex factorization model which has the advantage of being
interpretable in its parts and allows for further modeling extensions. A
further advantage of the proposed model is that unlike state of the art
optical flow based motion segmentation models that work on natural images[9],
the temporal coefficients of the model allow for capturing temporal
variability in appearance and brightness intensity of the objects tracked.

Thus we demonstrate that the proposed method works well in biological imaging
scenarios where the intensities of the moving objects fluctuate in time. We
demonstrated its effectiveness for extracting voltage traces from
multielectrode array data in the case of motion drift, as well as tracking and
extracting activity traces from neurons in C. elegans calcium imaging data.
One drawback to the method we have proposed is its computational cost –
storing the optimal transport plan in memory is prohibitively expensive for
high-resolution images. Just like in many of the optimal flow-based methods,
we avoid this problem by using dimensionality reduction techniques to reduce
the memory cost of computing optimal transport. Improving these dimensionality
reduction methods will be an important part of further work.

### Acknowledgements

The authors would like to acknowledge the sources of funding: NSF NeuroNex
Award DBI-1707398, and The Gatsby Charitable Foundation.

## References

  * [1] S. Anthwal and D. Ganotra.  An overview of optical flow-based approaches for motion segmentation.  The Imaging Science Journal, 67(5):284–294, 2019.
  * [2] J. Berclaz, F. Fleuret, E. Turetken, and P. Fua.  Multiple object tracking using k-shortest paths optimization.  IEEE transactions on pattern analysis and machine intelligence, 33(9):1806–1819, 2011.
  * [3] P. Bideau, A. RoyChowdhury, R. R. Menon, and E. Learned-Miller.  The best of both worlds: Combining cnns and geometric constraints for hierarchical motion segmentation.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 508–517, 2018.
  * [4] M. Cuturi.  Sinkhorn distances: Lightspeed computation of optimal transport.  In Advances in neural information processing systems, pages 2292–2300, 2013.
  * [5] M. Cuturi and A. Doucet.  Fast computation of wasserstein barycenters.  In International Conference on Machine Learning, pages 685–693, 2014.
  * [6] A. Dave, P. Tokmakov, and D. Ramanan.  Towards segmenting everything that moves.  arXiv preprint arXiv:1902.03715, 2019.
  * [7] M. Elad.  Sparse and redundant representations: from theory to applications in signal and image processing.  Springer Science & Business Media, 2010.
  * [8] E. Ilg, N. Mayer, T. Saikia, M. Keuper, A. Dosovitskiy, and T. Brox.  Flownet 2.0: Evolution of optical flow estimation with deep networks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2462–2470, 2017.
  * [9] M. Keuper, B. Andres, and T. Brox.  Motion trajectory segmentation via minimum cost multicuts.  In Proceedings of the IEEE International Conference on Computer Vision, pages 3271–3279, 2015.
  * [10] D. D. Lee and H. S. Seung.  Algorithms for non-negative matrix factorization.  In Advances in neural information processing systems, pages 556–562, 2001.
  * [11] S. Lloyd.  Least squares quantization in pcm.  IEEE transactions on information theory, 28(2):129–137, 1982.
  * [12] M. E. J. Obien, K. Deligkaris, T. Bullmann, D. J. Bakkum, and U. Frey.  Revealing neuronal function through microelectrode array recordings.  Frontiers in neuroscience, 8:423, 2015.
  * [13] P. Ochs, J. Malik, and T. Brox.  Segmentation of moving objects by long term video analysis.  IEEE transactions on pattern analysis and machine intelligence, 36(6):1187–1200, 2013.
  * [14] R. Prevedel, Y.-G. Yoon, M. Hoffmann, N. Pak, G. Wetzstein, S. Kato, T. Schrödel, R. Raskar, M. Zimmer, E. S. Boyden, et al.  Simultaneous whole-animal 3d imaging of neuronal activity using light-field microscopy.  Nature methods, 11(7):727, 2014.
  * [15] W. Qian, B. Hong, D. Cai, X. He, X. Li, et al.  Non-negative matrix factorization with sinkhorn distance.  In IJCAI, pages 1960–1966, 2016.
  * [16] A. Rolet, M. Cuturi, and G. Peyré.  Fast dictionary learning with a smoothed wasserstein loss.  In Artificial Intelligence and Statistics, pages 630–638, 2016\.
  * [17] M. A. Schmitz, M. Heitz, N. Bonneel, F. Ngole, D. Coeurjolly, M. Cuturi, G. Peyré, and J.-L. Starck.  Wasserstein dictionary learning: Optimal transport-based unsupervised nonlinear dictionary learning.  SIAM Journal on Imaging Sciences, 11(1):643–678, 2018.
  * [18] R. Sinkhorn and P. Knopp.  Concerning nonnegative matrices and doubly stochastic matrices.  Pacific Journal of Mathematics, 21(2):343–348, 1967.
  * [19] M. E. Spira and A. Hai.  Multi-electrode array technologies for neuroscience and cardiology.  Nature nanotechnology, 8(2):83, 2013.
  * [20] B. Taylor, V. Karasev, and S. Soatto.  Causal video object segmentation from persistence of occlusions.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 4268–4276, 2015.
  * [21] C. Tomasi and T. Kanade.  Shape and motion from image streams under orthography: a factorization method.  International journal of computer vision, 9(2):137–154, 1992.
  * [22] C. Villani.  Optimal transport: old and new, volume 338.  Springer Science & Business Media, 2008.
  * [23] C. Xie, Y. Xiang, Z. Harchaoui, and D. Fox.  Object discovery in videos as foreground motion clustering.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 9994–10003, 2019.
  * [24] E. Yemini, A. Lin, A. Nejatbakhsh, E. Varol, R. Sun, G. E. Mena, A. D. Samuel, L. Paninski, V. Venkatachalam, and O. Hobert.  Neuropal: A neuronal polychromatic atlas of landmarks for whole-brain imaging in c. elegans.  bioRxiv, page 676312, 2019.
  * [25] L. Zappella, X. Lladó, and J. Salvi.  Motion segmentation: A review.  In Proceedings of the 2008 conference on Artificial Intelligence Research and Development: Proceedings of the 11th International Conference of the Catalan Association for Artificial Intelligence, pages 398–407. IOS Press, 2008.
  * [26] P. Zhou, S. L. Resendez, J. Rodriguez-Romaguera, J. C. Jimenez, S. Q. Neufeld, A. Giovannucci, J. Friedrich, E. A. Pnevmatikakis, G. D. Stuber, R. Hen, et al.  Efficient and accurate extraction of in vivo calcium signals from microendoscopic video data.  Elife, 7:e28728, 2018.